PaperlayerView exact revision on arXiv

Source-linked AI summary

Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth

Yihe Dong, Jean-Baptiste Cordonnier, Andreas Loukas

arXiv:2103.03404v2cs.LG

TL;DR

The paper addresses limited understanding of why attention-based architectures work by decomposing self-attention networks into paths of attention heads across layers. It proves that pure self-attention converges doubly exponentially to token-uniform rank-1 outputs, while skip connections and MLPs counteract this degeneration. Experiments verify these convergence phenomena across transformer variants, although larger MLP Lipschitz constants can reduce robustness and complicate optimization.

  • Problem

    Understanding the operation and inductive bias of attention-based architectures remains limited despite their widespread use.

  • Method

    The paper decomposes SAN outputs into weakly interdependent paths, each corresponding to a sequence of attention heads across layers, and analyzes rank convergence.

  • Results

    Pure SANs converge doubly exponentially to rank-1, token-uniform outputs, while skip connections mitigate rank collapse and MLPs slow convergence.

  • Takeaways & Limitations

    Self-attention creates a token-uniformity inductive bias, and skip connections have a vital role beyond optimization and gradient flow in counteracting rank collapse.

  • Takeaways & Limitations

    Increasing MLP Lipschitz constants slows residual convergence but makes models less robust to input perturbations and can increase optimization difficulty through larger gradient variance.

Abstract

from arXiv · show

Attention-based architectures have become ubiquitous in machine learning, yet our understanding of the reasons for their effectiveness remains limited. This work proposes a new way to understand self-attention networks: we show that their output can be decomposed into a sum of smaller terms, each involving the operation of a sequence of attention heads across layers. Using this decomposition, we prove that self-attention possesses a strong inductive bias towards "token uniformity". Specifically, without skip connections or multi-layer perceptrons (MLPs), the output converges doubly exponentially to a rank-1 matrix. On the other hand, skip connections and MLPs stop the output from degeneration. Our experiments verify the identified convergence phenomena on different variants of standard transformer architectures.

1 Introduction

The paper develops tools to understand self-attention and finds that pure self-attention loses expressive power doubly exponentially with depth, while skip connections and MLPs counteract rank collapse.

  • Self-attention architectures are widespread, motivating tools to understand their operation and inductive biases.
  • Pure SANs lose expressive power doubly exponentially with depth, with outputs converging at a cubic rate to rank-1 matrices with identical rows.
  • Skip connections mitigate rank collapse, while MLPs slow convergence by increasing their Lipschitz constant.
  • The paper decomposes SANs into weakly interdependent paths, each corresponding to a sequence of attention-head choices across layers.
  • Its contributions include analyzing opposing architectural effects, proposing path decomposition, and experimentally verifying the theory on common transformer architectures.

2 Attention doubly exponentially loses rank

The analysis decomposes multi-head SANs into paths and shows that token mixing drives rapid convergence toward rank-1, token-uniform outputs as depth increases.

  • SANs converge exponentially with depth to rank-1 matrices that make all tokens identical.
  • Path decomposition: A multi-head SAN can be expressed as a sum of single-head networks called paths, each represented by a sequence of attention heads.
  • Path decomposition: Each attention head uses an input-dependent row-stochastic matrix, while value transformations and biases provide token-wise components.
  • Path decomposition: Without skip connections, a depth-L, width-H SAN contains H^L paths, whose combined degeneration yields rank-1 output.
  • Convergence: The convergence guarantees require bounded weight conditions and sufficiently small residuals, while experiments observe rank-1 convergence more broadly than the bound predicts.
  • Convergence: Low-rank inputs make attention heads mix tokens faster, creating a cascading effect that produces cubic-rate and doubly exponential convergence.
  • Architectural implications: For Xformers, low-rank attention approximations may accelerate path convergence, whereas sparse attention patterns may slow it by increasing output rank.

3 Mechanisms that counteract rank collapse

