1 Introduction↩︎

Transformers have been widely adopted in language modeling [1], vision tasks [2], [3], and speech recognition [4]. A crucial building block in transformers is the attention mechanism, dating back to [5], which was initially designed to capture long-range signals in sequential inputs by mixing individual tokens but has also been leveraged to capture general structures of input data. After the fully-attention-based language model has appeared [6], the research community gets interested in the functionality and benefits of the attention. To mention a few, transformers implicitly prefer hierarchical interpretations of input sequences [7]; store relational knowledge in MLP layers as an associative memory [8]; its computational graphs tend to be tree-structured [9]; suddenly capture tree structures of inputs after long training epochs [10]. Theoretically, training dynamics analysis explains how to learn spatially correlated patches by vision transformers (ViT) [11], select dominant tokens [12], store information as an associative memory [13], and select max-margin tokens [14], whereas [15] explains the in-context learning as a process of concept learning in Bayesian inference.

Among many aspects of attention, we specifically focus on localization—for a query token, self-attention can select a few relevant tokens only (which we call localized attention) or select many tokens uniformly. As attention can be regarded as a token mixer, it plays a pivotal role in studying how it selects tokens to reveal the characteristics of the token embeddings. To this end, we have the following research questions: (Q1) When is self-attention localized or uniform? (Q2) How does localization affect model performances?

Along this line, previous studies mainly investigated from the model expressivity and training stability perspectives. On the one hand, [16] and [17] initiated the discussion of attention localization and theoretically showed that a network with self-attention layers without skip connections exponentially loses the rank of hidden layers; the fact indicates that the model expressivity shall be immediately lost with more self-attention layers stacked. On the other hand, [18] empirically found that attention entropy—averaged Shannon entropy of an attention probability matrix—correlates with training stability. Specifically, a training loss curve tends to fall into a plateau when attention entropy is low. Since higher entropy indicates near-uniform attention weights, their finding apparently suggests that localized attention may lead the learning dynamics to a plateau. Up until now, these two failure modes have been discussed independently with slightly different notions of attention localization, and hence, our understanding of the blessing and curse of attention localization remains elusive.

To better comprehend, we characterize self-attention patterns by attention parameter matrices to reconcile the two collapse modes. We formulate the concept of localization by signal propagation probability (3), which describes how likely the signal of a specific input token propagates to the gradient of a training objective. If the signal propagation probability is high for a few numbers of tokens only, attention is regarded to be localized. We show that the localization mode can be characterized by the eigenspectrum of attention weight matrices (4). Specifically, attention is localized in the above sense when the eigenspectrum of the query-key parameter matrix has a non-zero mean and a small variance. Furthermore, the small eigenspectrum variance is relevant to both the rank collapse and entropy collapse (5), and thus, we give a unified perspective of the two notions of attention collapse. For this reason, we argue that attention collapse and its performance can be viewed more transparently based on the eigenspectrum variance. Lastly, we verified the correlation of the eigenspectrum and the model performance in the experiments with the WikiText dataset [19] by introducing a regularization scheme called LocAteR.

2 Setup↩︎

2.0.0.1 Notation.

We write vectors with all zeros and ones by \(\mathbf{0}\) and \(\mathbf{1}\), respectively, whereas the \(i\)-th one-hot vector is written as \(\mathbf{e}^i\). A vector is written in bold-face like \(\mathbf{a}\), and its \(i\)-th scalar element is written by non-bold \(a_i\). A matrix is written by capital bold-face like \(\mathbf{A}\), and \(\mathbf{A}_i\) denotes its \(i\)-th column vector unless otherwise noted. The identity matrix is denoted by \(\mathbf{I}\). The Hadamard product of \(\mathbf{A}\) and \(\mathbf{A}\) is written by \(\mathbf{A}^{\odot2} \mathrel{\vcenter{:}}=\mathbf{A}\odot \mathbf{A}\). We write the error function as \({\mathop{\mathrm{erf}}}\) and use \({\mathop{\mathrm{erf}}}(-z) = -{\mathop{\mathrm{erf}}}(z)\) without explicitly mentioning it. Infinitesimal asymptotic orders \(o(\cdot)\) are with respect to the sequence length \(T\) unless otherwise noted.

2.0.0.2 Transformer.

Let \(\mathbf{X}\mathrel{\vcenter{:}}=[\mathbf{x}_1 \; \mathbf{x}_2 \; \dots \; \mathbf{x}_T] \in \mathbb{R}^{d \times T}\) be an input with \(T\) tokens, defined later shortly. We suppose that all input sequences have the same length \(T\), and \(T\) is occasionally taken sufficiently large. The \(\ell\)-th (single-head) self-attention layer is defined as \[\begin{align} \tag{1} \mathbf{A}^\ell &\mathrel{\vcenter{:}}=\mathbf{S}\bigg(\frac{(\mathbf{X}^{\ell-1})^\top\mathbf{W}_{\mathrm{QK}}\mathbf{X}^{\ell-1}}{\lambda}\bigg), \\ \tag{2} \mathbf{U}^\ell &\mathrel{\vcenter{:}}=\mathbf{W}_{\mathrm{V}}\mathbf{X}^{\ell-1}\mathbf{A}^\ell, \end{align}\] where \(\mathbf{W}_{\mathrm{V}}\in \mathbb{R}^{d \times d}\) is the value parameters, \(\mathbf{W}_{\mathrm{QK}}(\mathrel{\vcenter{:}}=\mathbf{W}_{\mathrm{Q}}^\top\mathbf{W}_{\mathrm{K}}) \in \mathbb{R}^{d \times d}\) is the query-key parameters (with joint parametrization), \(\lambda > 0\) is temperature, commonly \(\lambda = \sqrt{d}\), and \(\mathbf{S}: \mathbb{R}^T \to \mathbb{R}^T\) is the softmax applied for each row. In this way, each input token in \(\mathbf{X}^{\ell-1}\) (embedded by \(\mathbf{W}_{\mathrm{V}}\)) is mixed by \(\mathbf{A}^\ell\). Then, the transformer block (without layer normalization [20]) is defined as \[\begin{align} & \mathbf{Z}^\ell \mathrel{\vcenter{:}}=\mathbf{U}^\ell + \mathbf{X}^{\ell-1}, \\ & \mathbf{H}^\ell \mathrel{\vcenter{:}}=\mathbf{W}_{\mathrm{F}_2}\sigma(\mathbf{W}_{\mathrm{F}_1}\mathbf{Z}^\ell), \\ & \mathbf{X}^\ell \mathrel{\vcenter{:}}=\mathbf{H}^\ell + \mathbf{Z}^\ell, \end{align}\] where \(\mathbf{H}^\ell\) is a feed-forward net with parameters \(\mathbf{W}_{\mathrm{F}_1}, \mathbf{W}_{\mathrm{F}_2} \in \mathbb{R}^{d \times d}\) and an (element-wise) activation \(\sigma: \mathbb{R}\to \mathbb{R}\). We omit the token embedding layer and set \(\mathbf{X}^0 \mathrel{\vcenter{:}}=\mathbf{X}\). There are two common variants of layer normalization positions, Post-LN [1] and Pre-LN [21], which are applied token-wise after the residual connections (\(\mathbf{Z}^\ell\) and \(\mathbf{X}^{\ell+1}\)) and before the inputs (\(\mathbf{X}^\ell\) and \(\mathbf{Z}^\ell\)), respectively. Then, the transformer block \(\mathbf{X}^\ell\) is stacked \(L\) times and \(\mathbf{F}(\mathbf{X}) \mathrel{\vcenter{:}}=\mathbf{X}^{L} \in \mathbb{R}^{d \times T}\) is the output.

