Training Dynamics of Transformers on Unbounded Inputs: Attention Learning Can Lag Output Layer Learning

Yaomengxi Han; Debarghya Ghoshdastidar

cs.LGOct 14, 2026

Training Dynamics of Transformers on Unbounded Inputs: Attention Learning Can Lag Output Layer Learning

Authors: Yaomengxi Han, Debarghya Ghoshdastidar

Organizations: Technical University of Munich

Abstract

The success of transformers is attributed to several factors, including their capacity to express formal languages over variable-length inputs, and ability of the attention mechanism to focus on relevant input features. However, there is limited understanding of the training dynamics of shallow transformers, specifically how to decouple the learning of the attention and the output (feedforward) layer. A natural extension of the theory of one-step gradient feature learning in feedforward networks suggests that one could first learn the attention weights with a single gradient step, and then train the output layer. We show that this fails for a shallow transformer trained to recognize a star-free language from unbounded-length input. We prove that the transformer can be trained to recognize the language with a three-phase gradient descent: the output layer is updated first, followed by an attention update, and the output layer is fine-tuned in the third phase. In contrast, the language cannot be learned by training only the attention or output layer, or by a single attention-layer update followed by finetuning the output layer. We empirically show that ``attention learning lagging output layer learning" is beneficial in other language recognition, regression, and in-context learning tasks as well.

Figures & tables

Figure 1: (a) Automaton of the occurrence language

A_{s}=\Sigma^{*}s\Sigma^{*}

with alphabet

\Sigma=\{1,\ldots,d\}

and unknown

s\in\Sigma

. The language can be expressed by a transformer with one attention head and linear output layer ( 1 ), which encodes word

w=(w_{1},\ldots,w_{m})

(\mathbf{e}_{w_{1}},\ldots,\mathbf{e}_{w_{m}})

, where

\mathbf{e}_{i}\in\mathbb{R}^{d}

i

-th standard basis vector. The optimal shallow transformer has weights of the output layer

\boldsymbol{\beta}=2\mathbf{e}_{s}\in\mathbb{R}^{d}

, and attention score for symbol

i\in\Sigma

\gamma_{i}=1

i=s

, and 0 otherwise. (b) A shallow transformer is trained to recognize

A_{s}

using gradient descent jointly on both layers. The output gap

(\beta_{s}-\max\limits_{i\neq s}\beta_{i})/2

and attention gap

(\gamma_{s}-\max\limits_{i\neq s}\gamma_{i})

, plotted across training epochs, show that output gap increases before attention gap appears, indicating that attention learning lags output layer learning. (c) Several training plans are considered, which decouple gradient descent on attention layer (A) and output layer (O), sometimes using a single gradient step (1x). “One-step attention learning”, where one attention update is followed by training output layer till loss converges, A(1x)

\to

O, does not correctly learn the language. We later prove that a 3-phase training, O(1x)

\to

A(1x)

\to

O, correctly learns to recognize

A_{s}

. (d)–(e) show how our 3-phase plan reaches the ideal gaps in both layers, but training attention/output only, or attention before output do not.

Figure 2: Joint GD (left), O-A-O training (middle), and A-O training (right). Attention gap, output gap, and test accuracy for

d=50

are shown with

5d^{2}

training and

d^{2}

test samples. Joint GD has no phases; for comparison, we start its Phase 2 when the attention starts to move (gap =

0.1

, at step

31

) and end its Phase 2 when the attention almost saturates (gap =

0.9

, at step

127

). In the 3-phase plan, the test accuracy reaches

1

by the end of Phase 2. Phase 3 only reduces training loss by fine-tuning output layer.

Training plan		Occurrence		OR		AND		Even-Pair		ICL bigram		SLR
Training plan		Train Loss	Test Acc.	Train Loss	Test Acc.	Train Loss	Test Acc.	Train Loss	Test Acc.	Train Loss	Test Acc.	Train Loss	Test Loss
Coupled with rand init.	Joint GD	$10^{-3}$	$1.00$	$10^{-3}$	$1.00$	$0.04$	$0.99$	$0.01$	$1.00$	$10^{-3}$	$1.00$	$0.01$	$0.01$
Decoupled 1/2/3-phase with random init.	Attention (A) only	$1.64$	$0.50$	$1.67$	$0.50$	$0.24$	0.94	$0.69$	$0.51$	$3.90$	$0.02$	$0.50$	$0.50$
	Output (O) only	$0.74$	$0.86$	$0.68$	$0.85$	$0.79$	$0.75$	$0.63$	$0.50$	$2.70$	$0.16$	$0.48$	$0.48$
	A(1x) $\to$ O	$0.72$	$0.76$	$0.69$	$0.75$	$0.30$	$0.91$	$0.62$	$0.50$	$2.58$	$0.26$	$0.41$	$0.43$
	O(1x) $\to$ A(1x)	$0.47$	$\mathbf{1.00}$	$0.64$	$0.78$	$0.52$	$0.87$	$0.52$	$0.50$	$3.20$	$\mathbf{1.00}$	$0.28$	$0.28$

Table 1: Train loss and test accuracy (loss for SLR ), with mean over

50

seeds . Standard deviations are reported in Appendix B.4 . Values below

10^{-3}

are reported with their upper bounds; those that match or surpass the joint GD baseline are bold. For AND , all attention-training phases have multiple steps instead of one.

Figure 3: Attention and output gaps on four tasks. Mean

\pm

standard deviation over

50

runs; dotted lines mark the ideal gaps (see Appendix B.2 ). Phases are defined as in Figure 2 ; single-layer plans end after Phase 1 and A

\to

O after Phase 2. For AND , we plot the head with the largest gap for each plan (all heads: Appendix B.3 ). For even pairs and SLR , the output gap is