The analysis identifies skip connections and MLPs as counteracting rank collapse, while layer normalization does not mitigate it. Skip connections create short paths that preserve residual information, whereas MLP effectiveness depends on their Lipschitz constants and involves robustness trade-offs.

  • 3.1 Skip connections are crucial: Skip connections diversify paths by allowing layers to be skipped, creating paths of lengths 0 through L rather than only length L.The path decomposition marks a skipped layer with h = 0; the resulting distribution includes short paths.
  • 3.1 Skip connections are crucial: Short paths are hypothesized to prevent rank-1 degeneration, and skip connections can preserve at least the input residual even as depth tends to infinity.There are infinitely many parameterizations with residual norm at least the input residual, even for L →∞ and arbitrarily small β.
  • 3.1 Skip connections are crucial: A tight lower bound for the residual with skip connections remains an open challenge, while available upper bounds are described as vacuously large.The paper presents a simple lower bound aligned with practice but leaves a tight lower bound unresolved.
  • 3.2 Multi-layer perceptrons (MLP) help: MLPs slow convergence through their Lipschitz constants, with more powerful nonlinearities able to increase rank and counteract self-attention's collapse.The paper characterizes this interaction as a tug-of-war between self-attention and MLPs.
  • 3.2 Multi-layer perceptrons (MLP) help: Increasing MLP Lipschitz constants also reduces robustness to input perturbations and may complicate optimization through larger gradient variance.Thus, using MLPs to counteract rank collapse introduces sensitivity and optimization costs.
  • 3.3 Layer normalization plays no role: Layer normalization does not mitigate rank collapse because its right multiplication cannot increase matrix rank.Layer normalization rescales and shifts inputs across the feature dimension, producing input-dependent parameters without changing this rank argument.

4 Experiments

Experiments verify rank collapse in self-attention variants and show that skip connections, MLPs, width, and path length affect convergence and expressive power.

  • Rank collapse in real architectures: Removing skip connections causes rapid rank collapse across BERT, Albert, and XLNet outputs.The experiments use relative residuals measured on 32 samples of 128-token Wikipedia biography excerpts.
  • Visualizing the bias of different architectures: 32-to-128 hidden-dimension increases slow convergence, making rank collapse hardly observable at dimension 128.The reported slowdown is associated with larger β at initialization.
  • Visualizing the bias of different architectures: Skip connections slow or stop convergence, whereas MLPs do not show the same positional bias and may not mitigate convergence in the tested setting.The authors caution that removing skips shifts MLP inputs and expect retraining to slow convergence.
  • Path effectiveness: Length-1 paths exceed 0.8, 0.6, and 0.65 accuracy on memorization, sorting, and convex-hull tasks, respectively.Longer paths perform near random guessing, while length-zero paths contain little useful information.
  • Path effectiveness: Performance consistently deteriorates with path length, indicating that short paths carry most predictive power across the three tasks.The evaluation metric is test-set per-token label prediction accuracy.
  • Path effectiveness: The transformer behaves like an ensemble of shallow networks, while long paths retain underutilized capacity associated with rank loss.The authors suggest preventing long paths from losing rank as one direction for improving their effectiveness.

5 Related works

The paper relates skip connections to optimization and residual-network path interpretations, distinguishes its output-rank focus from efficient-attention work, and connects rank collapse to stochastic-matrix convergence.

  • Skip connections: Skip connections are known to facilitate optimization by allowing gradients to bypass skipped layers during backpropagation.This paper identifies an additional role: preventing rapid rank-one degeneration of transformer outputs.
  • Path interpretations: Residual networks have been interpreted as collections of varying-length paths, with effective paths shorter than total depth; this work connects short paths to avoiding SAN rank collapse.The comparison links path length with both optimization behavior and expressive rank preservation.
  • Relation to efficient attention: The paper studies the rank of the network output, unlike efficient-attention methods that approximate the attention matrix with low-rank or kernel techniques.The two research directions are described as orthogonal.
  • Stochastic-matrix perspective: Products of stochastic matrices provide intuition for rank-one convergence, but deep attention networks involve more complex matrix interactions.Softmax attention makes each path component contain row-stochastic matrices.

6 Conclusion

The conclusion frames rank collapse as the result of competing forces in self-attention networks and identifies path decomposition as a modular analytical tool.

  • Conclusion: Self-attention promotes rank collapse, while skip connections and MLPs counteract it.The conclusion presents these as competing forces over rank collapse in SANs.
  • Conclusion: The path decomposition modularizes SAN analysis and is presented as independently useful for additional applications.The conclusion also points to future questions about token uniformity, long paths, width-depth trade-offs, and convergence bounds.

A Deferred Proofs

The deferred proofs build the analysis from single-head self-attention to deeper multi-head networks and then extend it to skip connections and MLPs.

  • Proof strategy: The proofs begin with a single-head self-attention layer before treating deeper single- and multi-head networks.The main rank-collapse result is developed after a proof sketch and the path-based formulation.
  • Architectural extensions: The analysis is extended to architectural variants incorporating skip connections and MLPs.These extensions characterize components that counteract self-attention’s rank-collapse behavior.

A.1 Single-layer and single-head