2.0.0.3 Learning task.

We focus on causal language modeling, where a model predicts the next token given contextual tokens. Formally, given \(T\) contextual tokens \(\mathbf{X}\in \mathbb{R}^{d \times T}\), the prediction target is the \((T+1)\)-th token \(\mathbf{y}\mathrel{\vcenter{:}}=\mathbf{x}_{T+1} \in \mathbb{R}^d\). With the squared loss, the objective is written as follows: \[J(\boldsymbol{\Theta}) \mathrel{\vcenter{:}}=\frac{1}{2}\mathop{\mathrm{\mathbb{E}}}\|\mathbf{y}- \mathbf{F}(\mathbf{X})_T\|^2,\] where \(\boldsymbol{\Theta}\mathrel{\vcenter{:}}=(\mathbf{W}_{\mathrm{V}}, \mathbf{W}_{\mathrm{QK}}, \mathbf{W}_{\mathrm{F}_1}, \mathbf{W}_{\mathrm{F}_2})\) denotes the model parameter set, and the expectation is taken over input sequences \((\mathbf{X}, \mathbf{y})\). Here, our decoding procedure in consideration is to simply choose the embedded last token \(\mathbf{F}(\mathbf{X})_T \in \mathbb{R}^T\). The parameters \(\boldsymbol{\Theta}\) are learned by minimizing \(J\). Note that our analysis considers optimizing the query-key parameters jointly. Although such joint parameterization is less common in practice, it is convenient for the theoretical derivation of the gradients and has been used in several previous studies [11], [12]. Interested readers may refer to a recent work revealing that the joint and separate QK-parametrization lead to different implicit regularizations [22].

2.0.0.4 Picewise linear approximation of softmax.

In this article, we choose to approximate the softmax function \(\mathbf{S}\) by linearization. For a \(T\)-dimensional input \(\boldsymbol{\omega}\in \mathbb{R}^T\), the softmax function is defined as \[S(\boldsymbol{\omega})_i \mathrel{\vcenter{:}}=\frac{\exp(\omega_i)}{\sum_{j \in [T]} \exp(\omega_j)} \text{~~~for all } i \in [T].\] For linearization, the Taylor expansion of \(S(\boldsymbol{\omega})_i\) around the origin \(\left\langle{\boldsymbol{\gamma}^i},{\omega}\right\rangle + \gamma_0^i\) is used, where \[\boldsymbol{\gamma}^i \mathrel{\vcenter{:}}=\nabla_i\mathbf{S}(\mathbf{0}) = \frac{1}{T}\mathbf{e}^i - \frac{1}{T^2}\mathbf{1}, \quad \gamma_0^i \mathrel{\vcenter{:}}= S(\mathbf{0})_i = \frac{1}{T}.\] Then, we approximate \(\mathbf{S}\) by the piecewise linear function such that \(S(\boldsymbol{\omega})_i \approx \max\{0, \min\{1, \left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i\}\} = \left\langle{\widetilde{\boldsymbol{\gamma}}^i},{\boldsymbol{\omega}}\right\rangle + \widetilde{\gamma}_0^i\), where \[(\widetilde{\boldsymbol{\gamma}}^i, \widetilde{\gamma}_0^i) = \begin{cases} (\mathbf{0}, 0) & \text{if } \left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i < 0, \\ (\boldsymbol{\gamma}^i, \gamma_0^i) & \text{if } \left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i \in [0, 1], \\ (\mathbf{0}, 1) & \text{if } \left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i > 1. \\ \end{cases}\] In the vector form, the piecewise approximation \(\mathbf{S}(\boldsymbol{\omega}) \approx \widetilde{\mathbf{S}}(\boldsymbol{\omega})\) is given by \[\widetilde{\mathbf{S}}(\boldsymbol{\omega}) = \boldsymbol{\Gamma}^\top\boldsymbol{\omega}+ \widetilde{\boldsymbol{\gamma}}_0, \text{~~where~} \begin{cases} \boldsymbol{\Gamma}\mathrel{\vcenter{:}}=[\widetilde{\boldsymbol{\gamma}}^1 \; \widetilde{\boldsymbol{\gamma}}^2 \; \dots \widetilde{\boldsymbol{\gamma}}^T], \\ \widetilde{\boldsymbol{\gamma}}_0 = [\widetilde{\gamma}_0^1, \widetilde{\gamma}_0^2, \dots, \widetilde{\gamma}_0^T]^\top. \end{cases}\] For notational simplicity, the column vectors of \(\boldsymbol{\Gamma}\) are exceptionally denoted by \(\widetilde{\boldsymbol{\gamma}}^i\) with superscripts, for which the \(\alpha\)-th element is written by \(\widetilde{\gamma}^i_\alpha\). The difference between \(\mathbf{S}\) and \(\widetilde{\mathbf{S}}\) is illustrated in 1. Note that a popular alternative to softmax, sparsemax [23], is also a piecewise linear function, although the functional form is slightly different from \(\widetilde{\mathbf{S}}\).

3 Signal propagation probability↩︎

We analyze how much each token contributes to the learning dynamics. To this end, we formalize how much the signal of a specific input token \(\mathbf{x}_i\) propagates to the gradient \(\nabla J\). Remark that this notion is slightly different from the contribution of an input token \(\mathbf{x}_i\) to the model output \(\mathbf{F}(\mathbf{X})_j\) analyzed by [24] recently.

3.0.0.1 Uniform vs. localized softmax.

The piecewise linear approximation implies that the \(i\)-th input signal is propagated to the subsequent blocks when \(\left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i \in [0,1]\); otherwise, \(\widetilde{S}(\boldsymbol{\omega})_i = \left\langle{\widetilde{\boldsymbol{\gamma}}^i},{\boldsymbol{\omega}}\right\rangle + \widetilde{\gamma}_0^i = \widetilde{\gamma}_0^i\), which hinders the input token \(\mathbf{x}_i\) from contributing to the self-attention layer 2 . Thus, we will focus on the following quantity.

When only a few \(\rho_i\) are significantly larger than zero, we can interpret it as localized softmax; in this case, the self-attention 2 is dominated by a small number of tokens. By contrast, most of the tokens contribute to self-attention almost equally if \(\rho_i\) takes a similar value across different \(i\); this situation is interpreted as uniform softmax.

3.0.0.2 Through the lens of gradient.

