Understanding the intricate non-convex training dynamics of softmax-based models is crucial for explaining the empirical success of transformers. In this article, we analyze the gradient flow dynamics of the value-softmax model, defined as L(Vσ(a)), where V and a are a learnable value matrix and attention vector, respectively. As the matrix times softmax vector parameterization constitutes the core building block of self-attention, our analysis provides direct insight into transformer's training dynamics. We reveal that gradient flow on this structure inherently drives the optimization toward solutions characterized by low-entropy outputs. We demonstrate the universality of this polarizing effect across various objectives, including logistic and square loss. Furthermore, we discuss the practical implications of these theoretical results, offering a formal mechanism for empirical phenomena such as attention sinks and massive activations.
Figures & tables
Figure 1: Representative attention patterns of the 2nd Transformer layer solving an induction task. Sigmoid and linear attentions are trained without additional normalization (see Table 1 ). An attention sink clearly emerges for softmax but not other attentions.
Figure 2: Experimental verification of Theorem 3.3 using logistic loss. The plot shows the evolution of attention scores σ(a) and value projections u over time. As predicted, the attention scores converge to a one-hot vector (blue line goes to 1, others to 0) and the value projections u diverge, preserving their order.
Figure 3: Evolution of attention logits in a regression problem with condition number κ=1 and 5.
Name
sim(ϕ(qi),ϕ(kj))
Normalization
softmax
exp(dhqikj⊤)
j′=1∑mexp(dhqikj′⊤)
sigmoid
σ(dhqikj⊤)
j′=1∑mσ(dhqikj′⊤)
elu
elu(dhqikj⊤)+1
j′=1∑m(elu(dhqikj′⊤)+1)
elu ⋅ elu
dh(elu(qi)+1)(elu(kj)+1)⊤
j′=1∑mdh(elu(qi)+1)(elu(kj′)+1)⊤
mlp
dhmlp(qi)mlp(kj)⊤
maxj′=1∑mdhmlp(qi)mlp(kj′)⊤,1
linear
dhqikj⊤
maxj′=1∑mdhqikj′⊤,1
Table 1: Attention functions we test in Section 4 . When computing the normalization, m= sequence length for non-causal models and m=i for causal models. For all functions but softmax, we test their variants both with and without normalization.
Figure 4: Proportion of sink heads at the 2nd layer of Transformer trained for the induction task.
Figure 5: Evolution of the attention scores for two differently labeled samples in the classification experiment.
Figure 6: Success rates of adversarial token 1-token flips in 1-layer 1-head Transformer trained for classification.
Figure 7: Distribution of attention head sparsity scores for pretrained softmax/sigmoid LLMs.
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Task format in the induction heads experiment. The arrows illustrate the induction head mechanism.
Hyperparameter
Value
Hyperparameter
Value
Layers
2
Epochs
10
Hidden dim
128
Train size
100,000
Feedforward dim
512
Eval size
1,000
Vocab size
51∗
Sequence length
13
Causal
true
Batch size
256
Use LayerNorm
true
Use MLP
true
Appendix
Table 2: Configuration and hyperparameters for the induction experiment. The bold values indicate the setup shown in Figures 1 and 4 . Results for other setups are additionally presented in Appendix F.2 Vocabulary size is increased by 50 if ti′=ti .
Figure 9: Task format in the classification experiment and the token swap procedure. Cell brightness illustrates the attention weight assigned to each token.
Hyperparameter
Value
Hyperparameter
Value
Layers
{1,2}
Epochs
10
Hidden dim
128
Train size
10,000
Feedforward dim
512
Eval size
1,000
Vocab size
{41,81,161}
Sequence length
11
Labels
4
Causal
false
Batch size
256
Positional type
no
Appendix
Table 3: Configuration and hyperparameters for the classification experiment. The bold values indicate parameters used for the setup in Figures 6 and 5 . Results for other setups are additionally provided in Appendix F.3
Figure 10: Training dynamics of Value-Softmax model with KL-divergence as loss function.
Figure 11: Training dynamics of tied Value-Softmax model ( Uσ(Ua) ) with logistic loss for different dimensionality D . We denote the column of U corresponding to the largest element of a as U(m) . In all cases, a sparse solution is found, and the norm of U(m) diverges.
Figure 12: Training dynamics in regression problems with condition number κ from 1 to 5. Higher condition number leads to slower optimization and stronger sparsity.
Figure 13: Training dynamics of value-softmax model for square and logistic losses and various activation functions. Note that square loss does not give rise to sparse solutions. Additionally, sum and sigmoid do not converge to sparse solutions for logistic loss as well.
Figure 14: Sink proportions in induction-trained Transformer with λ=0 . Each row represents a type of positional encodings. Setups that result in training failures are marked by red crosses. We observe that under this setup, most heads duplicate each other and do not form sinks.
Figure 15: Sink proportions in induction-trained Transformer with λ=0.2 . Each row represents a type of positional encodings. Setups that result in training failures are marked by red crosses. Interestingly, we observe that setups with no explicit positional encodings lead to an increased formation of attention sinks. Note that in this case the model has access to positional information through causal masking, so the task is still solvable.
Figure 16: Average induction accuracy for models trained with λ=0 . Each row represents a type of positional encodings.
Figure 17: Average induction accuracy for models trained with λ=0.2 . Each row represents a type of positional encodings.
Figure 18: Flip-rate plots for vocabulary size ∣V∣=41 . Rows: MLP vs. no-MLP. Columns: LayerNorm vs. no-LayerNorm.
Figure 19: Flip-rate plots for vocabulary size ∣V∣=81 . Rows: MLP vs. no-MLP. Columns: LayerNorm vs. no-LayerNorm.
Figure 20: Flip-rate plots for vocabulary size ∣V∣=161 . Rows: MLP vs. no-MLP. Columns: LayerNorm vs. no-LayerNorm.
Figure 21: Flip-rate plots for ∣V∣=41 with no MLP and no LayerNorm. Rows: number of heads H∈{1,2,4} . Columns: number of layers L∈{1,2} . Increasing the number of heads or layers leads to the diminishing flip rate.
Transformers often display an attention sink: probability mass concentrates on a fixed, content-agnostic position. Are sinks a byproduct of the optimization/training regime? Or are they sometimes functionally necessary in softmax Transformers? We prove that, in some settings, it is the latter: computing a simple trigger-conditional behavior necessarily induces a sink in softmax self-attention models. Our results formalize a familiar intuition: normalization over a probability simplex must force attention to collapse onto a stable anchor to realize a default state (e.g., when the model needs to ignore the input). We instantiate this with a concrete task: when a designated trigger token appears, the model must return the average of all preceding token representations, and otherwise output zero, a task which mirrors the functionality of attention heads in the wild (Barbero et al., 2025; Guo et al., 2024). We also prove that non-normalized ReLU attention can solve the same task without any sink, confirming that the normalization constraint is the fundamental driver of sink behavior. Experiments validate our predictions and demonstrate they extend beyond the theoretically analyzed setting: softmax models develop strong sinks while ReLU attention eliminates them in both single-head and multi-head variants.
Transformer self-attention can be interpreted as a gradient flow on the unit sphere, in which tokens evolve under softmax interaction potentials and tend to form clusters. While prior work has established clustering behavior for single-head attention, the multi-head setting remains less understood due to geometric interference between heads, which invalidates standard monotonicity arguments. In this work, we develop a theoretical framework for multi-head self-attention dynamics and resolve several open questions. We show that, under suitable conditions on the score matrices, a natural multi-head energy functional is non-decreasing along both flat and spherical dynamics. We identify the key obstruction to per-head monotonicity as radial shadow terms, which are projections of each head's output onto token directions, persisting even under orthogonality assumptions. We introduce a sufficient condition ensuring monotonicity and establish robustness to approximate orthogonality. In a simplified scalar-head regime with equiangular token configurations, we derive a closed-form expression for the critical inverse temperature governing clustering behavior, and show that heterogeneous heads exhibit super-additive clustering rates. In this regime, we also prove a separation in clustering time between ReLU and softmax attention in the linearized dynamics. Finally, we establish an entropy production identity and show that attention entropy increases monotonically toward equilibrium as clustering progresses. Our results provide a unified perspective on the dynamics of multi-head attention and clarify the mechanisms underlying clustering and stability in transformer models.
Attention sinks and massive activations are recurring and closely related phenomena in Transformer models. Existing explanations have largely focused on the forward pass, yet in pre-norm Transformers, large residual-stream norms play only an indirect forward role because sublayers operate on normalized inputs. We study this relationship from the perspective of backpropagation. Empirically and theoretically, we show that under causal masking, attention sinks can induce pronounced gradient concentration, which we term gradient sinks. Since the RMSNorm Jacobian attenuates gradients roughly in inverse proportion to input norm, massive activations can be understood as adaptive regulators of this localized gradient pressure during training. This interpretation predicts that attenuating sink-induced gradients should weaken massive activations. We test this prediction with V-scale, a modification that adjusts backpropagated gradients on the value path. In V-scale models, attention sinks are preserved, whereas massive activations are suppressed. These results identify gradient sinks as a backward-pass counterpart of attention sinks, and massive activations as an adaptive RMSNorm-mediated response that attenuates the resulting localized training pressure. Our code is available at https://anonymous.4open.science/r/GradientSinkCode-B309.