What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformers

We thank the reviewer for their positive comments, and are glad to know they found our paper well written and easy to follow. We address their concerns below.

Weaknesses:

1. We would like to emphasize that our work contributes towards a unified understanding of abrupt learning in various algorithmic tasks for Transformers, through novel results about representation collapse in hidden states and repetitions in the output. We further show how learning attention map is a learning bottleneck, and causally verify this via training time interventions. Importantly, we also show how our findings about representation collapse and repetitions are applicable to LLM pretraining, highlighting the relevance of our work. Finally, these empirical results form the basis for future theoretical investigation into such phenomena in Transformer training.

a) “Why is this phenomenon emerging…”: As we also note in the concluding discussion in the paper, this is an important direction for future work. We present some hypotheses for why these phenomena are occurring:

Hypothesis for slow learning of attention map:

Our hypothesis for loss plateau is that the attention layer first searches and combines the correct input tokens, and subsequently during training the MLP layer and LM head learn to map those hidden states to the correct output tokens. This is verified via non-linear (2-layer ReLU MLP) probing the outputs of Attention block during various points in training, and we find that such a probe is able to map the hidden states to the correct output labels earlier in training than the final model output (more details below). We note that this is related to prior work on the “search-then-fit” mechanism in 2–layer neural nets [Abbe 23], where the first layer in the neural network searches for the support variables of the target function in the input vector, and subsequently the 2nd layer fits the target function using these support variables.

Probing details: We probe hidden state outputs $h_i$ of the Attention block corresponding to the output tokens, before they are added to the residual stream. We map these hidden states to the correct output tokens (i.e. $y_i$ as defined for MWS) using a 2-layer ReLU MLP probe ${R}^{256} \ni x \mapsto W_2 \sigma(W_1 x) \in {R}^{17}, W_1 \in R^{256 \times 256}, W_2 \in R^{17 \times 256}$ mapping to logits over relevant vocabulary {0,...,16} (no SEP=17 token since that is not relevant for $y_i$). We implement the probe using MLPClassifier from the sklearn.neural_network library, train it using 1024 samples, and evaluate its performance using a held-out test set of 1024 samples.

We will add this experiment to the updated version of the paper; a deeper theoretical understanding of this phenomenon is an important direction for future research.

Hypothesis for representation collapse: We hypothesize that representation collapse arises from a bias of Transformers towards repetitive outputs in autoregressive setup; towards this hypothesis, we showed that training on repetitive sequences (hence explicitly encouraging representation collapse / repetition bias) is easier in the sense of training loss decreasing rapidly and monotonically, without any loss plateaus (Sec 5.2).

b) “Can this be tackled…”: We have found that abrupt learning, representation collapse and repetition bias occur robustly across various tasks, hyperparameter choices and optimization settings, and how to tackle these phenomena without directly intervening on the attention map is an important question for future work.

c) "…results are unsurprising… uniformly distributed mixing of input tokens.”: As we have also pointed out in Line 145 and 156, repetition bias / representation collapse phenomena are not present at initialization, and hence are not due to uniform attention weights at initialization. Instead, it emerges in the initial ~50 steps of training, which indicates it is a consequence of early phase model training. To quantitatively verify this, note that in Fig 2(b) representation collapse (cosine similarity) increases from 0.1 to 0.9, while repetition frequency also increases from around 0.1 to 0.8 in the first 50 steps of training.

2. “Very limited setup … strong assumptions …”: We note that our paper includes results on multiple tasks and model configurations (number layers, number of heads, etc.) for standard decoder-only Transformer architectures used in practice, confirming a unifying underlying behavior in all these cases. We also use a realistic next-token-prediction training setup with cross entropy loss and Adam optimizer, as is used in practice. In response to reviewer U6ru, we have also performed multiple ablations on the optimization setup and found abrupt learning, representation collapse and repetitions in all cases, highlighting that the phenomena are not limited to a single task / model / training setup. Our results on representation collapse and repetitions in LLM pretraining also demonstrate the generality of our findings in the toy model / data setup.

   “In Line 292 they claim…": We would like to point out that we have not claimed to focus on mitigating repetition in LLM in our paper, and line 292 is part of the ‘Related Work’ section discussing other prior work on repetitions in LLM.