The signal propagation probability naturally arises in the gradient. Since the learning dynamics of causal language modeling is governed by the gradient flow of \(J\), we can benefit from deriving the gradient of \(J\) to see how attention affects the learning dynamics. To keep the derivation concise, we consider a \(1\)-layer transformer (where we drop the superscripts \(\ell\)) without layer normalization and simplify the feed-forward net \(\mathbf{H}\) by supposing the identity activation. With the approximated softmax \(\widetilde{\mathbf{S}}\), the transformer can be written as follows: \[\begin{align} \mathbf{F}(\mathbf{X})_T &= \mathbf{W}_{\mathrm{F}}\{\mathbf{W}_{\mathrm{V}}\mathbf{X}\widetilde{\mathbf{S}}(\boldsymbol{\omega}) + \mathbf{x}_T\} \\ &= \mathbf{W}_{\mathrm{F}}\{\mathbf{W}_{\mathrm{V}}\mathbf{X}\boldsymbol{\Gamma}^\top\boldsymbol{\omega}+ \mathbf{W}_{\mathrm{V}}\mathbf{X}\widetilde{\boldsymbol{\gamma}}_0 + \mathbf{x}_T\}, \end{align}\] where \(\mathbf{W}_{\mathrm{F}}\mathrel{\vcenter{:}}=\mathbf{W}_{\mathrm{F}_2}\mathbf{W}_{\mathrm{F}_1} + \mathbf{I}\) and \(\boldsymbol{\omega}\mathrel{\vcenter{:}}=\mathbf{X}^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T / \lambda\). For this architecture, the QK-gradient is computed: \[\label{equation:qk95grad} \begin{align} \nabla_{\mathbf{W}_{\mathrm{QK}}}J &= \lambda^{-2}\mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{P}\boldsymbol{\Gamma}^\top\mathbf{X}^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T\mathbf{x}_T^\top] \\ &\phantom{=} + \lambda^{-1}\mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{P}\widetilde{\boldsymbol{\gamma}}_0\mathbf{x}_T^\top] + \lambda^{-1}\mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{q}\mathbf{x}_T^\top], \end{align}\tag{3}\] where \[\begin{align} \mathbf{P}&\mathrel{\vcenter{:}}=\mathbf{X}^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X}, \\ \mathbf{q}&\mathrel{\vcenter{:}}=\mathbf{X}^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top(\mathbf{W}_{\mathrm{F}}\mathbf{x}_T - \mathbf{y}). \end{align}\] When \(T\) is sufficiently large, we can drop asymptotically negligible terms with respect to \(T\) as detailed in 10, and the QK-gradient 3 is simplified as follows: \[\label{equation:qk95grad951} \frac{1}{\lambda^2} \!\! \sum_{i,j,\alpha,\beta \in [T]} \!\! \mathop{\mathrm{\mathbb{E}}}\left[ \widetilde{\gamma}^i_\alpha\widetilde{\gamma}^j_\beta (\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_j)(\mathbf{x}_\beta^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T)\mathbf{x}_\alpha\mathbf{x}_T^\top \right],\tag{4}\] where \(\check\mathbf{P}\mathrel{\vcenter{:}}=\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\).

Now, in the gradient term 4 , the summands with \(\widetilde{\gamma}^i_i\) are asymptotically dominant over those with \(\widetilde{\gamma}^i_\alpha\) (with \(\alpha \ne i\)) because \(\widetilde{\gamma}^i_i = T^{-1}\) and \(\widetilde{\gamma}^i_\alpha = o(T^{-1})\). Additionally, the \(i\)-th signal propagates when \(\widetilde{\gamma}^i_i > 0\), which holds iff \(\left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i \in [0,1]\) by definition. Therefore, we are motivated to check the condition \(\left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i \in [0,1]\) to see whether \(\mathbf{x}_i\) contributes to the gradient 3 . The signal propagation probability \(\rho_i\) characterizes its strength.

3.0.0.3 Summary.

In this section, we introduced the signal propagation probability \(\rho_i\), which characterizes how likely a given token \(\mathbf{x}_i\) contributes to the learning dynamics. Specifically, \(\widetilde{\boldsymbol{\gamma}}^i \ne \mathbf{0}\) holds more likely with larger \(\rho_i\), where \(\mathbf{x}_i\) contributes to the QK-gradient 3 . Subsequently, we will analyze the quantity \(\rho_i\) to see the behavior of the probability vector \(\boldsymbol{\rho}\in [0,1]^T\). When does the mass of \(\boldsymbol{\rho}\) concentrate to only a few tokens or scatter across most of the tokens?

4 When does attention localize?↩︎

We derive the signal propagation probability \(\rho_i\) based on the following synthetic data model for the sake of clarity.

To derive \(\rho_i\), we resort to the Gaussian approximation of \(\left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i\). Define \[\mu^i \mathrel{\vcenter{:}}=\mathop{\mathrm{\mathbb{E}}}[\left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i] \text{~~and~~} v^i \mathrel{\vcenter{:}}=\mathbb{V}[\left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i].\] We approximately suppose that \(\left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma_0^i \sim \mathcal{N}(\mu^i, v^i)\). Then, the signal propagation probability is approximated: \[\rho_i \approx \frac{1}{2}\left\{{\mathop{\mathrm{erf}}}\left(\frac{1-\mu^i}{\sqrt{2v^i}}\right) + {\mathop{\mathrm{erf}}}\left(\frac{\mu^i}{\sqrt{2v^i}}\right)\right\}.\] To leverage this formula, we derive \(\mu^i\) and \(v^i\).

The proof is given in 11. The symmetry of \(\mathbf{W}_{\mathrm{QK}}\) is assumed for convenience. In the case of asymmetric \(\mathbf{W}_{\mathrm{QK}}\), we can redefine the signal propagation probability with the symmetrized matrix \((\mathbf{W}_{\mathrm{QK}}+ \mathbf{W}_{\mathrm{QK}}^\top)/2\).

Recall that \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}) = \sum_{i \in [d]} w_i\) if we write the eigenvalues of \(\mathbf{W}\) by \((w_1, w_2, \dots, w_d)\), and that \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})^2 \le d\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)\) holds due to Jensen’s inequality. This implies that \[\label{equation:tr95tr295inequality} -\sqrt{d\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)} \le \mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}) \le \sqrt{d\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)}.\tag{5}\] Moreover, \(\mu^i\) and \(v^i\) are determined by the relative token location \(i/T\). By continuously extending \(i/T\) to \(\theta \in [0,1]\), the signal propagation probability \(\rho_i\) can be extended to \(\rho: [0,1] \to [0,1]\), defined over relative token locations: \[\label{equation:spp95approx} \rho(\theta) \mathrel{\vcenter{:}}=\Phi\bigg(\Big(\theta-\frac{1}{2}\Big)\xi; \theta\bigg) - \Phi\bigg(\Big(\theta-\frac{1}{2}\Big)\xi - \frac{1}{\eta}; \theta\bigg),\tag{6}\] where \[\begin{align} &\xi \mathrel{\vcenter{:}}=\frac{\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})}{\sqrt{\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)}}, \quad \eta \mathrel{\vcenter{:}}=\frac{\sqrt{\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)}}{\lambda}, \\ &\Phi(z;\theta) \mathrel{\vcenter{:}}=\frac{1}{2}{\mathop{\mathrm{erf}}}\bigg(\frac{z}{\sqrt{\smash[b]{2(2\theta^2+\frac{7}{12})}}}\bigg), \end{align}\] and the parameter ranges are \(\xi \in [-\sqrt{d}, \sqrt{d}]\) (due to the bound 5 ) and \(\eta \in (0,\infty)\). Here, \(\xi\) and \(\eta\) can be regarded independent (when \(\mathbf{W}\) is independent from \(\mathbf{X}\)) because the eigenspectrum scale \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)\) can be modulated within the bound 5 once the eigenspectrum of \(\mathbf{W}\) is given.

2 numerically illustrates \(\rho(\theta)\) with different \(\xi\) and \(\eta\). From these figures, we obtain a couple of observations.

