Gradient Flow Polarizes Softmax Outputs towards Low-Entropy Solutions

Aditya Varre; Mark Rofin; Nicolas Flammarion

cs.LGMar 6, 2026

Gradient Flow Polarizes Softmax Outputs towards Low-Entropy Solutions

Authors: Aditya Varre, Mark Rofin, Nicolas Flammarion

Organizations: Theory of Machine Learning Lab, EPFL.

Abstract

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}(\mathbf{V} σ(\mathbf{a}))$ , where $\mathbf{V}$ and $\mathbf{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

Panel of Figure 1 — 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

\sigma(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

\kappa=1

and 5.

Name	$\mathrm{sim}(\phi(\mathbf{q}_{i}),\phi(\mathbf{k}_{j}))$	Normalization
softmax	$\displaystyle\exp\!\left(\frac{\mathbf{q}_{i}\mathbf{k}_{j}^{\top}}{\sqrt{d_{h}}}\right)$	$\displaystyle\sum_{j^{\prime}=1}^{m}\exp\!\left(\frac{\mathbf{q}_{i}\mathbf{k}_{j^{\prime}}^{\top}}{\sqrt{d_{h}}}\right)$
sigmoid	$\displaystyle\sigma\!\left(\frac{\mathbf{q}_{i}\mathbf{k}_{j}^{\top}}{\sqrt{d_{h}}}\right)$	$\displaystyle\sum_{j^{\prime}=1}^{m}\sigma\!\left(\frac{\mathbf{q}_{i}\mathbf{k}_{j^{\prime}}^{\top}}{\sqrt{d_{h}}}\right)$
elu	$\displaystyle\mathrm{elu}\!\left(\frac{\mathbf{q}_{i}\mathbf{k}_{j}^{\top}}{\sqrt{d_{h}}}\right)+1$	$\displaystyle\sum_{j^{\prime}=1}^{m}\Big(\mathrm{elu}\!\left(\frac{\mathbf{q}_{i}\mathbf{k}_{j^{\prime}}^{\top}}{\sqrt{d_{h}}}\right)+1\Big)$
elu $\cdot$ elu	$\displaystyle\frac{\big(\mathrm{elu}(\mathbf{q}_{i})+1\big)\big(\mathrm{elu}(\mathbf{k}_{j})+1\big)^{\top}}{\sqrt{d_{h}}}$	$\displaystyle\sum_{j^{\prime}=1}^{m}\frac{\big(\mathrm{elu}(\mathbf{q}_{i})+1\big)\big(\mathrm{elu}(\mathbf{k}_{j^{\prime}})+1\big)^{\top}}{\sqrt{d_{h}}}$
mlp	$\displaystyle\frac{\mathrm{mlp}(\mathbf{q}_{i})\,\mathrm{mlp}(\mathbf{k}_{j})^{\top}}{\sqrt{d_{h}}}$	$\displaystyle\max\!\left(\left\|\sum_{j^{\prime}=1}^{m}\frac{\mathrm{mlp}(\mathbf{q}_{i})\,\mathrm{mlp}(\mathbf{k}_{j^{\prime}})^{\top}}{\sqrt{d_{h}}}\right\|,\,1\right)$
linear	$\displaystyle\frac{\mathbf{q}_{i}\mathbf{k}_{j}^{\top}}{\sqrt{d_{h}}}$	$\displaystyle\max\!\left(\left\|\sum_{j^{\prime}=1}^{m}\frac{\mathbf{q}_{i}\mathbf{k}_{j^{\prime}}^{\top}}{\sqrt{d_{h}}}\right\|,\,1\right)$

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

t_{i}^{\prime}\neq t_{i}

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	$\{{\bf 1},2\}$	Epochs	10
Hidden dim	128	Train size	10,000
Feedforward dim	512	Eval size	1,000
Vocab size	$\{\bf{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\sigma(Ua)

) with logistic loss for different dimensionality

D

. We denote the column of

U

corresponding to the largest element of

a

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

\kappa

from 1 to 5. Higher condition number leads to slower optimization and stronger sparsity.

Panel of Figure 13 — 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.

Panel of Figure 14 — Figure 14: Sink proportions in induction-trained Transformer with $\lambda=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.

Panel of Figure 15 — Figure 15: Sink proportions in induction-trained Transformer with $\lambda=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.

Panel of Figure 16 — Figure 16: Average induction accuracy for models trained with $\lambda=0$ . Each row represents a type of positional encodings.

Panel of Figure 17 — Figure 17: Average induction accuracy for models trained with $\lambda=0.2$ . Each row represents a type of positional encodings.

Panel of Figure 18 — Figure 18: Flip-rate plots for vocabulary size $|\mathcal{V}|=41$ . Rows: MLP vs. no-MLP. Columns: LayerNorm vs. no-LayerNorm.

Panel of Figure 19 — Figure 19: Flip-rate plots for vocabulary size $|\mathcal{V}|=81$ . Rows: MLP vs. no-MLP. Columns: LayerNorm vs. no-LayerNorm.

Panel of Figure 20 — Figure 20: Flip-rate plots for vocabulary size $|\mathcal{V}|=161$ . Rows: MLP vs. no-MLP. Columns: LayerNorm vs. no-LayerNorm.

Panel of Figure 21 — Figure 21: Flip-rate plots for $|\mathcal{V}|=41$ with no MLP and no LayerNorm. Rows: number of heads $H\in\{1,2,4\}$ . Columns: number of layers $L\in\{1,2\}$ . Increasing the number of heads or layers leads to the diminishing flip rate.

Explore similar work

Attention Sinks Are Provably Necessary in Softmax Transformers: Evidence from Trigger-Conditional Tasks
Mar 12, 2026Yuval Ran-MiloAttention Sinks Attention Mechanisms
Tel Aviv University
Gradient Flow Structure and Quantitative Dynamics of Multi-Head Self-Attention
May 5, 2026Ayan PendharkarSoftmax Attention Attention Head Analysis
Attention Sinks Induce Gradient Sinks: Massive Activations as Gradient Regulators in Transformers
Mar 18, 2026Yihong Chen, Zhouchen Lin, Quanming YaoAttention Sinks
Tsinghua University · Peking University

Mar 12, 2026J/K move · Enter open · S save

Attention Sinks Are Provably Necessary in Softmax Transformers: Evidence from Trigger-Conditional Tasks

Yuval Ran-Milo

Tel Aviv University

Attention Sinks Attention Mechanisms Softmax Attention

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Attention Sinks Are Provably Necessary in Softmax Transformers: Evidence from Trigger-Conditional Tasks

Gradient Flow Structure and Quantitative Dynamics of Multi-Head Self-Attention

Attention Sinks Induce Gradient Sinks: Massive Activations as Gradient Regulators in Transformers