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 As=Σ∗sΣ∗ with alphabet Σ={1,…,d} and unknown s∈Σ . The language can be expressed by a transformer with one attention head and linear output layer ( 1 ), which encodes word w=(w1,…,wm) as (ew1,…,ewm) , where ei∈Rd is i -th standard basis vector. The optimal shallow transformer has weights of the output layer β=2es∈Rd , and attention score for symbol i∈Σ is γi=1 if i=s , and 0 otherwise. (b) A shallow transformer is trained to recognize As using gradient descent jointly on both layers. The output gap (βs−i=smaxβi)/2 and attention gap (γs−i=smaxγ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) → O, does not correctly learn the language. We later prove that a 3-phase training, O(1x) → A(1x) → O, correctly learns to recognize As . (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 5d2 training and d2 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
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) → 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) → A(1x)
0.47
1.00
0.64
0.78
0.52
0.87
0.52
0.50
3.20
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 ± 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 → 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 → A and the zero initialization are omitted, as they coincide with the first two phases of O → A → 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
α=β=0d
α,β∼N(0,0.12)
AND
α(h)=β(h)=0d for h∈[4]
see text
Even pairs
W=0D×D,u=0D
W,u∼N(0,0.12)
ICL-bigram
A=B=0d
A,B∼N(0,0.12)
SLR
k=v=0d
k,v∼N(0,0.12/d)
Appendix
Table 2: Initializations. The fixed initialization is used only for O → A → 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
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
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 ≤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 .
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.
Oğuz Kaan Yüksel, Rodrigo Alvarez Lucendo, Nicolas Flammarion
The transformer is the most popular neural architecture for language modeling. The cornerstone of the transformer is its global attention mechanism, which lets the model aggregate information from all preceding tokens before generating the next token. One common variant of attention is called local attention, which restricts each token to aggregating information from a bounded window of predecessors, reducing the quadratic cost of global attention to linear. Although this restriction is usually motivated by efficiency, it has also been found to improve model quality, a phenomenon that has so far lacked a satisfactory explanation. We provide a formal account of this phenomenon in terms of recognizer expressivity. It has been shown that fixed-precision transformers with global attention correspond to a fragment of linear temporal logic containing a single past operator. We additionally prove that adding local attention introduces a second temporal operator, strictly enlarging the class of recognizable regular languages. Moreover, global and local attention are expressively complementary: neither subsumes the other, and combining them yields the richest fragment. Experiments on formal language recognition and natural language modeling corroborate the theory, showing that hybrid global--local transformers outperform their global-only counterparts.
Feed-forward networks hold two thirds of a transformer's non-embedding parameters, yet the architecture has not received a necessity test that controls parameters, compute, and depth at once. We pretrain attention-only decoder transformers (Simple Attention Networks, SANs) against standard transformers matched separately for parameter count, training FLOPs, and depth (2 to 48 layers), for up to 105B tokens at 6M to 87M parameters. Deleting feed-forward layers in place is costly: the standard transformer leads by 0.47 nats at matched depth and 0.26 nats at matched FLOPs. Reallocating the freed budget into attention depth closes the gap: at matched parameters the difference is 0.006 nats (0.27 percent of loss), reproducible to one part in ten thousand across seed pairs, shrinking across 5B, 30B, and 105B budgets, and holding near 0.02 nats across a 29x size range. Three measurements localize the remaining gap to parametric recall: attention-only models are better on context-grounded answers and worse where knowledge must come from weights. Weight spectra show why: routing matrices (Q/K) crystallize early, content matrices accumulate rank slowly, and removing feed-forward layers relocates this accumulation to the attention output projection. QK-normalization, not feed-forward layers or residual gating, keeps 48-layer attention-only stacks trainable. The deficit concentrates on low-context query prediction and localizes there entirely by the largest budget. A pre-registered test confirms the account: it predicts a 0.02 to 0.05 nat gap on knowledge-dense web text; a matched pair trained on fineweb-edu measures 0.040. Within the tested regime, attention does the rest.