• Localization. \(\rho(\theta)\) concentrates on fewer tokens as \(\eta\) increases (see \(|\xi| \ge 5\)). By contrast, \(\rho(\theta)\) behaves relatively uniformly regardless of \(\eta\) for small \(|\xi| \le 1\).

• Late-/middle-/early-site focus. Focus on small \(\eta\) such as \(\eta = 0.001\). As \(\xi\) increases to a large positive, \(\rho(\theta)\) puts positive weights for only late-site tokens, i.e., \(\theta > 0.5\). By contrast, as \(\xi\) decreases to a negative, \(\rho(\theta)\) focuses on early-site tokens, i.e., \(\theta < 0.5\). When \(\eta\) increases (see \(\eta \ge 0.5\)), \(\rho(\theta)\) localizes around \(\theta = 0.5\) with sufficiently large \(\xi\) (say, \(|\xi| \ge 5\)), which indicates middle-site focus.

• Vanishing signal. As \(\eta\) increases, \(\rho(\theta)\) degenerates to zero for any \(\theta \in [0,1]\) regardless of \(\xi\).

4.0.0.1 How \(\rho\) behaves at the limit.

Subsequently, we claim the above observations formally, which is proved in 11.

From late-/middle-/early-site focus in [lemma:spp95property], we see interestingly that \(\rho(\theta)\) localizes when \(\xi\eta = \mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})/\lambda\) asymptotically deviates from zero significantly so that \(|r| \gg 2\). At this limit, \(\rho(\theta)\) concentrates on \(\theta = 0.5\), inducing the middle-site focus. Conversely, attention becomes relatively uniform when \(\xi\eta = \mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})/\lambda\) is kept close to zero. In 3, we numerically illustrate this regime: \(\rho\) localizes at the middle site when \(\eta = 0.02\) (i.e., \(\xi\eta = 5.12\)).

Let us investigate the limiting condition of \((\xi,\eta)\) for localization: When is \((\xi,\eta)\) close to the limit \((\infty,0)\) while \(\xi\eta \to r \gg 2\)? Here, we focus on the eigenspectrum of \(\mathbf{W}\) by regarding its eigenvalues \((w_i)_{i\in[d]}\) as being sampled from a distribution with the mean \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}) = \sum_{i\in[d]}w_i\) and scale \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2) = \sum_{i\in[d]}w_i^2\). First, \(\eta \to 0\) indicates that the scale \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)\) should be close to zero. Next, \(\xi\eta (=\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})/\lambda) \to r \gg 2\) means that \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}) \gg 2\lambda\) at the limit, i.e., \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})\) should be significantly away from zero. By combining them, we tell that \(\rho\) localizes when the eigenspectrum concentrates around a non-zero mean. This happens more likely when the embedding dimension \(d\) is excessively large to make the eigenvalue sum \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})\) bounded away from zero while keeping the scale \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)\) close to zero (i.e., keeping every eigenvalue close to zero). Thus, a larger embedding dimension \(d\) is beneficial to drive attention to localization.

Next, from the claim of uniformity in [lemma:spp95property], we tell that \(\rho\) fluctuates less and less with \(\xi\) closer to zero. Hence, \(\rho(\theta)\) takes a similar value across different token positions \(\theta\) in this limit. When \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}) \to 0\), \(\rho\) attains this limit.

4.0.0.2 Summary.

Wrapping up this section, we obtain answers to (Q1) posed in 1 under the random walk model.

5 Are different collapse regimes reconcilable?↩︎

We discuss the results of our analysis in 4 and the previous arguments related to attention uniformity.

5.0.0.1 Connection to rank collapse.

[16] showed that self-attention blocks \(\mathbf{U}^\ell\) (see 2 ) converges to a rank-\(1\) matrix \(\mathbf{z}\mathbf{1}^\top\) (for some \(\mathbf{z}\)) with \(L \to \infty\) without skip connections or feed-forward blocks, which is called rank collapse.¹ They argued the importance of avoiding rank collapse for better expressivity because each token embedding in a rank-\(1\) self-attention block degenerates to the same. Hence, the rank collapse is related to the failure mode attributed to the uniformity after mixing key tokens by attention, which is slightly different from what we are concerned about—how each token contributes during mixing by attention (through the gradient, as discussed in 3).

[16] proved that the convergence rate to a rank-\(1\) matrix slows down when the matrix \(\ell_1\)-norm \(\|\mathbf{W}_{\mathrm{QK}}\|_1\) is large. Because we can draw the following connection between \(\|\mathbf{W}_{\mathrm{QK}}\|_1\) and \(|\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})|\): \[\label{equation:l195norm95lower95bound} \frac{|\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})|}{\sqrt{d}\|\boldsymbol{\Sigma}\|_2} \le \|\mathbf{W}_{\mathrm{QK}}\|_2 \le \|\mathbf{W}_{\mathrm{QK}}\|_{\mathrm{F}} \le \|\mathbf{W}_{\mathrm{QK}}\|_1,\tag{7}\] where the first inequality is due to the bound 5 and the Cauchy–Schwarz inequality, it is advisable to increase \(|\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})|\) under fixed \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)\) to mitigate the rank collapse. This is equivalent to reducing the eigenspectrum variance \[\label{equation:eigenspectrum95variance} d^2\mathbb{V}[w_i] = d^2(\mathop{\mathrm{\mathbb{E}}}[w_i^2] - \mathop{\mathrm{\mathbb{E}}}[w_i]^2) = d\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2) - |\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})|^2.\tag{8}\] Hence, minimizing the \(\mathbf{W}\)-eigenspectrum variance leads to better expressivity.

5.0.0.2 Connection to entropy collapse.

[18] introduced a concept called entropy collapse, in which the average Shannon entropy of the columns of the attention matrix \(\mathbf{A}^\ell\) (see 1 ) shrinks. Intuitively speaking, low attention entropy induces localized attention. This notion of localization is akin to ours because the attention entropy measures the uniformity the attention is applied to input tokens during mixing. They empirically observed that the training loss falls into a plateau with low attention entropy, which causes training instability of transformers, and hence advocate for keeping attention less peaked during training.

In [18], the attention entropy is asymptotically lower-bounded for large \(T\) by \[\label{equation:entropy95lower95bound} \ln(1+T\exp(-\nu)) + \frac{\nu\exp(-\nu/2)}{T^{-1} + \exp(-\nu)},\tag{9}\] where \(\nu \mathrel{\vcenter{:}}=\|\mathbf{X}\mathbf{X}^\top\|_2\|\mathbf{W}_{\mathrm{QK}}\|_2\). This lower bound is unimodal in \(\nu\) and vanishes at \(\|\mathbf{W}_{\mathrm{QK}}\|_2 \to \infty\) (see 4), so the attention entropy tends to be higher when \(\|\mathbf{W}_{\mathrm{QK}}\|_2\) is small. If \(|\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})|\) is not too small, \(\|\mathbf{W}_{\mathrm{QK}}\|_2\) is lower-bounded (see 7 ) and the attention entropy may be kept close to the peak of the lower bound 9 . To mitigate the entropy collapse, it is natural to decrease \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)\) under fixed \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})\) (which is equivalent to minimizing the eigenspectrum variance by 8 ) because of the bound \[\label{equation:spectral95norm95upper95bound} \|\boldsymbol{\Sigma}^{-1}\|_{\mathrm{F}}\sqrt{\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)} \ge |\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}})| \ge \|\mathbf{W}_{\mathrm{QK}}\|_2,\tag{10}\] where the first inequality is due to the Cauchy–Schwarz inequality. Hence, minimizing the \(\mathbf{W}\)-eigenspectrum variance helps the model to avoid the entropy collapse.