\cos(u,w^{*})

and

\cos(v,v^{*})

, with

w^{*}

the max-margin and

v^{*}

the ground-truth direction. Since

w^{*}

depends on the attention weights, the output gap on even pairs can change in Phase 2 although the output is frozen. O

\to

A and the zero initialization are omitted, as they coincide with the first two phases of O

\to

\to

O and with the random initialization, respectively. The trajectories on OR resemble those on occurrence in Figure 2 and are deferred to Appendix B.5 .

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Fixed	Random
Occurrence, OR	$\boldsymbol{\alpha}=\boldsymbol{\beta}=\mathbf{0}_{d}$	$\boldsymbol{\alpha},\boldsymbol{\beta}\sim\mathcal{N}(0,0.1^{2})$
AND	$\boldsymbol{\alpha}^{(h)}=\boldsymbol{\beta}^{(h)}=\mathbf{0}_{d}$ for $h\in[4]$	see text
Even pairs	$W=\mathbf{0}_{D\times D},\ u=\mathbf{0}_{D}$	$W,u\sim\mathcal{N}(0,0.1^{2})$
ICL-bigram	$A=B=\mathbf{0}_{d}$	$A,B\sim\mathcal{N}(0,0.1^{2})$
SLR	$k=v=\mathbf{0}_{d}$	$k,v\sim\mathcal{N}(0,0.1^{2}/d)$

Appendix

Table 2: Initializations. The fixed initialization is used only for O

\to

\to

O; Gaussian entries are i.i.d.

Table 6

Figure 4: The learning progress of the AND task for each head, with average trajectory and std (shade). The heads are ranked by the gap they reach at the end, per head and per run.

Training plan		Occurrence		OR		AND		Even-Pair		ICL bigram		SLR
Training plan		Train Loss	Test Acc.	Train Loss	Test Acc.	Train Loss	Test Acc.	Train Loss	Test Acc.	Train Loss	Test Acc.	Train Loss	Test Loss
Coupled with rand init.	Joint GD	$10^{-3}$	$1.00$	$10^{-3}$	$1.00$	$0.04$	$0.99$	$0.01$	$1.00$	$10^{-3}$	$1.00$	$0.01$	$0.01$
		( $10^{-6}$ )	( $0.00$ )	( $10^{-5}$ )	( $0.00$ )	( $0.06$ )	( $10^{-3}$ )	( $10^{-4}$ )	( $0.00$ )	( $10^{-7}$ )	( $0.00$ )	( $10^{-3}$ )	( $10^{-3}$ )
Decoupled 1/2/3-phase with random init.	Attention (A) only	$1.64$	$0.50$	$1.67$	$0.50$	$0.24$	$\mathbf{0.94}$	$0.69$	$0.51$	$3.90$	$0.02$	$0.50$	$0.50$
		( $0.08$ )	( $0.01$ )	( $0.06$ )	( $0.01$ )	( $0.18$ )	( $0.07$ )	( $10^{-3}$ )	( $0.02$ )	( $0.00$ )	( $0.01$ )	( $0.02$ )	( $0.02$ )
	Output (O) only	$0.74$	$0.86$	$0.68$	$0.85$	$0.79$	$0.75$	$0.63$	$0.50$	$2.70$	$0.16$	$0.48$	$0.48$

Appendix

Table 5: Train loss and test accuracy (loss for SLR ) as mean over

50

seeds, with the standard deviation in parentheses. Values

\leq 10^{-3}

are shown in upper bounds; those that match or surpass joint GD are bold. The red entries are deterministic runs because the training and test sets are the entire instance space.

Figure 5: Attention and output gaps on occurrence and OR . Notation and phases are as in Figure 3 .

Explore similar work

Incremental Learning of Sparse Attention Patterns in Transformers
Feb 22, 2026Oğuz Kaan Yüksel, Rodrigo Alvarez Lucendo, Nicolas FlammarionSelf-Attention Transformer Attention
TML Lab, EPFL
Characterizing the Expressivity of Local Attention in Transformers
May 1, 2026Jiaoda Li, Ryan CotterellTransformer Expressivity Transformer
A Controlled Study of Attention-Only Transformers
Jul 20, 2026Henry Ndubuaku, Karen Mosoyan, Jakub Mroz +5Transformer FFNs Transformer
Cactus Compute, Inc.

Feb 22, 2026J/K move · Enter open · S save

Incremental Learning of Sparse Attention Patterns in Transformers

Oğuz Kaan Yüksel, Rodrigo Alvarez Lucendo, Nicolas Flammarion

TML Lab, EPFL

Self-Attention Transformer Attention Sparse Attention Incremental Learning Attention Mechanisms Neural Network Training Dynamics+2

Open paper AskPDF arXiv

This paper studies simple transformers trained on a high-order Markov chain, where the model must incorporate information from multiple past positions, each with different statistical importance. We show that transformers learn the task incrementally, with each stage corresponding to learning how to copy information from a subset of positions via a sparse attention pattern. Notably, the learning dynamics transition from a competitive phase, where all heads focus on the statistically most important positions, to a cooperative phase, where different heads specialize in different patterns. We model these dynamics with simplified differential equations and prove stage-wise convergence of the resulting system. Functionally, these stages correspond to a sequence of increasingly expressive misspecified models, with the full model class reached only at the end. Overall, we give a theoretical account of how structured attention patterns and head specialization emerge in stages without an explicit curriculum, with implications for generalization in sequential tasks.

Abstract

Explore similar work

Incremental Learning of Sparse Attention Patterns in Transformers

Characterizing the Expressivity of Local Attention in Transformers

A Controlled Study of Attention-Only Transformers