Attention is not all you need:
pure attention loses rank doubly exponentially with depth

Let $\bm{X}$ be a $n\times d_{\textit{in}}$ input matrix consisting of $n$ tokens. A SAN is built out of $L$ multi-head self-attention layers, each having $H$ heads. The output of the $h$-th self-attention head can be written as

Above, $\bm{W}_{V,h}$ is a $d_{\textit{in}}\times d_{\textit{v}}$ value weight matrix and the $n\times n$ row-stochastic matrix $\bm{P}_{h}$ is given by

where (1) the key and query weight matrices $\bm{W}_{K,h}$ and $\bm{W}_{Q,h}$ are of size $d_{\textit{in}}\times d_{\textit{qk}}$, (2) $\bm{W}_{QK,h}=\bm{W}_{Q,h}\bm{W}_{K,h}^{\top}$, and (3) the softmax operates independently on each row of its input. We obtain the final equation by noting that softmax is shift-invariant and disregarding terms that provide a constant contribution across rows [CLJ20a].

The output of each SAN layer is formed by concatenating the individual outputs of all $H$ attention heads (along the last dimension) and linearly projecting them onto a subspace of appropriate size:

where we set $\bm{W}_{h}=\bm{W}_{V,h}\,\bm{W}_{O,h}^{\top}$ and $\bm{b}_{O}=\sum_{h}\bm{b}_{O,h}$.

Let $\bm{X}^{l}$ be the output of the $l$-th layer and fix $\bm{X}^{0}=\bm{X}$. As is common practice, we let all layers consist of the same number of heads.

Excluding biases $\text{{1}}\bm{b}_{O,h}^{\top}$, the SAN output is given by

which, after unrolling the recursion backwards, yields:

The above equations have a clear interpretation if we think of the SAN as a directed acyclic graph, with nodes corresponding to self-attention heads and directed edge connecting heads of consecutive layers.

We formalize this intuition in the following:

Each term in (1) describes a path of length $L$ across heads of different layers

and there are a total of $H^{L}$ such paths without skip connections.

The path decomposition thus describes the action of a multi-head SAN as the combination of simpler single-head networks. To gain intuition on path interdependence, it helps to split the operations performed into two types: those that act across tokens (multiplication from left) and those that apply independently on each token (multiplication from right). As seen, though paths can interact through token mixing (since $P_{\textit{path}}$ matrices jointly depend on $X$), token-wise operations are independent. We can also notice that biases are not particularly meaningful: their total contribution amounts to the single term $\text{{1}}\bm{b}^{\top}$ independently of the number of layers or heads used.

In the following we show that each path converges rapidly (as a function of length) to a rank-1 matrix with identical rows. Interestingly, this convergence is so dominant that adding more layers to the SAN does not help: though the number of paths is increased exponentially, each path degenerates doubly exponentially, leading also to a rank-1 output.

Before tackling the full SAN, it is instructive to consider the behavior of each path separately. We examine, in particular, how the residual

changes during the forward pass.

As the following result shows, the residual norm converges to zero surprisingly quickly (doubly exponentially with a cubic rate):

For the full theorem, we refer the reader to the Appendix.

Note that the bound in Eq 2 guarantees $\|\text{res}(\text{SAN}(\bm{X}))\|_{1,\infty}$ convergence for all inputs of small residual whenever $4\gamma\beta<\sqrt{d_{\textit{qk}}}$. In practice, our experiments imply that the region for convergence can be much greater.

The identified cubic rate of convergence is significantly faster than what would be expected when analyzing products of stochastic matrices (linear rate). As a rule of thumb, to achieve a decline of three orders of magnitude, say from 1000 to 1, one could expect a linear rate of convergence to require roughly a dozen iterations, whereas a cubic rate can do so in just two or three iterations. The reason why we get a cubic rate is that the rank of attention matrices depends also on the rank of the input. As we show, the self-attention heads mix tokens faster when formed from a low-rank matrix. This phenomenon becomes stronger as we build deeper SANs, leading to a cascading effect.

We provide a proof sketch bellow. Detailed proofs can be found in the Appendix.

We now move on to analyse the convergence of SANs with multiple heads per layer.

Our main result is as follows:

The bound guarantees convergence of $\text{SAN}(\bm{X})$ to rank one when $4\gamma\beta H<\sqrt{d_{\textit{qk}}}$. Our experiments show that this is a rather pessimistic estimate, as, in practice, we observe widespread convergence of output to rank-1.

Remark 1. Implications for Xformers. There has been a surge of architectural variants – that we collectively refer to as Xformers – aimed to improve the vanilla transformer [VSP⁺17] by reducing the quadratic self-attention complexity. The rank collapse result of Theorem 2.3 carries interesting implications for these architectures. One such variant relies on low-rank or kernel-based approximations to the full attention matrix [KVPF20, WLK⁺20, CLD⁺20], in which case the paths likely converge even faster to rank one due to the imposed low-rankedness. Another variant only computes a subset of the attention matrix entries using particular patterns [ZGD⁺20, CGRS19], such as random patterns, in which case one expects the paths to converge more slowly, as randomization tends to increase the rank of the output.