5.0.0.3 Rank collapse vs. entropy collapse.

At first sight, the two notions of collapse seem to contradict each other because avoiding the rank collapse leads to diverse token embeddings, whereas avoiding the entropy collapse leads to a uniform token mixer. Indeed, the matrix \(\ell_1\)-norm (that decreases under the rank collapse) and the spectral norm (that increases under the entropy collapse) are equivalent norms, and the two modes appear to be incompatible.

However, as we discussed above, this trade-off is reconcilable from the viewpoint of the \(\mathbf{W}\)-eigenspectrum. Setting the eigenspectrum mean to be bounded away from zero, we can avoid the rank collapse owing to the bound 7 . Under a fixed eigenspectrum mean, minimizing the eigenspectrum scale (equivalently, minimizing its variance 8 ) leads to high attention entropy due to the bound 10 and the unimodal shape of the entropy lower bound 9 . This variance minimization is nothing else but the condition of attention localization. Eventually, \(\rho(\theta)\) localizes and attends to specific sites of tokens, as we showed in [lemma:spp95property]. Hence, the signal propagation probability offers us a better view of localization. Let us summarize our second take-home.

5.0.0.4 Numerical simulation.

To see the relationship between the eigenspectrum, \(\rho\), and attention entropy, we simulated the signal propagation probability \(\rho_i\) using synthesized random walks following 1 with the isotropic \(\boldsymbol{\Sigma}= \mathbf{I}\) and anisotropic \(\boldsymbol{\Sigma}\). To obtain an anisotropic \(\boldsymbol{\Sigma}\), we first sampled \(\mathbf{R}\in \mathbb{R}^{d \times d}\) from element-wise \(\mathrm{Unif}(-2.5,2.5)\) and computed \(\boldsymbol{\Sigma}= \mathbf{R}^\top\mathbf{R}/d\). We sampled \(300\) sequences with \(40\) tokens, and obtained \(\mathbf{W}_{\mathrm{QK}}\) by generating \(128\) eigenvalues \((w_i)_{i \in [d]}\) from \(\mathcal{N}(\text{mean}, \text{scale}^2)\) and composed with a sampled orthogonal basis matrix \(\mathbf{B}\), by the eigendecomposition formula \(\mathbf{W}_{\mathrm{QK}}= \mathbf{B}\mathop{\mathrm{\mathrm{diag}}}((w_i)_{i})\mathbf{B}^\top\). The signal propagation probability was averaged over \(300\) sequences.

5 shows the averaged \(\rho_i\) with different mean-scale pairs and the corresponding attention entropy in the right-most figure. As seen, \(\rho_i\) localizes with smaller scales and larger means, which is consistent with the conclusion in 4. This trend supports the validity of \(\mathbf{W}_{\mathrm{QK}}\)-eigenspectrum as a proxy to \(\mathbf{W}\)-eigenspectrum. Moreover, we observe that \(\mathbf{W}_{\mathrm{QK}}\)-eigenspectrum with a fixed mean and scale leads to higher attention entropy.

6 Intervening attention localization↩︎

To empirically see the impact of localization on the model performance, we propose a method to control the degree of attention localization. We focus on the eigenspectrum of \(\mathbf{W}_{\mathrm{QK}}\) instead of \(\mathbf{W}= \mathbf{W}_{\mathrm{QK}}\boldsymbol{\Sigma}\) because \(\boldsymbol{\Sigma}\) does not change during training, and the numerical simulation showed that \(\mathbf{W}_{\mathrm{QK}}\) behaves as a reasonable proxy to \(\mathbf{W}\) (see 5).

We minimize the loss function \(J\) while minimizing the eigenspectrum scale and maintaining the mean to a fixed level: \[\min_{\boldsymbol{\Theta}} \Big\{ J(\boldsymbol{\Theta}) + \kappa_1\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}}^\top\mathbf{W}_{\mathrm{QK}}) + \kappa_2(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}}) - 1)^2 \Big\},\] where \(\kappa_1, \kappa_2 > 0\) are the regularization strengths. Here, we allow \(\mathbf{W}_{\mathrm{QK}}\) to be asymmetric, and the eigenspectrum scale is represented by \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}}^\top\mathbf{W}_{\mathrm{QK}})\). The regularization terms can be optimized fairly easily thanks to the following derivative formulae (for \(\mathbf{W}_{\mathrm{QK}}= \mathbf{W}_{\mathrm{Q}}^\top\mathbf{W}_{\mathrm{K}}\)): \[\begin{align} &\nabla_{\mathbf{W}_{\mathrm{Q}}}\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}}^\top\mathbf{W}_{\mathrm{QK}}) = 2\mathbf{W}_{\mathrm{K}}\mathbf{W}_{\mathrm{K}}^\top\mathbf{W}_{\mathrm{Q}}, \\ &\nabla_{\mathbf{W}_{\mathrm{K}}}\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}}^\top\mathbf{W}_{\mathrm{QK}}) = 2\mathbf{W}_{\mathrm{Q}}\mathbf{W}_{\mathrm{Q}}^\top\mathbf{W}_{\mathrm{K}}, \\ &\nabla_{\mathbf{W}_{\mathrm{Q}}}[(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}})-1)^2] \!=\! [\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}})-1]\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}})\mathbf{W}_{\mathrm{K}}, \\ &\nabla_{\mathbf{W}_{\mathrm{K}}}[(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}})-1)^2] \!=\! [\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}})-1]\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}})\mathbf{W}_{\mathrm{Q}}. \end{align}\] Since this whole objective drives the eigenspectrum scale to a small value, the signal vanishing can be avoided automatically. We call this regularization scheme LocAteR (LOCalized ATtEntion Regularization).

7 Experiments↩︎

We aim to observe the correlation between the eigenspectrum and localization. To this end, we train transformers with LocAteR and varying \(\kappa_1,\kappa_2\), and see how the model performances and attention foci change over time.

7.0.0.1 Setup.

We used fairseq v0.12.2 [25], which is a toolkit oriented for sequence modeling, to implement and train transformers. The basic training scheme was inherited from fairseq-cli/train.py. The model is a \(1\)-layer transformer with a single-head self-attention and Post-LN (default), and the input embedding dimension, attention embedding dimension, and feed-forward net embedding dimension are set to \(128\) altogether (namely, \(d=128\)).² Input data were transformed into \(64\) tokens (namely, \(T=64\)) with batch size \(64\). The optimizer is Adam [26] with default parameters and no clip norm, and the weight decay with \(0.01\) is used. The learning rate is fixed to \(2.5\times10^{-5}\) without any scheduling. The FP16 quantizer was applied to reduce memory usage. All the other configs remain to be the same as the default in fairseq-cli/train.py. Under this config, we updated the model with \(50000\) iters.

7.0.0.2 Language modeling.

