The KL bound

How jointly minimizing the reconstruction loss and the Gram-matrix alignment loss bounds the KL divergence of a downstream linear classifier.

Let $\mathbf{A} \in \mathbb{R}^{N \times d}$ be the patch-token activations at layer $\ell$ and $\hat{\mathbf{A}} = \mathbf{U}\mathbf{W}^{\top}$ their Semi-NMF^{6Convex and Semi-Nonnegative Matrix Factorizations
Chris H.Q. Ding, Tao Li, Michael I. Jordan
IEEE TPAMI, 2010} reconstruction, with Gram matrices $\mathbf{G} = \mathbf{A}\mathbf{A}^{\top}$ and $\hat{\mathbf{G}} = \hat{\mathbf{A}}\hat{\mathbf{A}}^{\top}$ . The two error terms controlled by the objective are

\delta_{\text{rec}} = \|\mathbf{A} - \hat{\mathbf{A}}\|_F, \qquad \delta_{\text{gram}} = \sqrt{\mathcal{L}_{\text{gram}}} = \frac{\|\mathbf{G} - \hat{\mathbf{G}}\|_F}{\|\mathbf{G}\|_F}.

Theorem (Appendix C, Eq. 35). Let

h(z) = W_c z + b

be any linear classifier on the CLS token, applied after the remaining

L_{\text{suf}} = L - \ell

transformer blocks

g_{>\ell}

. Then

\begin{aligned} &D_{\mathrm{KL}}\big(\sigma(h \circ g_{>\ell}(\mathbf{A})) \,\big\|\, \sigma(h \circ g_{>\ell}(\hat{\mathbf{A}}))\big) \\ &\quad \le \frac{1}{2} \|W_c\|_{\text{op}}^2\, L_{\text{block}}^{2L_{\text{suf}}} \big( \sqrt{N}\, \delta_{\text{rec}} + \alpha\, \|\mathbf{G}\|_F \cdot \delta_{\text{gram}} \big)^2, \end{aligned}

where

L_{\text{block}}

is the per-block Lipschitz constant and

\alpha

collects attention-projection constants.

How the proof goes

Reconstruction controls the Gram matrix only indirectly. A lemma gives $\|\mathbf{G} - \hat{\mathbf{G}}\|_F \le 2\|\mathbf{A}\|_F\, \delta_{\text{rec}} + \delta_{\text{rec}}^2$ , a bound that grows with the activation scale $\|\mathbf{A}\|_F$ .
Attention logits. Splitting the query-key product as $W_Q W_K^{\top} = I + (W_Q W_K^{\top} - I)$ shows the Gram error controls the identity part of the logit deviation directly, and reconstruction controls the rest.
One block. Softmax is 1-Lipschitz, and LayerNorm and the FFN are Lipschitz, so each block amplifies the error by at most $L_{\text{block}}$ .
All suffix blocks. Composing over the $L_{\text{suf}}$ blocks bounds the CLS-token deviation by $L_{\text{block}}^{L_{\text{suf}}} (\sqrt{N}\, \delta_{\text{rec}} + \alpha \|\mathbf{G}\|_F\, \delta_{\text{gram}})$ .
To KL. The KL divergence between softmax outputs is a Bregman divergence of the log-partition function, whose Hessian has spectral norm at most 1, so $D_{\mathrm{KL}} \le \tfrac{1}{2}\|\mathbf{u} - \hat{\mathbf{u}}\|_2^2$ for logits $\mathbf{u}$ .

What it means

The two losses play complementary roles. $\mathcal{L}_{\text{gram}}$ constrains the attention-logit deviation directly and independently of scale, which reconstruction does only through $\|\mathbf{A}\|_F \cdot \delta_{\text{rec}}$ . $\mathcal{L}_{\text{rec}}$ supplies pointwise fidelity, which the Gram loss cannot constrain because it is invariant to orthogonal transformations of $\mathbf{A}$ . Neither needs a forward or backward pass through $g_{>\ell}$ or $h$ .

The bound is worst-case and compounds across suffix blocks, so its numerical value is loose. It tightens as $\ell$ approaches the last layer. It identifies the two quantities whose joint minimization is sufficient to control predictive consistency. The agreement with classifier-supervised concepts is measured on the results page.

Full proof in the paper (Appendix C)Back to the method

References

[6]Chris H.Q. Ding, Tao Li, Michael I. Jordan. Convex and Semi-Nonnegative Matrix Factorizations. IEEE TPAMI, 2010.