PaperlayerView exact revision on arXiv

Source-linked AI summary

Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse

Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, Aurelien Lucchi

arXiv:2206.03126v1cs.LG

TL;DR

Transformers can suffer rank collapse at initialization, leaving token representations highly aligned and making the training implications insufficiently understood. This paper analyzes the resulting gradients and signal propagation, showing that query/key gradients vanish, residual branches require depth-dependent scaling to preserve rank, and architectural factors create disproportionate gradient norms.

  • Problem

    The paper addresses how rank collapse and high token alignment at initialization affect Transformer training, a question it describes as insufficiently understood.

  • Method

    The paper computes expected attention-layer gradients and analyzes how randomly initialized Transformer blocks propagate sequence norms, inner products, and correlations.

  • Results

    Rank collapse causes vanishing query/key gradients; depth-dependent residual scaling preserves token-angle structure, while architectural factors produce unequal query/key and value gradient magnitudes.

  • Takeaways & Limitations

    Adequate initialization and residual scaling are central to maintaining signal propagation, while disproportionate parameter gradients may help explain adaptive-optimizer use in Transformers.

  • Takeaways & Limitations

    Extending the signal-propagation results from linear activations to ReLU remains technically difficult.

Abstract

from arXiv · show

Transformers have achieved remarkable success in several domains, ranging from natural language processing to computer vision. Nevertheless, it has been recently shown that stacking self-attention layers - the distinctive architectural component of Transformers - can result in rank collapse of the tokens' representations at initialization. The question of if and how rank collapse affects training is still largely unanswered, and its investigation is necessary for a more comprehensive understanding of this architecture. In this work, we shed new light on the causes and the effects of this phenomenon. First, we show that rank collapse of the tokens' representations hinders training by causing the gradients of the queries and keys to vanish at initialization. Furthermore, we provide a thorough description of the origin of rank collapse and discuss how to prevent it via an appropriate depth-dependent scaling of the residual branches. Finally, our analysis unveils that specific architectural hyperparameters affect the gradients of queries and values differently, leading to disproportionate gradient norms. This suggests an explanation for the widespread use of adaptive methods for Transformers' optimization.

1 Introduction

Transformers rely on stacked attention layers and have achieved broad success, but increasing depth can align token representations at initialization and inhibit training. This paper links that rank-collapse phenomenon to vanishing query/key gradients, proposes depth-dependent residual scaling, and identifies unequal gradient dependencies across parameter groups.

  • 1 Introduction: Transformers combine stacked attention, fully connected sub-layers, residual connections, and layer normalization for sequence modeling across language and vision tasks.Attention computes token-wise importance from queries and keys and combines values accordingly.
  • 1 Introduction: Increasing depth raises token alignment at initialization unless adequate residual scaling is used, which can inhibit training in POST-LN Transformers.The observation is illustrated for IWSLT’14 De-En translation with Adam.
  • 1 Introduction: Rank collapse or high token alignment at initialization causes vanishing gradients for query and key parameter matrices, hindering meaningful attention learning.The effect worsens in very deep networks because rank deficiency can affect several layers.
  • 1 Introduction: Depth-dependent residual scaling approximately preserves token-angle cosines and stabilizes sequence rank, including in the infinite-depth limit.This scaling does more than stabilize activation norms at initialization.
  • 1 Introduction: Query/key gradients depend cubically on the propagating sequence norm, whereas value gradients depend linearly, producing disproportionate gradient norms.The analysis also identifies distinct dependencies on embedding dimension and sequence length.

2 Background

This background defines Transformer blocks, their initialization, and the rank-collapse phenomenon in token representations. It motivates the paper’s analysis of how sequence geometry and gradient signals propagate through randomly initialized architectures.

  • 2 Background: A Transformer stacks attention blocks with token-wise layer normalization, yielding POST-LN or PRE-LN variants depending on where normalization is applied.POST-LN applies normalization after residual connections, whereas PRE-LN applies it to sub-layer inputs.
  • 2 Background: Self-attention forms queries, keys, and values from the input sequence, computes row-wise softmax attention, and combines values before residual and feed-forward processing.The block uses learnable projection matrices and residual strengths α1 and α2.
  • 2 Background: Weights are independently initialized with zero mean and Xavier- or He-style variances scaled by the relevant query, key, value, and feed-forward dimensions.The stated variances are σ²_v = 1/d_v for values and feed-forward weights and σ²_k = 1/d_k for queries and keys.
  • 2 Background: Without residual connections, token representations converge toward a rank-1 matrix in which all tokens become identical, producing perfect alignment.This rank collapse can occur doubly exponentially with depth, while layer normalization and fully connected layers only partially alleviate it.
  • 2 Background: Vanishing gradients hinder gradient-based optimization by providing too little backpropagating signal, and their role in Transformers was previously underexplained.The paper frames rank collapse as a setting in which query and key gradients can vanish.
  • 2 Background: The paper studies sequence propagation through expected norms and token inner products or correlations, connecting activation geometry to signal propagation and dynamical isometry.Its focus is relationships among tokens within the same sequence rather than between separate data points.

