Training Dynamics of Transformers on Unbounded Inputs: Attention Learning Can Lag Output Layer Learning
This paper asks how shallow transformers learn on unbounded-length inputs and whether layer-wise training can be decoupled. It proves that a one-step attention-then-output schedule fails to learn the occurrence language, while a three-phase output-attention-output schedule succeeds in both population and finite-sample settings. The finding is significant because it formally explains why attention weights often lag behind the output layer during training, challenging the feedforward intuition that feature learning must come first.
The paper delivers a clean theoretical result: on the occurrence language over variable-length strings, a shallow single-head transformer can be trained to arbitrarily small risk by a specific three-phase gradient-descent plan, yet provably cannot be learned by training either layer in isolation or by an attention-first two-phase schedule. The proofs rely on careful gradient calculations and coupling constructions, and the authors augment the theory with experiments showing the three-phase plan broadly matches joint training on several tasks. However, the theoretical guarantees assume a simplified parameterization and zero initialization, and the empirical validation is limited to small-scale synthetic tasks, leaving open how well the insights transfer to realistic deep transformers.
The strongest aspect is the tight combination of negative and positive results. The impossibility theorems establish that no simpler schedule can succeed, which makes the three-phase positive result surprising and non-obvious. The analysis is explicit: the authors derive closed-form gradient expectations, use geometric-length sampling to balance classes, and provide a finite-sample corollary with $O(d\ln d)$ sample complexity. The experiments in Table 1 further show that the three-phase plan matches joint GD on occurrence, OR, ICL bigram, and SLR, supporting the authors' hypothesis that output-first learning is beneficial more broadly.
The theoretical model is highly simplified in ways that limit external validity. The authors parameterize the transformer directly by $(\boldsymbol{\alpha}, \boldsymbol{\beta})$ rather than by key, query, and value matrices, which removes over-parameterization and alters the implicit bias of gradient descent. The restriction to population risk for the main theorems, with only a finite-sample corollary and no stochastic GD analysis, is another gap. Empirically, the tasks are small synthetic language-recognition problems ($d = 50$, single-layer models), and the three-phase plan already fails to match joint GD on the AND task—a multi-head problem where attention-only training is competitive.
The theoretical evidence is carefully matched to the claims: each impossibility result precisely corresponds to a restricted training plan, and the positive result specifies exact step-size regimes. The comparison to feedforward one-step feature learning is fair and well-motivated, since the paper explicitly frames its contribution as testing whether that analogy carries over to transformers. The empirical section supports the general qualitative claim—attention often lags output learning—but also reveals limits: on even pairs the attention grows almost as fast as the output, and on AND the three-phase plan underperforms. The authors honestly report these deviations rather than over-selling the universality of the three-phase rule.
The paper does not mention a public code release or repository. All hyperparameters for the occurrence-language theory are stated explicitly (learning rates $\eta_1, \eta_2, \eta_3$ and number of steps $T$), and the finite-sample corollary gives exact sample sizes $N_1$ and $N_3$, so the theoretical protocol is reproducible on paper. However, the experimental details for the six tasks in Section 5 are spread across the main text and Appendix B.1; without released code, independently reproducing the exact training loops and gap metrics requires substantial manual reconstruction. The model itself is idealized—single-layer, no positional encoding, simplified softmax attention—so even with code, the gap to standard transformer implementations is large.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.