A simple modification to the path decomposition argument for SAN suffices to take into account skip connections. Specifically, we indicate the event that a path has skipped a layer by setting $h=0$ on the corresponding notation:

where we have fixed $\bm{P}_{0}=\bm{I}$ and $\bm{W}_{0}=\bm{I}$.

As observed, skip connections dramatically diversify the path distribution. Denote by $\mathcal{P}_{l}$ the set of paths of length $l$. With skip connections enabled, we have

paths of length $l$ (whereas before we had only length $L$ paths). We hypothesize that it is the presence of short paths that stops SAN from degenerating to rank-1.

While we can derive an upper bound for the residual similar to above (which we do in the Appendix for completeness) such an upper bound is vacuously large. Indeed, it is more informative to have a lower bound on the residual, to align with practice, where SANs with skip connections do not suffer rank collapse. We present the following simple lower bound:

The proof is elementary: by the path decomposition, there is always a path that skips all layers, i.e. the path with length 0, preserving the residual. It then follows that, for any parametrization that renders the contribution of the SAN layers orthogonal to the input, we will have $\|\text{res}(\bm{X}^{L})\|\geq\|\text{res}(\bm{X})\|$. A simple example of such a parametrization can be recovered by setting $\bm{W}_{V}^{l}=0$ for every $l\in[L]$, in which case $\|\text{res}(\bm{X}^{L})\|=\|\text{res}(\bm{X})\|$.

A tight lower bound to the residual in the presence of skip connections is highly nontrivial, and we pose it as an open challenge to the community.

Remark 2. SANs as ensembles of shallow networks. It can be deduced from Theorem 2.3 that the SANs with skip connections enabled heavily rely on short paths (since the residual rapidly declines as the path length becomes larger). In other words, SANs behave like ensembles of shallow single-head self-attention networks. The phenomenon was previously identified for ResNets [VWB16a] (though the latter study didn’t study the rank-collapse phenomenon). Here, the components of this ensemble are inter-dependent, as each attention head participates in many paths of different lengths. Experimental results in §4 support this implication. The supplementary material also provides a study of the paths distribution across several common architectures.

We now study how using an MLP affects the residual. In particular, we focus on SANs with layers written as

Note that, to keep the notation compact, we use $f_{l}$ to denote both the MLP as well as the output bias.

In our subsequent analysis, we use $\lambda_{l,1,\infty}$ to denote the Lipschitz constant of $f_{l}$ with respect to $\ell_{1,\infty}$ norm. Note that, though finding the exact constant can be NP-hard even for shallow MLPs [SV18], since $f_{l}$ comprises of linear transformations with Lipschitz nonlinearities, $f_{l}$ is generally Lipschitz.

As seen, though the effect of MLP is less drastic than that of skip connections, the convergence rate in Cor 3.2 can be controlled by the Lipschitz constants $\lambda_{f,1,\infty}$ of the MLPs: the more powerful the MLPs are the slower the convergence becomes. This reveals a tug-of-war between the self-attention layers and the the MLPs, which due to their nonlinearity can increase the rank. §4 shows that indeed MLPs counteract convergence in experiments.

We should emphasize that using MLPs to counteract the rank-collapse is not without drawbacks: While increasing the Lipschitz constants slows down residual convergence, it also renders the model less robust and more sensitive to input perturbations [CKSN18]. Larger Lipschitz constants may also pose greater challenges to optimization, as they lead to larger gradient variance.

Layer normalization is accomplished by rescaling and shifting the input across the feature dimension:

where $\bm{b}_{\text{LN}}$ is the mean of each column $\text{SA}(\bm{X})$ and $\bm{D}_{LN}$ is a diagonal matrix with entries corresponding to the (possibly scaled or shifted) standard deviation of each column $\text{SA}(\bm{X})$.

By setting $\tilde{\bm{W}}_{h}=\bm{W}_{h}\bm{D}_{LN}^{-1}$ and $\tilde{\bm{b}}_{O}=\bm{b}_{O}-\bm{b}_{LN}$, the above is re-written as

which is identical to the equation before layer normalization was applied, though now $\tilde{\bm{W}}_{h}$ and $\tilde{\bm{b}}_{O}$ are input dependent. Since right multiplication cannot increase the rank of a matrix, we conclude that layer normalization does not mitigate the rank collapse.

