SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- First, we demonstrate a mechanism to learn representations via denoising within a idealized mixture of Gaussian data model for a single token (i.e., with sequence length  $N = 1$ ).
- Second, we illustrate the similarities between a such-derived representation learning scheme and existing self-attention layers within the transformer (with sequence length 1), thus demonstrating an interpretation of the self-attention layer as a generalized mechanism to denoise against a mixture-of-Gaussian-marginal model for a set of tokens.
Now we produce the proofs alluded to in Section 2.2, which mostly form the technical aspects of the first listed contribution. To simplify the proofs, we use the following notation correspondences:  $\boldsymbol{x} \mapsto \boldsymbol{z}^{\ell}$ ,  $\boldsymbol{z} \mapsto \boldsymbol{z}^{\ell + 1}$ , and  $\sigma \mapsto \sigma^{\ell}$ .
Proposition 1. Let  $\mathbf{u}_1, \ldots, \mathbf{u}_K \in \mathbb{R}^d$  be independent and have distribution  $\mathbf{u}_k \sim \mathcal{N}(\mathbf{0}, \mathbf{\Sigma}_k)$  for  $\mathbf{\Sigma}_k \succeq \mathbf{0}$ , and let  $\mathbf{z}$  take value  $\mathbf{u}_k$  with probability  $\pi_k > 0$ . Let  $\mathbf{w} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}_d)$  be independent of  $\mathbf{z}$ . Let  $\mathbf{x} \dot{=} \mathbf{z} + \sigma \mathbf{w}$ . Let  $\mathbf{x} \mapsto q(\mathbf{x})$  be the density of  $\mathbf{x}$ . We define
$$
\boldsymbol {M} _ {k} \doteq \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1 / 2} \tag {19}
$$
and assume that  $\pi_i\operatorname *{det}(M_i) = \pi_j\operatorname *{det}(M_j)$  for all  $1\leq i\leq j\leq K$  . Then we have
$$
\begin{array}{l} \nabla_ {\boldsymbol {x}} \log q (\boldsymbol {x}) (20) \\ = - \left[ M _ {1}, \dots , M _ {K} \right] \left[ \operatorname {d i a g} \left(\operatorname {s o f t m a x} \left(- \frac {1}{2} \left[ \begin{array}{c} \left\| M _ {1} ^ {*} \boldsymbol {x} \right\| _ {2} ^ {2} \\ \vdots \\ \left\| M _ {K} ^ {*} \boldsymbol {x} \right\| _ {2} ^ {2} \end{array} \right]\right)\right) \otimes I _ {d} \right] \left[ \begin{array}{c} M _ {1} ^ {*} \boldsymbol {x} \\ \vdots \\ M _ {K} ^ {*} \boldsymbol {x} \end{array} \right], (21) \\ \end{array}
$$
where  $\otimes$  denotes the Kronecker product, i.e., the block matrix defined by
$$
\boldsymbol {A} \otimes \boldsymbol {B} = \left[ \begin{array}{c c c} A _ {1 1} \boldsymbol {B} & \dots & A _ {1 n} \boldsymbol {B} \\ \vdots & \ddots & \vdots \\ A _ {m 1} \boldsymbol {B} & \dots & A _ {m n} \boldsymbol {B} \end{array} \right] \tag {22}
$$
Proof. Let  $u$  be the multinomial random variable such that  $z = z_{u}$ , so that  $u$  has probability mass function  $\pi$ . Then by the law of total probability, we have
$$
\begin{array}{l} \nabla_ {\boldsymbol {x}} \log q (\boldsymbol {x}) = \nabla_ {\boldsymbol {x}} \log \sum_ {k = 1} ^ {K} q (\boldsymbol {x} \mid k) \pi_ {k} (23) \\ = \frac {\sum_ {k = 1} ^ {K} \pi_ {k} \nabla_ {\boldsymbol {x}} q (\boldsymbol {x} \mid k)}{\sum_ {k = 1} ^ {K} q (\boldsymbol {x} \mid k) \pi_ {k}} (24) \\ \end{array}
$$
where  $q(\pmb{x} \mid k)$  is the conditional density of  $\pmb{x}$  given the event  $\{u = k\}$ . To compute this quantity, note that conditional on the value of  $u$ , we have
$$
\boldsymbol {x} = \boldsymbol {z} _ {u} + \sigma \boldsymbol {w} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {\Sigma} _ {u} + \sigma^ {2} \boldsymbol {I} _ {d}). \tag {25}
$$
$$
q (\boldsymbol {x} \mid k) = \frac {1}{\sqrt {(2 \pi) ^ {d}} \det \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right)} \exp \left(- \frac {1}{2} \boldsymbol {x} ^ {*} \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1} \boldsymbol {x}\right), \tag {26}
$$
$$
\nabla_ {\boldsymbol {x}} q (\boldsymbol {x} \mid k) = - q (\boldsymbol {x} \mid k) \cdot \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1} \boldsymbol {x}. \tag {27}
$$
Putting this all together, we get
$$
\begin{array}{l} \nabla_ {\boldsymbol {x}} \log q (\boldsymbol {x}) (28) \\ = - \frac {\sum_ {k = 1} ^ {K} q (\boldsymbol {x} \mid k) \pi_ {k} \cdot \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1} \boldsymbol {x}}{\sum_ {k = 1} ^ {K} q (\boldsymbol {x} \mid k) \pi_ {k}} (29) \\ = - \frac {\sum_ {k = 1} ^ {K} \pi_ {k} \det  \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1 / 2} \exp \left(- \frac {1}{2} \boldsymbol {x} ^ {*} \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1} \boldsymbol {x}\right) \cdot \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1} \boldsymbol {x}}{\sum_ {k = 1} ^ {K} \pi_ {k} \det  \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1 / 2} \exp \left(- \frac {1}{2} \boldsymbol {x} ^ {*} \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} I _ {d}\right) ^ {- 1} \boldsymbol {x}\right)}. (30) \\ \end{array}
$$
Now define  $M_{k}\dot{=} (\pmb{\Sigma}_{k} + \sigma^{2}\pmb{I}_{d})^{-1 / 2}$ . With this notation, we have
$$
\begin{array}{l} \nabla_ {\boldsymbol {x}} \log q (\boldsymbol {x}) = - \frac {\sum_ {k = 1} ^ {K} \pi_ {k} \det  \left(\boldsymbol {M} _ {k}\right) \exp \left(- \frac {1}{2} \boldsymbol {x} ^ {*} \boldsymbol {M} _ {k} \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x}\right) \cdot \boldsymbol {M} _ {k} \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x}}{\sum_ {k = 1} ^ {K} \pi_ {k} \det  \left(\boldsymbol {M} _ {k}\right) \exp \left(- \frac {1}{2} \boldsymbol {x} ^ {*} \boldsymbol {M} _ {k} \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x}\right)} (31) \\ = - \frac {\sum_ {k = 1} ^ {K} \pi_ {k} \det \left(\boldsymbol {M} _ {k}\right) \exp \left(- \frac {1}{2} \| \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x} \| _ {2} ^ {2}\right) \cdot \boldsymbol {M} _ {k} \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x}}{\sum_ {k = 1} ^ {K} \pi_ {k} \det (\boldsymbol {M} _ {k}) \exp \left(- \frac {1}{2} \boldsymbol {x} ^ {*} \boldsymbol {M} _ {k} \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x}\right)}. (32) \\ \end{array}
$$
Given our assumption that each  $\pi_k\operatorname*{det}(M_k)$  is the same, we have
$$
\begin{array}{l} \nabla_ {\boldsymbol {x}} \log q (\boldsymbol {x}) (33) \\ = - \frac {\sum_ {k = 1} ^ {K} \pi_ {k} \det \left(\boldsymbol {M} _ {k}\right) \exp \left(- \frac {1}{2} \| \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x} \| _ {2} ^ {2}\right) \cdot \boldsymbol {M} _ {k} \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x}}{\sum_ {k = 1} ^ {K} \pi_ {k} \det \left(\boldsymbol {M} _ {k}\right) \exp \left(- \frac {1}{2} \| \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x} \| _ {2} ^ {2}\right)} (34) \\ = - \frac {\sum_ {k = 1} ^ {K} \exp \left(- \frac {1}{2} \| M _ {k} ^ {*} \boldsymbol {x} \| _ {2} ^ {2}\right) \cdot M _ {k} M _ {k} ^ {*} \boldsymbol {x}}{\sum_ {k = 1} ^ {K} \exp \left(- \frac {1}{2} \| M _ {k} ^ {*} \boldsymbol {x} \| _ {2} ^ {2}\right)} (35) \\ = - \sum_ {k = 1} ^ {K} e _ {k} ^ {*} \operatorname {s o f t m a x} \left(- \frac {1}{2} \left[ \begin{array}{c} \| M _ {1} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \\ \vdots \\ \| M _ {K} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \end{array} \right]\right) M _ {k} M _ {k} ^ {*} \boldsymbol {x} (36) \\ = - \left[ M _ {1}, \dots , M _ {K} \right] \left[ \operatorname {d i a g} \left(\operatorname {s o f t m a x} \left(- \frac {1}{2} \left[ \begin{array}{c} \| M _ {1} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \\ \vdots \\ \| M _ {K} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \end{array} \right]\right)\right) \otimes I _ {d} \right] \left[ \begin{array}{c} M _ {1} ^ {*} \boldsymbol {x} \\ \vdots \\ M _ {K} ^ {*} \boldsymbol {x} \end{array} \right]. (37) \\ \end{array}
$$
Now we provide a final justification for the result cited in Section 2.2.
Approximation 2. In the setting of Proposition 1, diagonalize  $\pmb{\Sigma}_{k} = \pmb{U}_{k}\pmb{\Lambda}_{k}\pmb{U}_{k}^{*}$  where  $\pmb{U}_k\in \mathbb{R}^{d\times p}$  is orthogonal and  $\pmb{\Lambda}_k\succ \mathbf{0}\in \mathbb{R}^{p\times p}$  is diagonal. Then we have the approximation
$$
\mathbb {E} [ \boldsymbol {z} \mid \boldsymbol {x} ] \approx [ U _ {1}, \dots , U _ {K} ] \left[ \operatorname {d i a g} \left(\operatorname {s o f t m a x} \left(\frac {1}{2 \sigma^ {2}} \left[ \begin{array}{c} \| U _ {1} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \\ \vdots \\ \| U _ {K} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \end{array} \right]\right)\right) \otimes \boldsymbol {I} _ {p} \right] \left[ \begin{array}{c} \boldsymbol {U} _ {1} ^ {*} \boldsymbol {x} \\ \vdots \\ \boldsymbol {U} _ {K} ^ {*} \boldsymbol {x} \end{array} \right]. \tag {38}
$$
$$
\begin{array}{l} \nabla_ {\boldsymbol {x}} \log q (\boldsymbol {x}) = - \sum_ {k = 1} ^ {K} e _ {k} ^ {*} \operatorname {s o f t m a x} \left(- \frac {1}{2} \left[ \begin{array}{c} \| M _ {1} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \\ \vdots \\ \| M _ {K} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \end{array} \right]\right) M _ {k} M _ {k} ^ {*} \boldsymbol {x} (39) \\ = - \sum_ {k = 1} ^ {K} \boldsymbol {e} _ {k} ^ {*} \operatorname {s o f t m a x} \left(- \frac {1}{2 \sigma^ {2}} \left[ \begin{array}{c} \| \sigma \boldsymbol {M} _ {1} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \\ \vdots \\ \| \sigma \boldsymbol {M} _ {K} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \end{array} \right]\right) \boldsymbol {M} _ {k} \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x} (40) \\ \end{array}
$$
<sup>9</sup>This assumption can be easily relaxed to  $\Lambda_{k} \succeq 0$  for all  $k$ , but requires some more notation to handle, and the form of the solution does not change. Thus we handle the case where all matrices are full rank for simplicity.
$$
= - \sum_ {k = 1} ^ {K} \boldsymbol {e} _ {k} ^ {*} \operatorname {s o f t m a x} \left(\frac {1}{2 \sigma^ {2}} \left[ \begin{array}{c} \| \boldsymbol {x} \| _ {2} ^ {2} - \| \sigma \boldsymbol {M} _ {1} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \\ \vdots \\ \| \boldsymbol {x} \| _ {2} ^ {2} - \| \sigma \boldsymbol {M} _ {K} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \end{array} \right]\right) \boldsymbol {M} _ {k} \boldsymbol {M} _ {k} ^ {*} \boldsymbol {x}. \tag {41}
$$
Now define  $\pmb{P}_k \doteq \pmb{I}_d - \sigma \pmb{M}_k$ , and let  $U_k^\perp \in \mathbb{R}^{d \times (d - p)}$  be an orthogonal complement of  $U_k$ . Then we have
$$
\begin{array}{l} \boldsymbol {P} _ {k} = \boldsymbol {I} _ {d} - \sigma \boldsymbol {M} _ {k} (42) \\ = \boldsymbol {I} _ {d} - \sigma \left(\boldsymbol {\Sigma} _ {k} + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1 / 2} (43) \\ = \boldsymbol {I} _ {d} - \sigma \left(\left[ \begin{array}{l l} \boldsymbol {U} _ {k} & \boldsymbol {U} _ {k} ^ {\perp} \end{array} \right] \left[ \begin{array}{l l} \boldsymbol {\Lambda} _ {k} & \boldsymbol {0} \\ \boldsymbol {0} & \boldsymbol {0} \end{array} \right] \left[ \begin{array}{l} \boldsymbol {U} _ {k} ^ {*} \\ (\boldsymbol {U} _ {k} ^ {\perp}) ^ {*} \end{array} \right] + \sigma^ {2} \boldsymbol {I} _ {d}\right) ^ {- 1 / 2} (44) \\ = \boldsymbol {I} _ {d} - \sigma \left(\left[ \begin{array}{l l} \boldsymbol {U} _ {k} & \boldsymbol {U} _ {k} ^ {\perp} \end{array} \right] \left[ \begin{array}{c c} \boldsymbol {\Lambda} _ {k} + \sigma^ {2} \boldsymbol {I} _ {p} & \boldsymbol {0} \\ \boldsymbol {0} & \sigma^ {2} \boldsymbol {I} _ {d - p} \end{array} \right] \left[ \begin{array}{c} \boldsymbol {U} _ {k} ^ {*} \\ (\boldsymbol {U} _ {k} ^ {\perp}) ^ {*} \end{array} \right]\right) ^ {- 1 / 2} (45) \\ = \boldsymbol {I} _ {d} - \left[ \begin{array}{l l} \boldsymbol {U} _ {k} & \boldsymbol {U} _ {k} ^ {\perp} \end{array} \right] \left[ \begin{array}{c c} \sigma \left(\boldsymbol {\Lambda} _ {k} + \sigma^ {2} \boldsymbol {I} _ {p}\right) ^ {- 1 / 2} & \boldsymbol {0} \\ \boldsymbol {0} & \sigma \cdot (\sigma^ {2}) ^ {- 1 / 2} \boldsymbol {I} _ {d - p} \end{array} \right] \left[ \begin{array}{c} \boldsymbol {U} _ {k} ^ {*} \\ \left(\boldsymbol {U} _ {k} ^ {\perp}\right) ^ {*} \end{array} \right] (46) \\ = \boldsymbol {I} _ {d} - \left[ \begin{array}{l l} \boldsymbol {U} _ {k} & \boldsymbol {U} _ {k} ^ {\perp} \end{array} \right] \left[ \begin{array}{c c} \left(\sigma^ {- 2} \boldsymbol {\Lambda} _ {k} + \boldsymbol {I} _ {p}\right) ^ {- 1 / 2} & \boldsymbol {0} \\ \boldsymbol {0} & \boldsymbol {I} _ {d - p} \end{array} \right] \left[ \begin{array}{c} \boldsymbol {U} _ {k} ^ {*} \\ \left(\boldsymbol {U} _ {k} ^ {\perp}\right) ^ {*} \end{array} \right] (47) \\ = \left[ \begin{array}{c c} \boldsymbol {U} _ {k} & \boldsymbol {U} _ {k} ^ {\perp} \end{array} \right] \left[ \begin{array}{c c} \boldsymbol {I} _ {p} - \left(\sigma^ {- 2} \boldsymbol {\Lambda} _ {k} + \boldsymbol {I} _ {p}\right) ^ {- 1 / 2} & \boldsymbol {0} \\ \boldsymbol {0} & \boldsymbol {0} \end{array} \right] \left[ \begin{array}{c} \boldsymbol {U} _ {k} ^ {*} \\ \left(\boldsymbol {U} _ {k} ^ {\perp}\right) ^ {*} \end{array} \right] (48) \\ \approx \left[ \begin{array}{l l} \boldsymbol {U} _ {k} & \boldsymbol {U} _ {k} ^ {\perp} \end{array} \right] \left[ \begin{array}{l l} \boldsymbol {I} _ {p} & \boldsymbol {0} \\ \boldsymbol {0} & \boldsymbol {0} \end{array} \right] \left[ \begin{array}{l} \boldsymbol {U} _ {k} ^ {*} \\ (\boldsymbol {U} _ {k} ^ {\perp}) ^ {*} \end{array} \right] (49) \\ = \boldsymbol {U} _ {k} \boldsymbol {U} _ {k} ^ {*}. (50) \\ \end{array}
$$
Thus  $P_{k}$  is approximately a projection when  $\sigma$  is small. Under this algebraic relation, we have
$$
= - \frac {1}{\sigma^ {2}} \boldsymbol {x} + \frac {1}{\sigma^ {2}} \left[ \boldsymbol {U} _ {1}, \dots , \boldsymbol {U} _ {K} \right] \left[ \operatorname {d i a g} \left(\operatorname {s o f t m a x} \left(\frac {1}{2 \sigma^ {2}} \left[ \begin{array}{c} \| \boldsymbol {U} _ {1} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \\ \vdots \\ \| \boldsymbol {U} _ {K} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \end{array} \right]\right)\right) \otimes \boldsymbol {I} _ {p} \right] \left[ \begin{array}{c} \boldsymbol {U} _ {1} ^ {*} \boldsymbol {x} \\ \vdots \\ \boldsymbol {U} _ {K} ^ {*} \boldsymbol {x} \end{array} \right]. \tag {59}
$$
Plugging this into Tweedie's formula, we have
$$
\mathbb {E} [ \boldsymbol {z} \mid \boldsymbol {x} ] \approx \left[ U _ {1}, \dots , U _ {K} \right] \left[ \operatorname {d i a g} \left(\operatorname {s o f t m a x} \left(\frac {1}{2 \sigma^ {2}} \left[ \begin{array}{c} \| U _ {1} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \\ \vdots \\ \| U _ {K} ^ {*} \boldsymbol {x} \| _ {2} ^ {2} \end{array} \right]\right)\right) \otimes \boldsymbol {I} _ {p} \right] \left[ \begin{array}{c} U _ {1} ^ {*} \boldsymbol {x} \\ \vdots \\ U _ {K} ^ {*} \boldsymbol {x} \end{array} \right]. \tag {60}
$$
Remark 3. Although Approximation 2 is stated as an approximation rather than as a proposition, we believe it should be possible without too much extra work to convert it into a statement of asymptotic equivalence as  $\sigma \rightarrow 0$  (in particular, holding for  $\sigma$  below the smallest (nonzero) eigenvalue of any  $\Sigma_{k}$ ). Most approximations taken in the derivation of Approximation 2 can immediately be turned into asymptotic claims; the only slightly delicate point is treating the softmax, which can be accomplished using standard "high temperature" convergence behavior of the softmax function (in particular, as  $\sigma \rightarrow 0$  in our expressions, the softmax concentrates on the "best head").
A.2 Companion to Section 2.3
We again wish to re-iterate the core contribution of our approach in Section 2.3. The application of a compression perspective to representation learning has been discussed before, for example in the line of maximal coding rate reduction works [46]. In Section 2.3, we provide the following contributions and developments to this perspective:
- We propose a generalized coding rate function  $R^{c}(\cdot; \mathbf{U}_{[K]})$  which measures the coding rate with respect to a set of subspaces  $\mathbf{U}_{[K]}$  as opposed to a set of classes (as in [46, 55]), making the underlying formulation unsupervised.
- We then show how if we adopt the framework of alternating minimization of the sparse rate reduction objective, then unrolling the first alternating step — gradient descent on this coding rate objective — nearly exactly recovers the common multi-head attention mechanism found in transformer networks (except that the query/key/value operators are all the same operation  $U_{k}^{*}$  now, which we interpret as projection onto a single subspace).
In the process of the second contribution, and in the following proofs, we make some simple approximations and technical assumptions. The validity of these assumptions may be explored, and the approximations refined, altogether providing a more complex (and possibly more performant) resulting self-attention like operator. For the sake of technical clarity and simplicity in this work, we make perhaps the simplest possible choices. As a result, we do not claim that our network is optimally designed, but rather that the principles we develop in this work (compression, denoising, sparsification, unrolled optimization) can provide the backbone for far superior and more interpretable network architectures in the future on sundry tasks. As it is, with our straightforward, simple, and interpretable design, we still obtain meaningful conceptual results and very solid empirical performance.
We now give the derivation of the approximation alluded to in Section 2.3.
Approximation 4. Let  $\mathbf{Z} \in \mathbb{R}^{d \times N}$  have unit-norm columns, and  $\mathbf{U}_{[K]} = (\mathbf{U}_1, \dots, \mathbf{U}_K)$  such that each  $\mathbf{U}_k \in \mathbb{R}^{d \times p}$  is an orthogonal matrix, the  $(\mathbf{U}_k)_{k=1}^K$  are incoherent, and the columns of  $\mathbf{Z}$  approximately lie on  $\bigcup_{k=1}^K \operatorname{Span}(\mathbf{U}_k)$ . Let  $\gamma = \frac{p}{N\epsilon^2}$ . Let  $\kappa > 0$ . Then
$$
\boldsymbol {Z} - \kappa \nabla_ {\boldsymbol {Z}} R ^ {c} (\boldsymbol {Z} \mid \boldsymbol {U} _ {[ K ]}) \approx (1 - \kappa \gamma) \boldsymbol {Z} + \kappa \gamma \operatorname {M S S A} (\boldsymbol {Z} | \boldsymbol {U} _ {[ K ]}), \tag {61}
$$
where as in Section 2.3 we have
$$
\mathit {S S A} (\boldsymbol {Z} | \boldsymbol {U} _ {k}) = \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) \operatorname {s o f t m a x} \left(\left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) ^ {*} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right)\right), \tag {62}
$$
$$
M S S A (\boldsymbol {Z} | \boldsymbol {U} _ {[ K ]}) = \gamma [ \boldsymbol {U} _ {1}, \dots , \boldsymbol {U} _ {K} ] \left[ \begin{array}{c} S S A (\boldsymbol {Z} | \boldsymbol {U} _ {1}) \\ \vdots \\ S S A (\boldsymbol {Z} | \boldsymbol {U} _ {K}) \end{array} \right], \tag {63}
$$
where  $\mathrm{softmax}(\cdot)$  is the softmax operator (applied to each column of an input matrix), i.e.,
$$
\operatorname {s o f t m a x} (\boldsymbol {v}) = \frac {1}{\sum_ {i = 1} ^ {n} e ^ {v _ {i}}} \left[ \begin{array}{c} e ^ {v _ {1}} \\ \vdots \\ e ^ {v _ {n}} \end{array} \right], \tag {64}
$$
$$
\operatorname {s o f t m a x} \left(\left[ \boldsymbol {v} _ {1}, \dots , \boldsymbol {v} _ {K} \right]\right) = \left[ \operatorname {s o f t m a x} \left(\boldsymbol {v} _ {1}\right), \dots , \operatorname {s o f t m a x} \left(\boldsymbol {v} _ {K}\right) \right]. \tag {65}
$$
Proof. According to (9), the gradient  $\nabla_{\mathbf{Z}}R^{c}(\mathbf{Z};\mathbf{U}_{[K]})$  is
$$
\nabla_ {\boldsymbol {Z}} R ^ {c} (\boldsymbol {Z}; \boldsymbol {U} _ {[ K ]}) = \gamma \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} \left(\boldsymbol {I} + \gamma \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) ^ {*} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right)\right) ^ {- 1}. \tag {66}
$$
Notice that according to [55], the gradient is precisely the residual of a ridge regression for each (projected) token  $U_{k}^{*}z_{i}$  using other projected tokens  $U_{k}^{*}z_{j}$  as the regressors, hence being the residual of an auto-regression.
However, as we have seen in the work of ReduNet [55], computing the inverse  $(\mathbf{I} + \gamma (\mathbf{U}_k^*\mathbf{Z})^* (\mathbf{U}_k^*\mathbf{Z}))^{-1}$  can be expensive. Hence for computational efficiency, we may approximate it with the first order term of its von Neumann expansion:
$$
\begin{array}{l} \nabla_ {\boldsymbol {Z}} R ^ {c} (\boldsymbol {Z}; \boldsymbol {U} _ {[ K ]}) = \gamma \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} \left(\boldsymbol {I} + \gamma \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) ^ {*} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right)\right) ^ {- 1} (67) \\ \approx \gamma \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} \left(\boldsymbol {I} - \gamma \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) ^ {*} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right)\right) (68) \\ = \gamma \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} - \gamma \boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} \left[ \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) ^ {*} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) \right]\right) (69) \\ \end{array}
$$
Notice that the term  $(U_k^* Z)^*(U_k^* Z)$  is the auto-correlation among the projected tokens. As the tokens  $Z$  may be from different subspaces, we would prefer to use only tokens that belong to the same subspace to regress and compress themselves. Hence we may convert the above correlation term into a subspace-membership indicator with a softmax operation, whence (69) becomes
$$
\begin{array}{l} \nabla_ {Z} R ^ {c} \left(\boldsymbol {Z}; \boldsymbol {U} _ {[ K ]}\right) \approx \gamma \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} - \gamma \boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} \left[ \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) ^ {*} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) \right]\right) (70) \\ \approx \gamma \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} - \gamma^ {2} \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} \operatorname {s o f t m a x} \left(\left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) ^ {*} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right)\right)\right) (71) \\ \end{array}
$$
Then, we can rewrite the above approximation to the gradient of  $R^c$  as:
$$
\begin{array}{l} \nabla_ {Z} R ^ {c} \left(\boldsymbol {Z}; \boldsymbol {U} _ {[ K ]}\right) \approx \gamma \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} - \gamma^ {2} \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} \operatorname {s o f t m a x} \left(\left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right) ^ {*} \left(\boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z}\right)\right)\right) (72) \\ = \gamma \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \boldsymbol {U} _ {k} ^ {*} \boldsymbol {Z} - \gamma^ {2} \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \operatorname {S S A} (\boldsymbol {Z} \mid \boldsymbol {U} _ {k}) (73) \\ = \underbrace {\left(\gamma \sum_ {k = 1} ^ {K} \boldsymbol {U} _ {k} \boldsymbol {U} _ {k} ^ {*}\right) \boldsymbol {Z}} _ {\approx \gamma Z} - \gamma^ {2} \left[ \boldsymbol {U} _ {1}, \dots , \boldsymbol {U} _ {K} \right] \left[ \begin{array}{c c} \operatorname {S S A} (\boldsymbol {Z} & \boldsymbol {U} _ {1}) \\ & \vdots \\ \operatorname {S S A} (\boldsymbol {Z} & \boldsymbol {U} _ {K}) \end{array} \right] (74) \\ \approx \gamma \boldsymbol {Z} - \gamma^ {2} \left[ \boldsymbol {U} _ {1}, \dots , \boldsymbol {U} _ {K} \right] \left[ \begin{array}{c c} \operatorname {S S A} (\boldsymbol {Z} & \boldsymbol {U} _ {1}) \\ & \vdots \\ \operatorname {S S A} (\boldsymbol {Z} & \boldsymbol {U} _ {K}) \end{array} \right]. (75) \\ \end{array}
$$
Thus the gradient descent step with learning rate  $\kappa > 0$  gives