Title: What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformers
Venue: NeurIPS 2025
Accepted: yes
Keywords: Abrupt learning

transformer training dynamics

interpretability

science of language models
Abstract: Training Transformers on algorithmic tasks frequently demonstrates an intriguing *abrupt learning* phenomenon: an extended performance plateau followed by a sudden, sharp improvement. This work investigates the underlying mechanisms for such dynamics, primarily in shallow Transformers. We reveal that during the plateau, the model often develops an interpretable *partial solution* while simultaneously exhibiting a strong *repetition bias* in their outputs. This output degeneracy is accompanied by *internal representation collapse*, where hidden states across different tokens become nearly parallel. We further identify the slow learning of optimal attention maps as a key bottleneck. Hidden progress in attention configuration during the plateau precedes the eventual rapid convergence, and directly intervening on attention significantly alters plateau duration and the severity of repetition bias and representational collapse. We validate that these identified phenomena—repetition bias and representation collapse—are not artifacts of toy setups but also manifest in the early pre-training stage of large language models like Pythia and OLMo.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors investigate the phenomenon of abrupt learning (i.e., the sudden learning of algorithmic tasks such as moving-window-sum, prefix sum, etc.) in shallow Transformer models. In this setting, they empirically observe different phenomena during the loss plateau preceding the task solution: collapse of the internal representations, output repetition bias, and convergence to partially correct solutions. Additionally, the authors show that some of the identified behaviors (e.g., repetition bias and representation collapse) can also be observed in the early stages of the pre-training of larger models, such as Pythia and OLMo.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nThe paper is overall well-written, and the narration is easy to follow. The authors evaluate a sufficiently wide choice of the algorithmic tasks (even if results for only one of them are included in the main text, see Weaknesses). The tasks and the [Transformer] models are very simple, allowing to make more tailored and precise observations on the nature of the observations. Furthermore, I also like that, at this level of complexity, it is possible to easily interpret some components of the model (such as the attention matrix) and intervene on them to validate hypotheses and conjectures (for instance, in the experiments on programming the correct attention mechanism vs. programming the optimal MLP/embeddings). \n\n\n**Weaknesses**\n\n1. The main weakness that I can identify in this paper is the relevance of the findings and the absence of a deeper investigation and understanding of them. For instance: I like the analysis on the internal representation collapse for MWS. However, why is this phenomenon emerging? Can this be tackled in a way which is not directly programming the “true” attention matrix for the model? The same goes for the repetition bias: I like how you somewhat exclude that this phenomenon purely emerges because of correlations in the data, but what is causing it then? Also, some of the results are unsurprising: the role of the attention mechanism in the representation collapse, for instance, since the attention weights at initialization will be uniformly distributed, resulting in an equally uniformly distributed mixing of the input tokens.\n2. More in general, I do not see how the observations in the paper could contribute to a better understanding of grokking-like phenomena. While the authors claim in the related work “to aim to understand a unifying reason behind such observations […] without restrictive assumptions on our model or training setup”, they only provide a set of (valid) empirical observations and, in fact, use a very limited setup, which could be interpreted as using strong assumptions. Furthermore, in Line 292 they claim to also focus on “how to mitigate such phenomena” (repetition in language), which is not the case in this paper.\n3. Claiming that the studied models converge to “interpretable partial solutions” in the early stages of the training is a bit stretched in my opinion. The claim relies on a single observation in MWS, where the model initially learns to copy the first token, which is probably not enough to generalize the statement so much. To claim this, the authors should include evidence at least all the other algorithmic tasks considered in the paper.\n4. The extrapolation of the results to LLMs is also handwavy. The statements on repetition bias are not supported by any quantitative evidence. Also, the experimental setup (e.g., from which layer are the hidden states extracted, how are the 100 questions from ARC-Easy selected, etc.) for the representation collapse are not sufficiently described in the manuscript.  \n5. The results are not self-contained in the main paper. This is not a major weakness, but I believe that spending some time to fit more complete results directly in the main paper (and not in the appendix) could improve the manuscript.\n6. It is not mandatory, but the authors could have included the experimental code in the submission. Given the simple nature of the setup, it would have been easy for me to replicate some of the experiments, which would have in turn increased my confidence in the reported results.\n\nQUESTIONS:\n1. Why did you not consider the sequence entropy as a measure of repetition bias from the beginning? If the number of tokens which are equal to the followings is high, I would also expect the entropy-based measure to capture this information. What is the point in introducing a flawed measure to then just change it afterwards?\n2. How do you bias the attention matrix? Lines 183-186 are not clear.\n3. It would be good to avoid color overlap between Figures 3 and 4 (which could be probably merged into a single figure as well?).\n4. How do you tokenize the inputs? This is an important detail missing from the paper.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work investigates what happens during the loss plateaus. They find that the model learns a partial solution that solves easier parts of the task (e.g., predicting the 1st element in the modulo cumulative task); Sec. 3. However, at the same time, there’s a repetition bias (models output the same token over and over again) and representational collapse (tokens are nearly parallel); Secs. 3 & 5. They reaffirm that the slow learning is due to the attention map learning (Sec. 4). While most experiments are done on algorithmic toy tasks, they validate their findings in early phases of LLM pretraining (Sec. 6).\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n* S1: The paper addresses an often overlooked part of training: what happens during early-training loss plateaus. Especially the connection between representational collapse and the loss plateaus is a meaningful finding.\n\n* S2: The findings are intuitive (this is a plus!) and interesting. In particular, the observation of repetition bias and the analysis of the role of attention learning are interesting (Sec. 4)\n\n* S3: The experiments are well-motivated, designed, and sound to support the findings. They allow fine-grained analysis to clearly work out the core mechanisms and behaviors.\n\n* S4: The analysis is interesting and a connection to LLMs is made.\n\n* S5: The paper is well-written.\n\n## Weaknesses\n\n* W1: It is unclear whether the observed phenomena (repetition bias) generalize to multiple loss plateaus or if it only happens for the first plateau. This would limit the paper’s findings. (Note that the title is a more general and I’d suggest to revise it, e.g., “What Happens During the *Early-Training* Loss Plateau?” or similar.)\n\n* W2: It remains unclear why transformers have repetition bias in the early stages of training. The authors argue that this is an “inductive bias coming from gradient-based training” (l. 146-147) but do not further elaborate this nor provide experimental evidence.\n* W3: It is somewhat expected that the attention hinders the algorithmic task to work well, as it implements the core mechanism for such tasks. It’d be interesting to also consider tasks that are less reliant on attention (though, I have to admit that I could not come up with one myself that is somewhat sensible).\n\n* W4: The attention progress measure is not always as smooth as for the MVS task (e.g., see the multi-digit addition task in Fig. 9 right).\n\n* W5: Code is not (yet) provided.\n\n## Comments\n\n* C1: References are often missing the publication venue.\n\n* C2: The attention biasing is a bit unclear and it would be good to elaborate it in a bit more detail.\n\n* C3: Equations should be numbered.\n\n* C4: The related work in the appendix puts those two works - https://arxiv.org/abs/2310.12956 and https://arxiv.org/abs/2309.07311 - in the category of “grokking” (delayed generalization, i.e., low train loss with late sudden jump of val/test loss). However, these works study “breakthroughs”, “phase transitions”, or “Eureka moments” (i.e., high train & val/test loss with sudden jump of train loss). Thus, it’d suggest discussing them in the main text (along with their more related work on abrupt learning).\n\nQUESTIONS:\n* Q1: Is it ensured that each sample is only seen once during training? (train loss = test loss?)\n\n* Q2: Is the repetition bias not just a result of representational collapse? I.e., as all hidden representations of tokens become parallel (i.e., same direction), the same token will be output. Fig. 2b would be in support of this (i.e., the curves for repetition bias and representational collapse follow a similar shape).\n\n* Q3: Why is the 1st token position not included in the averaging of $COS_{i,j}$?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe main text examines a shallow transformer trained on the **moving-window-sum (MWS)** algorithmic task, using this “fruit-fly” setting to probe the broader phenomenon of abrupt jumps in loss and accuracy during transformer training. (More architectural tweaks and additional algorithmic tasks are presented in the appendices.)\n\nIn MWS, a given input sequence $x_1,\\dots,x_n$ must be completed by producing $y_1 = x_1$ and, for $i \\ge 2$,\n$y_i = (x_i + x_{i-1}) \\bmod p$.\n\nThroughout training, the authors track quantitative metrics that capture early-stage **degeneracies**—notably repetition bias and representation collapse—which coincide with the well-known loss plateau.  By systematically tracking these signals **and demonstrating that targeted manipulations of attention scores can lengthen or remove the plateau**, the work provides initial empirical evidence that slow query–key alignment acts as a causal bottleneck, offering a concrete lead for future efforts to shorten such plateaus.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n* **Transparent setup.** Task, metrics, and training regimen are described clearly, so the experiments are easy to replicate.\n* **Causal interventions.** The paper’s *Biasing the Attention Map* and *Training with Optimal Attention* experiments in Section 4 show that directly manipulating the attention map lengthens or removes the plateau, supporting the bottleneck claim.\n* **Quantitative metrics.** Repetition rate, cosine-collapse, and APM turn vague “degeneracy” into reproducible numbers.\n* **Evidence is quite extensive.** The same patterns appear on several algorithmic tasks and are glimpsed in early LLM checkpoints, hinting at broader relevance.\n\n**Weaknesses**\n\n* The idea of a partial solution feels a bit arbitrary. A solution can be partial in myriad ways. \n* The APM metric is not readily available in most tasks as we might not know the optimal $\\Omega$. Furthermore, the metric just assesses whether the attention has the right support, not the optimal attention values, as far as I can tell.\n\nQUESTIONS:\n**Questions on APM**\n\n1. **Generality** – Can APM be computed for the other algorithmic tasks in Table 1?\n   If so, could you list the corresponding sets $\\Omega$ explicitly for each task?\n\n2. **Support vs. weights** – APM reflects only the *support* of the attention map, not its optimal weights. This makes it a rather blunt tool. (Is this why the curve in Fig. 2 (b) tops out at \\~0.8?)\n\n3. **Placement in the paper** – APM is introduced in §2, whereas the other metrics (repetition bias $\\rho$, representation-collapse **COS**) appear in §3.\n   Should we view APM as conceptually distinct from $\\rho$ and **COS**, or would it be clearer to define all three together?\n\n---\n\n**Miscellaneous**\n\n4. **Terminology — “degeneracy.”**  The term can sometimes have a precise mathematical/technical meaning, yet here it seems to label undesirable model behaviours (repetition bias, representation collapse, etc.).  Would a word like “pathology” or “failure mode” be clearer?\n\n5. **Styling** – Loss and accuracy are plotted with different styles in Figs. 1 and 2; please make them consistent.\n\n6. **Colour reuse** – In Fig. 2 the two panels recycle the same colours for different quantities, which can be confusing; consider distinct palettes.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the learning dynamics of Transformers on algorithmic tasks, focusing on the common pattern of a long performance plateau followed by a sudden improvement. The authors find that during the plateau, models develop partial solutions but suffer from repetition bias and representation collapse. These issues stem from slow learning of attention patterns, which, once resolved, lead to rapid gains in performance. The paper also shows that similar phenomena appear in early training of LLMs like Pythia and OLMo, highlighting their broader relevance.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. The paper is very well written, with good organization and clarity. \n\n2. The characterization of repetition bias, representation collapse, and hidden progress in attention learning provides a useful lens through which to understand transformer training dynamics.\n\n**Weaknesses**\n\n1. While the paper compellingly identifies phenomena like repetition bias, representation collapse, and attention as the underlying bottleneck, it doesn't go deeper into why this bottleneck arises. I understand this is noted by the authors as a future direction, but without at least some theoretical framing or hypothesis testing, the insight feels limited. The paper would be stronger with a toy model or a mechanistic explanation.\n\n2. Some of the findings feel like they might be sensitive to training details. What happens if you change the optimizer, add regularization, or tweak the learning rate schedule? Right now it’s hard to tell whether the observed effects are fundamental or just artifacts of a particular training setup. A few ablations on this would really help clarify robustness. I’d be open to raising my score if this point is addressed convincingly in the rebuttal.\n\n- A good example see When Attention Sink Emerges in Language Models: An Empirical View https://arxiv.org/abs/2410.10781.\n\nQUESTIONS:\n1. Can you provide any hypotheses or theoretical intuitions for why attention learning is slow or why representation collapse occurs during the plateau phase? \n\n2. How sensitive are the observed dynamics to initialization and optimizer choice? It would be helpful to know whether these behaviors are optimizer- or hyperparameter-specific, or if they robustly generalize across settings.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}