3. “...claim relies on a single observation in MWS…” : The partial solution claim is not based on a single task, and we have indeed included evidence for 4 more algorithmic tasks in the appendix (Table 1 + Appendices B.1-B.6), also pointed out in Lines 136-138. Other recent work has also shown similar results for partial solutions in n-gram estimation tasks with Transformers [Varre 25], corroborating our findings about partial solutions during Transformer training.

   To further illustrate this point, we train the Transformer on concatenated sequences of increasing complexity, as below,

   $x_1, x_2, \ldots, x_{16}, {\rm SEP}, y_1, y_2, \ldots, y_{13}$

   $y_i = \begin{cases} x_i & i\in [1, 4] \\\\ (x_i+x_{i+1}) \mod 17 & i\in [5, 9] \\\\ (x_{i+1}+x_{i+2}+x_{i+3}) \mod 17 & i\in [10, 13] \\\\ \end{cases}$

   I.e. a concatenation of Copy, MWS over 2 variables (MWS2) and MWS over 3 variables (MWS3). We find that there are separate plateaus during training, each corresponding to learning a part of the above concatenated sequences (similar to Figure 1 in [Abbe 23], and Figure 1 in [Varre 25]) – highlighting partial solutions based on “difficulty” of the respective sub-sequence.

4. We used the first 100 questions of the ARC-easy test dataset, and the last layer hidden states for all models. Further, we also measured the repetition frequency ($\rho$) of Pythia outputs, and found that it follows similar trends as cosine similarity (like the MWS case). Concretely, the table below gives the evolution of repetition frequency during the first 10000 training steps, across the 4 Pythia models discussed in the submission.

We will add these results through a plot similar to Fig. 8 in the updated version of the manuscript.

5. While we agree that it would be ideal to have all the results in the main paper, we have prioritized discussing the main results on a single task (MWS) in the main paper for clarity, and postpone details about other tasks and ablations to the appendix.

6. We will provide code to replicate all experiments at a later stage.

Questions:

1. We note that repetition frequency ($\rho$) is not a flawed measure, but a clearer / stronger one since it directly measures contiguous repetitions (i.e. substrings of the form v,v,v,… for some v in vocabulary) instead of non-contiguously over the full sequence. By definition, a sequence with high repetition frequency ($\rho$) would also have a correspondingly low seq. entropy. Seq. entropy was only used to show that repetition may exist in other non-obvious forms as well, and we have not ‘changed it afterwards’ in that results for the other algorithmic tasks (i.e. except prefix sum) in the appendix still use repetition frequency ($\rho$).

2. For each step in training after $t_0$, for attention map $A \in \mathbb{R}^{(L-1) \times (L-1)}$, we use the mask $M \in \mathbb{R}^{(L-1) \times (L-1)}, M_{ij} = c$ for $(i,j) \in \Omega, M_{ij} = 1$ otherwise, except for $i=1$ (since that is a partial solution and converges early in training). Then, we use the modified attention map $A \odot M$ (elementwise product) for training and inference. We will elaborate on this in the updated version of the manuscript as well.

3. We will make necessary changes in plot styling / color palette in the updated version.

4. Token ID = integer as in the sequence, and token ID for [SEP] is already clarified wherever a task is defined. We will clarify this in the experiments section as well.

We hope that our response has clarified the reviewer’s concerns about our work.

[Abbe 23]: Abbe et al., 2023. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. COLT 2023.

[Varre 25]: Varre et al., 2025. Learning In-context n-grams with Transformers: Sub-n-grams Are Near-Stationary Points. ICML 2025.

We thank the reviewer for their positive feedback on our work, and are happy to know they found our findings intuitive and interesting! We address the reviewer's concerns below.

Weaknesses:

1. We only found 1 plateau in MWS (after partial solution plateau) - once it escapes the plateau, it converges to the optimal solution. Similarly, we did not observe multiple plateaus in the tasks we consider in the paper i.e. once the model exits the first significant plateau, it converges to the optimal solution.

   To address this issue, we also discuss training on a target function that shows staircase-like learning with multiple loss plateaus, each plateau corresponding to learning a (progressively more ‘difficult’) part of the sequence. Specifically, sequences of the form,