The single-head, single-layer analysis bounds how far self-attention outputs are from being constant across tokens. These bounds support the residual recursion used to establish rank collapse with depth.

  • Single-layer and single-head: Self-attention output SA(X) = P_XWV is bounded in its entry-wise distance from being constant across tokens.The attention matrix P_X is generated by softmax attention probabilities.
  • Single-layer and single-head: The ℓ1 residual satisfies ∥SA(X) − 1(r′)⊤∥1 ≤ 2∥D1∥1∥R∥1∥WV∥1.The bound follows from ∥softmax(r)∥1 = 1 and submultiplicativity of the matrix norm.
  • Single-layer and single-head: The analysis also derives an analogous bound for the ℓ∞ residual and combines the two norm bounds into the main claim.The residual measures deviation from a matrix with identical rows.
  • Single-layer and multiple-layers: For multiple layers, recursively applying the single-layer bound yields a residual bound that amounts to doubly exponential convergence to a rank-1 matrix.The recursion is unfolded backward from the last layer to the first.
  • Single-layer and multiple-layers: The single-head lemma assumes bounded query-key weight norms, with ∥W^l_QK,1∥1 ≤ β for every layer l ∈ [L].This norm condition is part of the setting for the multilayer corollary.

A.4 Multiple-head and multiple-layers

The multi-head, multi-layer analysis extends the single-head residual recursion across heads and layers. Under bounded head weights, the output still converges doubly exponentially to a rank-1 matrix.

  • Multiple-head and multiple-layers: A depth-L SAN with H heads per layer has an output residual bounded under uniform bounds on every head’s weight norm.The condition is ∥W^l_h∥1,∞ ≤ β for all heads h and layers l.
  • Multiple-head and multiple-layers: The resulting bound indicates that the output converges doubly exponentially to a rank-1 matrix.The proof recursively applies the single-layer multi-head residual bound.

A.5 SAN with skip connections

Adding skip connections changes the residual recursion for SANs. The resulting upper bound does not indicate convergence, unlike the bound for SANs without skip connections.

  • SAN with skip connections: For SANs with skip connections, the residual analysis provides an upper bound intended to illustrate residual growth.The bound is described as vacuously large and is motivated by practice, where skip-connected SANs do not suffer rank collapse.
  • SAN with skip connections: The depth-L, width-H skip-connected SAN bound does not indicate convergence.The derivation unfolds the residual recursion backward across layers and accounts for the factor of H.
  • SAN with skip connections: The residual bound is obtained by selecting terms along paths in a complete binary tree and taking the maximum leaf term.The ordering of the choices does not matter; only how often each term is selected matters.

A.6 SAN with MLP

The MLP analysis tracks how an MLP changes the residual while preserving row constancy on already uniform inputs. With bounded Lipschitz constants, the resulting depth-dependent bound still has a doubly exponential convergence rate.

  • SAN with MLP: The SAN-with-MLP theorem assumes ∥W^l_h∥1,∞ ≤ β and λl,1,∞ ≤ λ across heads and layers.Here λl,1,∞ denotes the Lipschitz constant of the layer’s function f with respect to the ℓ1,∞ norm.
  • SAN with MLP: The resulting residual bound amounts to a doubly exponential rate of convergence with respect to the ℓ1,∞ norm.The proof recursively unrolls the bound across layers.
  • SAN with MLP: In the MLP setting, the combined value-output transformation is represented by Wh := WVWO rather than the value weight WV alone.This substitution is used in the residual analysis of the multi-head layer.
  • SAN with MLP: The function f representing the MLP and output bias preserves constancy across rows when applied to a matrix of identical rows.Weight multiplications, bias addition, and entrywise nonlinearities all preserve this property.

B.1 The path length distribution of transformers

Figure 5 visualizes path-length distributions across common encoder-only transformer architectures with varying depths and widths. The authors hypothesize that diverse path distributions may be more efficient than architectures dominated by long paths.

  • Architecture comparison: Figure 5 compares path-length distributions across GPT-3, T5, Bert, ViT, DistilBert, and MobileBert encoder architectures.The architectures vary in depth and width, and the legends are sorted by total heads, L×H.
  • Ensemble structure: The figure examines transformers as ensembles of paths with different lengths.This visualization is intended to provide practical insight into the ensemble structure of transformer architectures.
  • Interpretation: Models overwhelmingly concentrated on long paths are hypothesized to be less efficient than models with more diverse path distributions.The authors also hypothesize that long-path-focused models may be less robust because they require larger MLP Lipschitz constants to counteract self-attention’s token-uniformity bias.
Loading 2103.03404v2…