We conducted the language modeling task. The dataset we used is WikiText-2 [19], which is a collection of high-quality Wikipedia articles. We conduced the experiments with fixed scale regularization strength \(\kappa_1=100\) and varying mean regularization strength \(\kappa_2\) from \(0, 10^{-3}, 10^{-2}, \dots, 10^{0}\).

The results are shown in 6, in which stronger regularizers tend to make the eigenspectrum scale smaller. This, in turn, maintains the attention entropy higher during the updates entirely, and eventually, the model achieves better perplexity. While the better model performance with higher attention entropy has already been observed by [18], we also showed that a smaller eigenspectrum scale contributes to higher attention entropy. This empirically corroborates that attention localization leads to better model performance, probably because the attention mechanism appropriately selects relevant tokens during training.

7 shows the signal propagation probability at each training iteration. We compute the signal propagation probability of token \(i\) by counting the frequency of \(\left\langle{\boldsymbol{\gamma}^i},{\boldsymbol{\omega}}\right\rangle + \gamma^i_0 \in [0,1]\) in a given batch. LocAteR entails salient horizontal stripes, each corresponding to attended tokens. Yet, the stripes do not appear in “bulk” as we analyzed in 4 because our synthetic data model in 1 does not perfectly align with real datasets. Nevertheless, our experiments evidently contrast the localized and uniformed attention depending on the eigenspectrum scale because no salient stripes are observed without LocAteR.

In 6 7, we observe different learning phases for the first \(10^4\) and the rest iters. Indeed, [27] observed similar phenomena and explained that it is due to the different convergence speeds between attention weights corresponding to informative and non-informative tokens. The relationship between the \(\mathbf{W}_{\mathrm{QK}}\)-eigenspectrum and this dynamics is beyond our scope and left for future work.

8 Conclusion and limitation↩︎

We revealed that attention localizes when the eigenspectrum of \(\mathbf{W}\) concentrates to a non-zero mean, or equivalently, with larger eigenspectrum mean \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W})\) and smaller scale \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}^2)\). Based on it, LocAteR was proposed to shrink the scale \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}}^2)\) while maintaining the mean \(\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}})\). Interestingly, maximizing the scale is related to mitigating both rank collapse and entropy collapse, and hence, the two apparently contradictory failure modes can be reconciled. The experiments on a real-world dataset corroborate it, though the random walk model is not perfectly satisfied.

We recognize three limitations of this work. First, we rely on the strong random walk model. Although the Gaussianity may be reasonable because of usual initialization schemes of transformer embedding layers, it is interesting to consider an alternative model to capture token correlations better. Second, the formal analysis is mainly restricted to \(1\)-layer transformers. Recent studies often consider gradient explosion in the large-depth limit from the viewpoint of layer normalization [21], [28] and initialization [29], [30]. It must be fruitful to integrate these perspectives to our gradient analysis through 3 . Lastly, why attention localization leads to better model performance still remains elusive. Whereas localization is related to avoiding rank collapse (and hence higher model expressivity), we need additional effort to fully understand the mechanism.

Acknowledgments↩︎

A part of the experiments of this research was conducted using Wisteria/Aquarius in the Information Technology Center, the University of Tokyo.

9 Helper lemmas↩︎

The formulae in 1 are standard and cropped from [31].

The formulae in 2 can be shown by recursively applying 1. We omit the proofs since they are elementary.

10 Omitted derivations↩︎

10.1 QK-gradient↩︎

Here, we complement the derivations of the QK-gradient terms shown in 3. To get 3 , we compute \(\nabla_{\mathbf{W}_{\mathrm{QK}}}J\): \[\begin{align} \nabla_{\mathbf{W}_{\mathrm{QK}}}J &= \frac{1}{2}\mathop{\mathrm{\mathbb{E}}}[\nabla_{\mathbf{W}_{\mathrm{QK}}}\|\mathbf{y}- \mathbf{F}(\mathbf{X})_T\|^2] \\ &= \frac{1}{2}\mathop{\mathrm{\mathbb{E}}}[\nabla_{\mathbf{W}_{\mathrm{QK}}}\{\mathbf{F}(\mathbf{X})_T^\top\mathbf{F}(\mathbf{X})_T - 2\mathbf{y}^\top\mathbf{F}(\mathbf{X})_T\}] \\ &= \frac{1}{2}\mathop{\mathrm{\mathbb{E}}}[\boldsymbol{\omega}^\top\boldsymbol{\Gamma}\mathbf{X}^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X}\boldsymbol{\Gamma}^\top\boldsymbol{\omega}+ 2(\mathbf{W}_{\mathrm{V}}\mathbf{X}\widetilde{\boldsymbol{\gamma}}_0 + \mathbf{x}_T)^\top\mathbf{W}_{\mathrm{F}}^\top\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X}\boldsymbol{\Gamma}^\top\boldsymbol{\omega}- 2\mathbf{y}^\top\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X}\boldsymbol{\Gamma}^\top\boldsymbol{\omega}] \\ &= \frac{1}{2\lambda^2}\mathop{\mathrm{\mathbb{E}}}[\mathbf{x}_T^\top\mathbf{W}_{\mathrm{QK}}^\top\mathbf{X}\boldsymbol{\Gamma}\mathbf{P}\boldsymbol{\Gamma}^\top\mathbf{X}^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T] + \frac{1}{\lambda}\mathop{\mathrm{\mathbb{E}}}[(\widetilde{\boldsymbol{\gamma}}_0)^\top\mathbf{P}\boldsymbol{\Gamma}^\top\mathbf{X}^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T] + \frac{1}{\lambda}\mathop{\mathrm{\mathbb{E}}}[\mathbf{q}^\top\boldsymbol{\Gamma}^\top\mathbf{X}^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T] \\ &= \frac{1}{\lambda^2}\mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{P}\boldsymbol{\Gamma}^\top\mathbf{X}^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T\mathbf{x}_T^\top] + \frac{1}{\lambda}\mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{P}\widetilde{\boldsymbol{\gamma}}_0\mathbf{x}_T^\top] + \frac{1}{\lambda}\mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{q}\mathbf{x}_T^\top]. \end{align}\] By expanding the first term of \(\nabla_{\mathbf{W}_{\mathrm{QK}}}J\), we get the following: \[\label{equation:qk95grad95195} \frac{1}{\lambda^2} \!\! \sum_{i,j,\alpha,\beta \in [T]} \!\! \mathop{\mathrm{\mathbb{E}}}\left[ \widetilde{\gamma}^i_\alpha\widetilde{\gamma}^j_\beta (\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_j)(\mathbf{x}_\beta^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T)\mathbf{x}_\alpha\mathbf{x}_T^\top \right],\tag{11}\] where \(\check\mathbf{P}\mathrel{\vcenter{:}}=\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\). Similarly, by expanding the second and third terms of \(\nabla_{\mathbf{W}_{\mathrm{QK}}}J\), we get the following terms, respectively: \[\begin{align} \tag{12} &\frac{1}{\lambda} \sum_{i,\alpha,\beta \in [T]} \mathop{\mathrm{\mathbb{E}}}\left[\widetilde{\gamma}^i_\alpha\widetilde{\gamma}_0^\beta (\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_\beta)\mathbf{x}_\alpha\mathbf{x}_T^\top\right], \\ \tag{13} &\frac{1}{\lambda} \sum_{i,\alpha \in [T]} \mathop{\mathrm{\mathbb{E}}}\left[\widetilde{\gamma}^i_\alpha \{\mathbf{x}_i^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top(\mathbf{W}_{\mathrm{F}}\mathbf{x}_T - \mathbf{y})\}\mathbf{x}_\alpha\mathbf{x}_T^\top\right]. \end{align}\]