$x_1, x_2, \ldots, x_{16}, {\rm SEP}, y_1, y_2, \ldots, y_{13}$

$y_i = \begin{cases} x_i & i\in [1, 4] \\\\ (x_i+x_{i+1}) \mod 17 & i\in [5, 9] \\\\ (x_{i+1}+x_{i+2}+x_{i+3}) \mod 17 & i\in [10, 13] \\\\ \end{cases}$

I.e. a concatenation of Copy, MWS over 2 variables (MWS2) and MWS over 3 variables (MWS3). We find that there are separate plateaus during training, each corresponding to learning a part of the above concatenated sequences (similar to Figure 1 in [Abbe 23], and Figure 1 in [Varre 25]) – highlighting partial solutions based on “difficulty” of the respective sub-sequence.

We found that in this case, the tokens which are not yet predicted correctly during each plateau, also show representation collapse i.e. have high pairwise cosine similarity, indicating that the observations are not limited to a single loss plateau.

2. We provide empirical justification for gradient based training by showing that at initialization, the model does not have repetition bias or representation collapse, but only appears after a few (~20-30) steps of training. Please see Fig 2 where the cosine similarity is quite low at the initial stages and rapidly increases during the early phase of training, and similarly for repetition frequency. This highlights that these pathologies are not an artifact of random initialization, but that the training process leads to these issues in the early stages. Understanding why these issues occur at the early stage of training is an important direction for future research, as we also note in our concluding discussion.

3. Indeed, the main idea behind our work is to show dependence on attention and how it might cause a loss plateau, which is why the tasks are designed to depend crucially on learning the correct attention map. Since one of the key features of transformers is the attention mechanism that appropriately selects relevant tokens for computing the output, we focus on this aspect in our work on learning dynamics as well.

4. We agree, although the main idea is not smoothness but that there is progress before the sudden drop, evident in the increase in attention progress measure before the sudden drop in loss. We also note that for other tasks, APM increases continuously throughout training.

5. We will provide code to replicate all experiments at a later stage.

Comments:

1. , 3., 4. We will make the necessary changes as suggested by the reviewer in the updated version of the manuscript.

2. For each step in training after $t_0$, for attention map $A \in \mathbb{R}^{(L-1) \times (L-1)}$, we use the mask $M \in \mathbb{R}^{(L-1) \times (L-1)}, M_{ij} = c$ for $(i,j) \in \Omega, M_{ij} = 1$ otherwise, except for $i=1$ (since that is a partial solution and converges early in training). Then, we use the modified attention map $A \odot M$ (elementwise product) for training and inference. We will elaborate on this in the updated version of the manuscript as well.

Questions:

1. Yes, we work in an online training setup (i.e. sample a new batch of 256 samples from the distribution at every step of training), so train loss = test loss (also noted in lines 104-107).

2. We agree that there is a correlation between representation collapse and repetitions, though would also like to point out that in the multi-digit addition task (Appendix B.1), repetition frequency remains quite low (~0.5) despite the cosine similarity being ~= 1.0. Also note that the cosine similarity is close to 1.0 but not exactly 1.0, so some hidden state representations might not exactly coincide, leading to lower repetition frequency.

3. The first output token $y_1$ is the partial solution for the MWS task, and converges to the correct output early on in training. Hence, there is no representation collapse for that token (with the other output tokens $y_2, \ldots, y_n$ that are still not correctly predicted), and hence we do not use it for our hidden state analysis.

[Abbe 23]: Abbe et al., 2023. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. COLT 2023.

[Varre 25]: Varre et al., 2025. Learning In-context n-grams with Transformers: Sub-n-grams Are Near-Stationary Points. ICML 2025.

We thank the reviewer for their positive evaluation of our work! We address their concerns (about partial solution, attention progress measure, etc.) below.

Weaknesses:

1. While we agree that a partial solution may not be unique for a given task (i.e. multiple partial solutions could be defined for the same task), we have focused on partial solutions that occur consistently during training in our setup, and also found that they can be understood intuitively w.r.t. the original task. In other words, a partial solution is simply the capability the model has already learned before the sudden drop in loss.