3 Theoretical Results

The analysis links rank collapse to vanishing query-key gradients and shows that depth-dependent residual scaling preserves token geometry, while architectural factors create unequal gradient magnitudes.

  • Assumptions: The uniform-attention assumption is exact almost surely as d_k →∞ and remains a faithful approximation in simulations with d_k = d_v ≈512.This assumption underlies the gradient analysis and propagation results.
  • 3.1 Vanishing Gradients for Queries and Keys under Rank Collapse: Rank-1 token representations place initialization in a region where query-key gradient norms vanish, making meaningful attention learning difficult, especially in deep networks.The proof uses gradient-norm bounds under uniform attention and rank-1 representations; the resulting quantities vanish almost surely.
  • 3.2 Forward Signal Propagation and the Importance of Scaling the Residual Branches: Without depth-dependent scaling, token correlations approach one and rank collapse emerges at infinite depth; setting residual strengths to O(1/L) approximately preserves norms and inner products.The scaling maintains input-dependent correlations even in the infinite-depth limit.
  • Limitations: Extending the propagation analysis from linear activations to ReLU remains technically difficult because shared position-wise ReLU layers complicate inner-product propagation.Exact formulas are available only for restricted ReLU settings.
  • 3.3 Gradient-Scale Dependencies: Gradient norms depend differently on input variance, token correlation, embedding size, and sequence length for queries/keys versus values, producing disproportionate parameter updates.Query gradients have cubic dependence on input variance, whereas value gradients have linear dependence; their correlation dependence is also opposite.

4 Related Work

The paper builds on work about signal propagation, rank collapse, layer normalization, and residual scaling in Transformers, extending these lines through backward-pass analysis and rank-preserving initialization.

  • The paper extends Transformer rank-collapse research by showing that collapse leads to vanishing query and key gradients.This backward-pass analysis distinguishes the work from prior identification of the rank-collapse phenomenon.
  • Layer normalization is analyzed as a forward-pass variance stabilizer, while the paper argues its inclusion is not strictly necessary for controlling forward signal propagation.
  • The analysis connects differing query and value gradient magnitudes to the empirically observed effectiveness of adaptive optimization methods.The paper presents this as a rigorous explanation of a discrepancy previously hypothesized to matter for Transformer training.
  • The proposed 1/√L residual scaling provides a theoretically grounded way to stabilize rank propagation, beyond prior residual-network stabilization results.The paper characterizes the effect of skip connections on rank propagation rather than only activation stability.
  • Prior studies found properly rescaled residual connections beneficial for Transformer training, but did not characterize their impact on rank propagation.

5 Conclusions and Future Work

The paper concludes that rank collapse at initialization causes vanishing query and key gradients, while appropriately scaled residual connections can stabilize rank propagation and reveal disproportionate parameter-gradient magnitudes.

  • Rank collapse and high token correlation at initialization cause vanishing query and key gradients in Transformer architectures.Residual connections mitigate collapse at finite depth but do not prevent high token alignment without appropriate scaling.
  • Appropriate 1/√L scaling of residual branches is required to prevent high alignment of token representations as depth increases.The conclusion distinguishes finite-depth mitigation from stability in deeper networks.
  • Input variance, embedding size, and sequence length can produce large differences between query/key and value gradient magnitudes.These dependencies provide a theoretical basis for disproportionate gradient scales across parameter subsets.
  • The paper concludes that Transformer effectiveness depends partly on architectural design together with adequate initialization.

A.1.3 Proof of Theorem 3.2

Under uniform attention and rank-1 token representations, the query and key gradients vanish in expectation and therefore almost surely. Gradients for values and inputs remain nonzero, so backward information can still propagate.

  • Theorem 3.2: Rank-1 token representations under uniform attention make the query and key gradient norms vanish almost surely.The proof first establishes zero expected squared gradient norms; non-negativity then implies almost-sure vanishing.
  • Implication: Unlike query and key gradients, gradients with respect to values and inputs do not vanish in the rank-collapse regime.Thus rank collapse does not eliminate all backward information, although the paper states that it still substantially affects training.
  • Proof strategy: The vanishing result follows by conditioning on the layer representation and bounding the remaining network-dependent terms.Softmax derivatives are almost surely bounded, and the analogous linear-network expectation is finite.
  • Rank-collapse structure: The proof represents every token as a scalar multiple of one vector and uses equal coefficients to characterize rank collapse.This reduces the representation matrix to Xℓ = 1_n x^T and enables the gradient calculation.
  • Moment calculation: The calculation uses fourth-order moments of correlated Gaussian variables, expanded through Isserlis’ theorem and separate equal- and unequal-index cases.The proof decomposes the terms according to whether the feature indices coincide and then sums the resulting contributions.