To get 11 , we expand the first term of \(\nabla_{\mathbf{W}_{\mathrm{QK}}}J\): \[\begin{align} \mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{P}\boldsymbol{\Gamma}^\top\mathbf{X}^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T\mathbf{x}_T^\top] &= \mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{X}^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X}\boldsymbol{\Gamma}^\top\mathbf{X}^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T\mathbf{x}_T^\top] \\ &= \mathop{\mathrm{\mathbb{E}}}\left[\{(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X})(\mathbf{X}\boldsymbol{\Gamma})^\top\}^\top\{(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X})(\mathbf{X}\boldsymbol{\Gamma})^\top\}\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T\mathbf{x}_T^\top\right] \\ &= \mathop{\mathrm{\mathbb{E}}}\bigg[\bigg\{\sum_{i \in [T]}(\mathbf{X}\widetilde{\boldsymbol{\gamma}}^i)(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{x}_i)^\top\bigg\} \bigg\{\sum_{j \in [T]}(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{x}_j)(\mathbf{X}\widetilde{\boldsymbol{\gamma}}^j)^\top\bigg\}\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T\mathbf{x}_T^\top\bigg] \\ &= \mathop{\mathrm{\mathbb{E}}}\bigg[\sum_{i,j\in[T]} (\mathbf{X}\widetilde{\boldsymbol{\gamma}}^i)\{(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{x}_i)^\top(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{x}_j)\}(\mathbf{X}\widetilde{\boldsymbol{\gamma}}^j)^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T\mathbf{x}_T^\top\bigg] \\ &= \mathop{\mathrm{\mathbb{E}}}\bigg[\sum_{\alpha,\beta\in[T]}\bigg\{\sum_{i,j\in[T]}\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_j\bigg\}\widetilde{\gamma}^i_\alpha\widetilde{\gamma}^j_\beta\mathbf{x}_\alpha\mathbf{x}_\beta^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T\mathbf{x}_T^\top\bigg] \\ &= \sum_{i,j,\alpha,\beta\in[T]} \mathop{\mathrm{\mathbb{E}}}[\widetilde{\gamma}^i_\alpha\widetilde{\gamma}^j_\beta(\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_j)(\mathbf{x}_\beta^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T)\mathbf{x}_\alpha\mathbf{x}_T^\top]. \end{align}\] To get 12 , we expand the second term of \(\nabla_{\mathbf{W}_{\mathrm{QK}}}J\): \[\begin{align} \mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{P}\widetilde{\boldsymbol{\gamma}}_0\mathbf{x}_T^\top] &= \mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{X}^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X}\widetilde{\boldsymbol{\gamma}}_0\mathbf{x}_T^\top] \\ &= \mathop{\mathrm{\mathbb{E}}}[\{(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X})(\mathbf{X}\boldsymbol{\Gamma})^\top\}^\top\{(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X})(\widetilde{\boldsymbol{\gamma}}_0\mathbf{x}_T^\top)\}] \\ &= \mathop{\mathrm{\mathbb{E}}}\bigg[\bigg\{\sum_{i\in[T]}(\mathbf{X}\widetilde{\boldsymbol{\gamma}}^i)(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{x}_i)^\top\bigg\} \bigg\{\sum_{\beta\in[T]}(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{x}_\beta)(\widetilde{\gamma}_0^\beta\mathbf{x}_T^\top)\bigg\}\bigg] \\ &= \mathop{\mathrm{\mathbb{E}}}\bigg[\sum_{i,\beta\in[T]}(\mathbf{X}\widetilde{\boldsymbol{\gamma}}^i)\{(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{x}_i)^\top(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{x}_\beta)\}(\widetilde{\gamma}_0^\beta\mathbf{x}_T^\top)\bigg] \\ &= \mathop{\mathrm{\mathbb{E}}}\bigg[\sum_{\alpha\in[T]}\bigg\{\sum_{i,\beta\in[T]}\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_\beta\bigg\}\widetilde{\gamma}^i_\alpha\widetilde{\gamma}_0^\beta\mathbf{x}_\alpha\mathbf{x}_T^\top\bigg] \\ &= \sum_{i,\alpha,\beta\in[T]} \mathop{\mathrm{\mathbb{E}}}[\widetilde{\gamma}^i_\alpha\widetilde{\gamma}_0^\beta(\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_\beta)\mathbf{x}_\alpha\mathbf{x}_T^\top]. \end{align}\] To get 13 , we expand the third term of \(\nabla_{\mathbf{W}_{\mathrm{QK}}}J\): \[\begin{align} \mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{q}\mathbf{x}_T^\top] &= \mathop{\mathrm{\mathbb{E}}}[\mathbf{X}\boldsymbol{\Gamma}\mathbf{X}^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top(\mathbf{W}_{\mathrm{F}}\mathbf{x}_T - \mathbf{y})\mathbf{x}_T^\top] \\ &= \mathop{\mathrm{\mathbb{E}}}[\{(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{X})(\mathbf{X}\boldsymbol{\Gamma})^\top\}^\top(\mathbf{W}_{\mathrm{F}}\mathbf{x}_T - \mathbf{y})\mathbf{x}_T^\top] \\ &= \mathop{\mathrm{\mathbb{E}}}\bigg[\bigg\{\sum_{i\in[T]}(\mathbf{X}\widetilde{\boldsymbol{\gamma}}^i)(\mathbf{W}_{\mathrm{F}}\mathbf{W}_{\mathrm{V}}\mathbf{x}_i)^\top\bigg\}(\mathbf{W}_{\mathrm{F}}\mathbf{x}_T - \mathbf{y})\mathbf{x}_T^\top\bigg] \\ &= \mathop{\mathrm{\mathbb{E}}}\bigg[\bigg\{\sum_{i,\alpha\in[T]}\widetilde{\gamma}^i_\alpha\mathbf{x}_\alpha\mathbf{x}_i^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top\bigg\}(\mathbf{W}_{\mathrm{F}}\mathbf{x}_T - \mathbf{y})\mathbf{x}_T^\top\bigg] \\ &= \sum_{i,\alpha\in[T]} \mathop{\mathrm{\mathbb{E}}}[\widetilde{\gamma}^i_\alpha\{\mathbf{x}_i^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top(\mathbf{W}_{\mathrm{F}}\mathbf{x}_T - \mathbf{y})\}\mathbf{x}_\alpha\mathbf{x}_T^\top]. \end{align}\]

10.2 Order evaluation of QK-gradient↩︎