2. “...APM metric is not readily available in most tasks…”: We provide the optimal $\Omega$ for all tasks considered in our paper below.

   Here we assume the attention map is of the shape $A \in R^{(L-1) \times (L-1)}$ in the next token prediction setup on the full input sequence. Hence, for {MWS, Prefix sum, Histogram, Reverse, Copy}, $L=33$, for permutation $L=50$, and for Multi-digit addition $L=15$. Further, the $(i, j)$ below are in the set $[L-1] \times [L-1]$ to denote row / column of $A$.

• “...just assesses whether the attention has the right support, not the optimal attention values…”: We note that APM is intended to measure attention score magnitude at the correct (support) token positions relative to other positions, and can be considered a proxy to measuring attention probabilities in softmax Attention. The overall idea is that the model should attend to the correct tokens in the sequence, and not the other (irrelevant) tokens, and this requires that the attention score has higher magnitude at the support positions relative to other tokens. We note that for MWS the APM reaches 0.8 highlighting that the attention map attends to the correct tokens with ~80% of available attention scores.

We further note that the optimal attention map in our tasks is usually sparse, and hence measuring its support is already very informative of the progress. Moreover, the optimal attention score values may not be unique, so measuring the magnitude of these values without taking into account the support would not be too useful.

Questions on APM:

1. Yes indeed, we have reported APM dynamics for all tasks, as well as the respective attention heatmaps in Appendix B.1-B.6. We have summarized APM sets $\Omega$ for all tasks in response to Weaknesses.

2. Please see our response to weakness regarding APM above.

3. We note that APM is more a property of the attention map (hence model architecture), which is why we chose to define it with the architecture setup for clarity. COS is for hidden states and $\rho$ for output token sequences, and we chose to define them with the associated collapse / repetition phenomena. A conceptual distinction was not intended, since all these quantities are ultimately used for understanding the training dynamics.

Miscellaneous:

4. We agree and will replace degeneracy by a more suitable word as suggested.
5. and 6. We will make the required changes in plot styling / color palette in the updated version.

We hope our response has addressed reviewer concerns.

We thank the reviewer for their encouraging comments, and are happy to know they found our work well written! We address the reviewer’s concerns (about hypotheses for observed phenomena, robustness to hyperparameters, etc.) below.

Weaknesses:

1. Hypotheses for observed phenomena

Slow learning of attention map:

Our hypothesis for loss plateau is that the attention layer first searches and combines the correct input tokens, and subsequently during training the MLP layer and LM head learn to map those hidden states to the correct output tokens. This is verified via non-linear (2-layer ReLU MLP) probing the outputs of Attention block during various points in training, and we find that such a probe is able to map the hidden states to the correct output labels earlier in training than the final model output (more details below). We note that this is related to prior work on the “search-then-fit” mechanism in 2–layer neural nets [Abbe 23], where the first layer in the neural network searches for the support variables of the target function in the input vector, and subsequently the 2nd layer fits the target function using these support variables.

[Probing details]: We probe hidden state outputs $h_i$ of the Attention block corresponding to the output tokens, before they are added to the residual stream. We map these hidden states to the correct output tokens (i.e. $y_i$ as defined for MWS) using a 2-layer ReLU MLP probe ${R}^{256} \ni x \mapsto W_2 \sigma(W_1 x) \in {R}^{17}, W_1 \in R^{256 \times 256}, W_2 \in R^{17 \times 256}$ mapping to logits over relevant vocabulary {0,...,16} (no SEP=17 token since that is not relevant for $y_i$). We implement the probe using MLPClassifier from the sklearn.neural_network library, train it using 1024 samples, and evaluate its performance using a held-out test set of 1024 samples.

We will add this experiment to the updated version of the paper; a deeper theoretical understanding of this phenomenon is an important direction for future research.

Representation collapse: We hypothesize that representation collapse arises from a bias of Transformers towards repetitive outputs in autoregressive setup; towards this hypothesis, we showed that training on repetitive sequences (hence explicitly encouraging representation collapse / repetition bias) is easier in the sense of training loss decreasing rapidly and monotonically, without any loss plateaus (Sec 5.2).

2. Robustness to hyperparameters: We show that abrupt learning occurs across a range of optimization hyperparameter choices as below.