To verify our theoretical predictions, we examine the residual of three well-known transformer architectures: BERT [DCLT18], Albert [LCG⁺19], and XLNet [YDY⁺19]. Figure 2 plots the relative residual ${\|\text{res}(\text{SAN}(\bm{X}^{l})\|_{1,\infty}}/{\|\text{SAN}(\bm{X}^{l})\|_{1,\infty}},$ of each layer’s output before and after the networks have been trained. To compute these ratios we ran the network on 32 samples of 128 tokens excerpts of biographies from Wikipedia [LGA16] and display the mean and standard deviation.

The experiments confirm that, as soon as the skip connections are removed, all networks exhibit a rapid rank collapse. Though MLPs do not seem to help in the mitigation of convergence, we caution that the observation is not an accurate portrayal of how trained transformers behave: removing the skip connections introduces a drastic distribution shift in the MLP input. We expect that the convergence will slow down if the network is retrained.

To empirically investigate the inductive bias of the different components of the transformer architecture, we study the behavior of a single-layer transformer when applied recurrently (akin to the universal transformer [DGV⁺19]) to predict a simple 2D circular sequence.

Specifically, we train a single-layer transformer to sequentially predict two circular arcs in $\mathbb{R}^{2}$ of radius $0.3$, starting at $(-0.3,0)$ and $(0.3,0)$, respectively, each directed counter-clockwise and consisting of $1000$ points (illustrated as gray trajectories). An input sample consists of a sequence of two opposing points on the circle, one from the top arc and the other from the bottom arc. We apply teacher-forcing at each step, meaning we give the network the ground truth coordinates of the two current points, and train it to predict the next two points. The model attempts to minimize the MSE loss between the predicted points and the ground truth points on the trajectories. At inference time, we don’t apply teacher-forcing, and simply feed the model output as input for the next step.

Since this recurrent application of a single-layer transformer can be reparametrized to be equivalent to a multi-layer transformer without skip connections, we hypothesize that at inference time the predicted trajectories of the two arcs will converge to the same point (indicating a rank collapse), rather than following the training trajectories. Note that the setting has also been intentionally constructed to enable training even without skip connections (by using teacher forcing) and thus to disentangle the two distinct benefits of skip connections: their ability to improve optimization and their mitigation of rank collapse.

We trained the network until it could perfectly memorize the next step on the circular trajectories with near-zero loss. Figure 3 demonstrates the trajectories predicted at inference time (i.e., without teacher forcing). As seen on the top row, without MLP or skip connections the network exhibits rank collapse. Theorem 2.2 predicts that the convergence is slower when $\beta\geq\|\bm{W}_{QK}^{l}\|_{1}\|\bm{W}_{V}^{l}\|_{1,\infty}$ increases. Indeed, as the hidden dimension increases from 32 to 128 (leading to larger $\beta$ at initialization), the convergence slows down, becoming hardly observable for dimension 128.

We conclude that, in accordance to our analysis, adding MLP or skip connections either stops or drastically slows down rank collapse. As observed, skip connections tend to slow down points from moving. The latter phenomenon is because in this setting skip connections introduce a bias towards remaining in the same position. On the other hand, adding MLPs does not exhibit the same bias.

SANs can be seen as ensembles of paths of different lengths (from 0 to $L$), each involving a different sequence of self-attention heads. Our analysis of SAN with skip connections indicates that path expressivity decreases with path length, even if the number of non-linear operations involved increases. To test this hypothesis, we isolate paths of different lengths and evaluate their predictive power.

Tasks. We considered the following three tasks to test path effectiveness with respect to length:

• •

  Sequence memorization. To solve this task, a model needs to memorize a pre-determined mapping from natural language sentences and random label sequences of the same length. We use random tokens (rather than actual labels) to make this purely a test of expressiveness of a network by way of memorizing training data, rather than confounding effects such as generalizability. The models tested are trained to minimize the cross entropy loss between predicted and the ground truth labels. The training data consist of 500 English sentences from Wikipedia and News sources [DGM06, WSM⁺19], which are tokenized using the SentencePiece tokenizer [KR18] into a vocabulary of size $30522$ with 128 tokens per sequence. Each sequence is mapped to a random binary sequence of the same length.

• •

  Learning to sort. Given an input sequence of letters, this task learns to sort the letters in alphabetical ordering (similar task have been studied before [FOŠ19]). Specifically, the model’s output for each input letter is used to determine the position of that letter in the predicted ordering. Each input sequence, of length $8$, is created by sampling uniformly randomly, with replacement, from an alphabet of size $10$. The training and test sets consist of 1000 and 200 sequences, respectively. To ensure robustness with respect to hyperparameters, we experimented with a variety of settings (adjusting the model depth, number of heads, and the difficulty of the task by changing the alphabet size and sequence length) and observed consistent behavior.

• •

  Convex hull prediction. This task was inspired by the work of [VFJ15]. Given a sequence of $N$ points uniformly distributed in $[0,1]\times[0,1]$ and shifted by a random bivariate standard normal, this task predicts the convex hull of these points. Specifically, for each point in the set, the model predicts whether it’s part of the convex hull. The training set consists of $10,000$ sequences of points in $[0,1]\times[0,1]$, each of length $10$.

In all three tasks, we report the test-set per-token label prediction accuracy as the evaluation metric.