The orders of the QK-gradient terms 11 , 12 , and 13 are evaluated. In this subsection, we assume that the covariance matrix in 1 is \(\boldsymbol{\Sigma}= \mathbf{I}\) for simplicity. The following evaluation still applies with minor modifications for a general \(\boldsymbol{\Sigma}\). For 11 , the Cauchy–Schwarz inequality implies that \[\begin{align} \text{\ref{equation:qk95grad95195}} &= \bigg|\sum_{i,j,\alpha,\beta} \mathop{\mathrm{\mathbb{E}}}[\widetilde{\gamma}^i_\alpha\widetilde{\gamma}^j_\beta(\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_j)(\mathbf{x}_\beta^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T)\mathbf{x}_\alpha\mathbf{x}_T^\top]\bigg|^2 \\ &\le \underbrace{\bigg\{\sum_{i,j,\alpha,\beta} \mathop{\mathrm{\mathbb{E}}}[(\widetilde{\gamma}^i_\alpha\widetilde{\gamma}^j_\beta)^2]\bigg\}}_\text{(A)} \underbrace{\bigg\{\sum_{i,j,\alpha,\beta} \mathop{\mathrm{\mathbb{E}}}[(\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_j)^2]\bigg\}}_\text{(B)} \underbrace{\bigg\{\sum_{i,j,\alpha,\beta} \mathop{\mathrm{\mathbb{E}}}[(\mathbf{x}_\beta^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T)^2]\bigg\}}_\text{(C)} \underbrace{\bigg\{\sum_{i,j,\alpha,\beta} \mathop{\mathrm{\mathbb{E}}}[(\mathbf{x}_\alpha\mathbf{x}_T^\top)^{\odot2}]\bigg\}}_\text{(D)}. \end{align}\] For (A), we have \(\widetilde{\gamma}^i_\alpha, \widetilde{\gamma}^j_\beta \le T^{-1}\) by definition, and hence \(\text{(A)} = O(1)\). For (B), by using ?? , \[\begin{align} \text{(B)} = T^2\sum_{i,j} \mathop{\mathrm{\mathbb{E}}}[\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_j\mathbf{x}_j^\top\check\mathbf{P}\mathbf{x}_i] = T^2\sum_{i,j} \{(i^2+ij-3i-j+4)\mathop{\mathrm{\mathrm{tr}}}(\check\mathbf{P}^2) + (i^2-2i+2)\mathop{\mathrm{\mathrm{tr}}}(\check\mathbf{P})^2\} = O(T^6). \end{align}\] For (C), by following the same computation as 14 , \[\begin{align} \text{(C)} &= T^3\sum_\beta \mathop{\mathrm{\mathbb{E}}}[\mathbf{x}_\beta^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_T\mathbf{x}_T^\top\mathbf{W}_{\mathrm{QK}}\mathbf{x}_\beta] \\ &= T^3\sum_\beta \{(\beta^2+(T-3)\beta-(T-4))\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}}^2) + (\beta^2-2\beta+2)\mathop{\mathrm{\mathrm{tr}}}(\mathbf{W}_{\mathrm{QK}})^2\} = O(T^6). \end{align}\] For (D), its \((i,j)\)-element can be evaluated as follows (no matter whether \(i=j\) or not): \[\begin{align} \text{(D)}_{ij} = T^3\sum_\alpha \mathop{\mathrm{\mathbb{E}}}[x_{\alpha,i}^2x_{T,j}^2] = T^3\sum_\alpha \mathop{\mathrm{\mathbb{E}}}[x_{\alpha,i}^2\{(T-\alpha) + x_{\alpha,j}^2\}] = O(T^5)\sum_\alpha \mathop{\mathrm{\mathbb{E}}}[x_{\alpha,i}^2] + T^3\sum_{\alpha} \mathop{\mathrm{\mathbb{E}}}[x_{\alpha,i}^2x_{\alpha,j}^2] = O(T^6). \end{align}\] By plugging them back, we now confirmed that \(|\text{\ref{equation:qk95grad95195}}| = O(T^8)\).

The orders of 12 13 can be evaluated similarly and the detailed evaluations are omitted. \[\begin{align} |\text{\ref{equation:qk95grad952}}|^2 &= \bigg|\sum_{i,\alpha,\beta} \mathop{\mathrm{\mathbb{E}}}[\widetilde{\gamma}^i_\alpha\widetilde{\gamma}_0^\beta(\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_\beta)\mathbf{x}_\alpha\mathbf{x}_T^\top]\bigg|^2 \\ &\le \bigg\{\sum_{i,\alpha,\beta} \mathop{\mathrm{\mathbb{E}}}[(\widetilde{\gamma}^i_\alpha\widetilde{\gamma}_0^\beta)^2]\bigg\} \bigg\{\sum_{i,\alpha,\beta} \mathop{\mathrm{\mathbb{E}}}[(\mathbf{x}_i^\top\check\mathbf{P}\mathbf{x}_\beta)^2]\bigg\} \bigg\{\sum_{i,\alpha,\beta} \mathop{\mathrm{\mathbb{E}}}[(\mathbf{x}_\alpha\mathbf{x}_T^\top)^{\odot2}]\bigg\} \\ &= O(T) \cdot O(T^4) \cdot O(T^5) \\ &= O(T^{10}) \\ \implies |\text{\ref{equation:qk95grad952}}| &= O(T^5). \\ |\text{\ref{equation:qk95grad953}}|^2 &= \bigg|\sum_{i,\alpha} \mathop{\mathrm{\mathbb{E}}}[\widetilde{\gamma}^i_\alpha\{\mathbf{x}_i^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top(\mathbf{W}_{\mathrm{F}}\mathbf{x}_T-\mathbf{y})\}\mathbf{x}_\alpha\mathbf{x}_T^\top]\bigg|^2 \\ &\le \bigg\{\sum_{i,\alpha} \mathop{\mathrm{\mathbb{E}}}[(\widetilde{\gamma}^i_\alpha)^2]\bigg\} \bigg\{\sum_{i,\alpha} \mathop{\mathrm{\mathbb{E}}}[(\mathbf{x}_i^\top\mathbf{W}_{\mathrm{V}}^\top\mathbf{W}_{\mathrm{F}}^\top(\mathbf{W}_{\mathrm{F}}\mathbf{x}_T-\mathbf{y}))^2]\bigg\} \bigg\{\sum_{i,\alpha} \mathop{\mathrm{\mathbb{E}}}[(\mathbf{x}_\alpha\mathbf{x}_T^\top)^{\odot2}]\bigg\} \\ &= O(1) \cdot O(T^4) \cdot O(T^4) \\ &= O(T^8) \\ \implies |\text{\ref{equation:qk95grad953}}| &= O(T^4). \end{align}\]

Hence, we have \(|\text{\ref{equation:qk95grad95195}}| = O(T^8)\), \(|\text{\ref{equation:qk95grad952}}| = O(T^5)\), and \(|\text{\ref{equation:qk95grad953}}| = O(T^4)\), which implies that the QK-gradient 3 is asymptotically dominated by 11 .

11 Proofs↩︎

12 Additional experiments↩︎

Here, we show additional results of the language modeling task with \(1\)-/\(3\)-/\(6\)-layer transformers with different embedding dimensions \(d=32,128\). For \(d=128\), the configurations remain the same except for the number of decoder layers as in 7. For \(d=32\), we used the learning rate \(0.0001\) (instead of \(0.000025\) used for \(d=128\)), and the other configurations remain the same. The results are shown in 8 (\(d=32\), \(1\)-layers), 9 (\(d=32\), \(3\)-layers), 10 (\(d=32\), \(6\)-layers), 11 (\(d=128\), \(3\)-layers), and 12 (\(d=128\), \(6\)-layers). The overall trends are quite similar to the case of \(1\)-layer transformers with \(d=128\) as seen in 6: As \(\kappa_2\) increases, the eigenspectrum scale decreases, the attention entropy increases, and eventually, the perplexity improves (namely, decreases).

References↩︎