A.3 Forward Pass: Proofs of Lemma 3.2 and 3.3

The forward-pass proofs track norms and correlations through attention, linear, and skip layers. Depth-dependent residual scaling stabilizes activation norms even as depth becomes arbitrarily large.

  • Proof setup: The analysis derives forward moments under uniform attention by combining zero-mean residual terms with expectations from linear and attention layers.Conditional independence makes the cross terms vanish in the skip-connection calculation.
  • Attention contribution: Under uniform attention, the expected squared attention output is proportional to the input correlation and inversely proportional to the value dimension and sequence length squared.This relation is used to propagate the forward statistics through the attention block.
  • Depth limit: The norm recursion is unrolled across layers, yielding a finite-depth expression whose limit is controlled by the residual scaling.The proof applies the product law for limits after separating the two residual-block contributions.
  • Propagation of the norm: Depth-dependent residual scaling stabilizes the norm of propagated representations as network depth grows.The result follows by recursively combining the expected contributions of skip, linear, and attention layers.

A.4 Proof of Theorem 3.3: Correlations are Preserved under Residual Scaling

With equal-norm input tokens, the paper shows that depth-dependent residual scaling preserves token correlations. Constant nonzero residual coefficients instead produce a different depth-dependent correlation behavior.

  • Theorem 3.3: Depth-dependent residual scaling preserves the correlations between equal-norm token representations as depth increases.The theorem analyzes the expected norm and pairwise correlation under the stated scaling regime.
  • Comparison: When residual coefficients are fixed nonzero constants independent of depth, the resulting correlation behavior differs from the depth-scaled case.The theorem treats this constant-coefficient regime separately from the stabilizing scaling choice.
  • Proof setup: Rotational symmetry of Gaussian weight matrices ensures that equal-norm input tokens have equal expected norms at every layer.This symmetry reduces the correlation analysis to a shared per-token norm and pairwise correlation coefficient.
  • Correlation recursion: The correlation recursion is obtained by expanding the relation between pairwise inner products, token norms, and the correlation coefficient.Summing over token indices produces the expression used to take the depth limit.

A.5 Motivation for Assumption 3.1

The paper motivates uniform attention by showing that increasing the key dimension makes the attention matrix approach the uniform matrix. Empirically, the deviation decreases with a trend proportional to d_k^-2.

  • Assumption: The uniform-attention assumption treats attention as a matrix with every entry equal to 1/n.This assumption enables the subsequent expectation calculations involving attention and its derivatives.
  • Theoretical motivation: The theoretical argument establishes almost-sure convergence of the softmax attention matrix to the uniform matrix as d_k tends to infinity with d_v fixed.It uses Glorot initialization, convergence of the query-key product, and continuity of softmax.
  • Figure 7: Figure 7 measures the average Frobenius distance between the attention matrix and the uniform matrix as a function of d_k.The caption fixes d_v at 100.

B.2 Further Empirical Assessment of Assumption 3.1

The experiments test the uniform-attention assumption in forward and backward analyses, finding strong agreement under Xavier-like initialization but deviations at higher input variance and depth.

  • At input standard deviation 1/d, the uniform-attention assumption closely approximates common Xavier initialization.
  • Increasing input standard deviation causes deviations from the uniform-attention approximation.
  • As depth increases, the representation-norm distribution becomes more heavy-tailed.
  • Theoretical predictions for query and key gradient norms show very good agreement with automatically differentiated gradients across hidden dimension, sequence length, input correlation, and input variance.

B.3 Empirical Verification of the Gradient Analysis of Section 3.3

The empirical gradient study verifies the theory across architectural and task-specific parameters, showing that token correlation and input variance create distinct scaling for query and value gradients.

  • Increasing token correlation quickly diminishes query-gradient norms relative to value-gradient norms.
  • Input variance affects value gradients linearly but query gradients cubically, highlighting the importance of stabilizing the forward pass.
  • The dependence on sequence length and hidden dimension is more complicated and also varies with token correlation.
  • The experiments vary hidden dimension, sequence length, input correlation, and input variance while averaging results over 20 runs.

C.2 Translation Task

The translation experiments examine temperature scaling and trainable residual branches in a POST-LN Transformer, relating gradient balance to token correlation and activation propagation.

  • Temperature scaling is chosen with τfinal = 8.5 to match value and query gradient norms under the stated correlation and sequence-length assumptions.
  • The translation evaluation uses the best-performing model selected on a held-out validation set.
  • In the best trained model, residual scaling parameters increase while token correlation does not significantly increase, allowing efficient gradient propagation.
Loading 2206.03126v1…