To summarize, results for all hyperparameter and optimizer choices have abrupt learning, representation collapse and repetitions, with varying length of plateau / extent of representation collapse depending on learning rate, batch size, initialization scale, choice of optimizer. Note that for all ablations, the other hyperparameters remain fixed to default values as described in the paper.

• Optimizer: We have already reported results on SGD (that are similar to the Adam case) in our submission (Fig. 26 in Appendix E). In addition, we experimented with

a) AdamW (lr=1e-4, wd={1e-3, 1e-2, 5e-2, 1e-1}) and

b) Muon (lr=0.02) with {Adam (lr=1e-4), AdamW (lr=1e-4, wd=1e-2)} for embedding, LM head and all bias parameters.

and found that the results are similar to those with Adam (abrupt learning, representation collapse and repetitions occur), although Muon appears to have a shorter loss plateau.

• Learning rate schedule: We experimented with {linear warmup, linear warmup with cosine decay}, with warmup steps {20, 50, 100, 200} for a total 400 training steps and did not find any significant changes in training characteristics. (Note that cosine decay leads to 0 learning rate at the end of training). We also note that prior work [Chen 24] used linear warmup in their training setup and observed abrupt learning.

Additionally, we also tested the learning rate schedule in the suggested paper on attention sink with some modifications to total steps and warmup steps for our setup (i.e., total steps=400, warmup steps=20, rest of the hyperparams like AdamW weight decay, min / max learning rate, beta1, beta2, gradient clipping remain the same as the suggested paper).

• L2 Regularization: We experimented with Adam with weight decay = {1e-3, 1e-2, 5e-2}, and found that while 1e-3, 1e-2 cases are similar to the unregularized training results, for wd=5e-2, the plateau is slightly longer while the peak cosine similarity is slightly lower (~0.8)

• Adam + constant LR: We experimented with learning rate = {5e-7, 1e-6, 1e-5, 1e-3} and did not find significant changes in training characteristics except that with smaller learning rates, the plateau is longer as expected.

• Batch size: We vary batch size = {8, 32, 1024, 4096} and did not find changes in training characteristics except that with smaller batch sizes, the plateau is longer as expected.

• Initialization scale: Default implementation initializes entries of all linear / embedding layers i.i.d. from $\mathcal{N}(0, 0.02^2)$ except for entries of $W_O, W_2$ matrices initialized i.i.d. from $\mathcal{N}(0, (0.02^2)/2L)$, for $L$ layer model ($L=1$ in our case). We vary the scale of initialization by replacing $0.02$ by $0.02 \alpha$ for $\alpha \in$ {1e-3, 1e-2, 1e-1, 1e1}, and find that for $\alpha < 1$, the loss plateau is longer and the cosine similarity goes to $\approx 1.0$ faster. Whereas for $\alpha = 10$, the cosine similarity peaks at around $0.6$ and the loss plateau is longer than $\alpha=1$ case.

(RC= Representation collapse)

Questions:

1. Please see response to Weakness 1 above.

• Hyperparameter / optimizer choice: please see response to Weakness 2.

• Initialization choice: We note that we do not set random seed for any experiment, hence this training behavior is not limited to a specific random initialization. The behavior however is sensitive to initialization scale i.e. variance of the normal distribution used to sample initial weights (as expected based on prior work in optimization theory for neural nets); please see initialization scale results in response to Weakness 2 above.

We hope that our response has resolved the reviewer’s concerns about the robustness of our results to hyperparameter and optimizer choices, and that our hypothesis about why abrupt learning occurs (supported by preliminary experimental results) has helped towards gaining a theoretical / mechanistic intuition for this phenomenon.

[Abbe 23]: Abbe et al., 2023. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. COLT 2023.

[Chen 24]: Chen et al., 2024. Sudden Drops in the Loss: Syntax Acquisition, Phase Transitions, and Simplicity Bias in MLMs. ICLR 2024.

[Gu 25] Gu et al., 2025. When Attention Sink Emerges in Language Models: An Empirical View. ICLR 2025.

[Loshchilov 19]: Loshchilov et al. 2019. Decoupled Weight Decay Regularization. ICLR 2019