What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformers
Authors: Pulkit Gopalani
Wei Hu
Venue: NeurIPS 2025 poster ()
Keywords: Abrupt learning
transformer training dynamics
interpretability
science of language models
Primary area: deep_learning
TL;DR: Early phase training of Transformers on algorithmic tasks shows a plateau in loss, repetition bias and representation collapse before sudden drop in loss.

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.

What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformers
Pulkit Gopalani
Wei Hu
University of Michigan, Ann Arbor gopalani @umich. ch, edu
University of Michigan, Ann Arbor vvh@umich. edu
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 dynamies, primarily in shallow Transformers. We reveal that during the plateau, the model often develops an interpretable partial solution while simultane- ously exhibiting 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 optimal attention maps key bottleneck. Hidden progress in attention config- uration during the plateau precedes the eventual rapid convergence, and directly intervening attention significantly alters plateau duration and the severity of repetition bias and representational collapse. We validate that these identified phenomena reeentmena-repetiiooo bias and representation collapse- setups but also manifest the early pre-training stage of large language models like Pythia and OLMo.
not artifacts of toy
Introduction
Training Transformers mathematical algorithmic tasks often exhibits intriguing "abrupt learning" phenomenon in their training dynamics, where the model's performance plateaus at a suboptimal level for extended period before suddenly and rapidly converging the optimal solution (Figures]
[ [and This is often considered an example of the broader phenomenon of "emergence," where model capabilities appear to arise discontinuously and unpredictably with increasing amount of parameters, training data, or training steps [49]. Understanding these sharp phase transitions in learning trajectories crucial for gaining deeper insights into how Transformer models learn and develop their sophisticated capabilities.
3 36
47 54)
Despite recent progress in understanding such abrupt learning dynamics in Transformers for specific [7 (54), parity learning (3). Markov chains (13), (22, and matrix completion [18], a unifying account of the model evolution
tasks like in-context learning (10) 15 41
grammar learning during loss plateau is still missing. Further, many of these works require specific assumptions on model and data that limit their generality.
The goal of this paper is to uncover universal characteristics and underlying mechanisms that define these training dynamics that are broadly applicable to wide range of setups and tasks. Specifically, what common patterns manifest the model's 's input-output behavior and internal representations during the extended plateau phase, and what critical changes precede the sudden shift towards higher performance? Is there hidden progress accumulating beneath the surface of the loss platcau? Answering these questions is pivotal for building a comprehensive picture of the nature of phase transitions in Transformer training dynamics.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Train/Test Loss Accuracy
Before Sudden Loss
28-
10
24-
Map
18-
14
54
100
280
250
350
750
400
After Sudden Drop Loss
(h, Rd)
Hidber
h.
1,6.6,7.2211118888
12.
11. 1,, 10,
Example Outpur Soqwence
Feampio Outpur
Map
Figure 1: Abrupt learning and related characteristics. Training shallow Transformer algorith- mic tasks like moving-wiing nooingg-window-sig window-sum exhibits an abrupt learning curve: performance plateaus for an extended number of steps, before suddenly and sharply improving to optimum. Before the sudden drop loss, the attention map cannot be interpreted easily, whereas the post-sudden-dry attention map is clearly interpretable w.r.t. the task. Furthermore, the model exhibits degenerate patterns before the sudden drop, including output repetitions and collapse of its hidden representations.
To investigate these questions systematically and within a controlled setting, we focus on training layers) suite of simple algorithmic tasks. The
small, shallow Transformers (typically reduced model size allows for more tractable analysis and clearer interpretation of internal model mechanisms, avoiding the obfuscation that arise from the interplay of countless factors in large models. Furthermore, algorithmic tasks such as moviing-widdow-sin addition (as detailed later) have well-defined optimal solutions, thus allowing precisely measure the model's progress against a known ground truth and to readily interpret which aspects of the problem the model is succeeding that key findings from these controlled small-scale studies extend to the pre-training dynamics of actual Large Language Models (LLMs).
prefix-sum, and multi-digit
different training stages. Interestingly, will demonstrate
Our Contributions. We identify implicit biases that underlie the early plateau period of Transformer training: the model learns partial solution while being biased toward degenerate patterns in its outputs and internal representations. We further study the pivotal role attention map learning in driving these phenomena and overcoming the performance plateau. See Figure | for an overview of our findings. Our specific contributions are:
We show that during the initial loss plateau, the model often learns partial solution, which correctly predicts subset of easier tokens within a sequence- those that might be intuitively simpler to learn (e.g., copying the first element in a moving-windoo sum, or predicting the final carry-over multi-digit addition)-whhh parts of the task. This pattern is observed across diverse algorithmic tasks [T.
Partial solutions during plateau:
failing on more complex
Repetition bias in outputs: We identify strong repetition bias during the plateau, where the model tends to output repetitive tokens. This bias can be quantified by metrics such as direct count of repeated subsequent tokens the entropy of the output token distribution. Such repetitions significantly increase during the early training steps and then markedly decrease as the performance starts to improve (Figure (h).
The output repetition bias is accompanied by representation collapse, where hidden representations for different tokens become nearly parallel (e.g., cosine similarity often exceeds 0.9), indicating degenerate representational geometry inside Transformers. Subsequently, this representational similarity drops significantly as the model's performance improves, signifying diversification of internal representations (Figure 2b
Internal representation collapse:

---

Accuracy
Attention Progress Cosine Similarity
-07
0.8
08
Sm. RRppetoonn
-05
Costne Smiarry
Accumrry
Repeliton Frequency
Portaal Soluton Accurscy
Attantonn Progress Messur
-04
02
150
200
250
300
400
350
200
250
310
Staps
Steps
(a)
(b)
Abrupt learning dynamics for the MWS task. (a): Train/Test loss and Train/Test Accuracy (note that both train and data metrics are near-identical the online training setup, and thus we only report train metrics); (b): Attention Progress, Repetition Frequency, and Representation Cosine Similarity between hidden states. Increase attention progress is gradual and happens before the sudden loss drop. Repetition frequency and representation cosine similarity rapidly increase at beginning and decrease to low values later on.
Figure 2:
Crucial role of attention map learning: We find that gradual learning of the optimal attention pattern can commence during the loss plateau, before the sudden drop in loss (Figure 2b). By directly intervening scores towards or away from the optimal configuration we can observe tangible changes in the duration of loss plateau and the severity of degenerate behaviors like repetition bias and representation collapse.
the attention map during training for instance, by biasing the attention
Validation in LLMs: Our identified phenomena of repetition bias and representation collapse are not limited to small Transformers on synthetic algorithmic tasks. We further validate their occurrence in the early pre-training phases of L.MMs like Pythia and OLMo, suggesting these general characteristics of Transformer training dynamics.
Setup, Abrupt Learning, and Attention Map
We mainly present results for the moving noving-window-sse window-sum task in the main text, which we define below; we also validate our findings on various other algorithmic tasks like multi-digit addition, permutations, histogram, prefix-sum, etc., in Appendix
Data. The moving--wndoww (MWS) task involves computing the sliding- window sum (modulo p) of a length-n sequence over windows of size 2; that is, sequences in MWS are
31, I2 Tn, SEP, 1/1
Un
Φ (xi-1 x:;) mod
i =
y:
i2 2
Here, 31, .2......2 In are the input sequence, SEP a separator token, and the task is complete the sequence with outputs y...... In. In the experiments in the main paper, we use = 16, 17. Ii Unif (aif(1,27 2.. 16] and SEP 17. We denote the full vocabulary == {0, 1, 2, ... 17), and directly use these integers as token IDs when generating token embeddings for input to the model.
.....
We use a 1-layer, 1-head Transformer with causal masking and linear attention. This simple architecture can already solve the MWS task to perfect accuracy. Formally, for sequence of tokens .......,,,)), the Transformer output is,
Model Architecture.
TF . sL) LM o (Id MLP) (Id + Attn) Embed $1, .
(I)

---

where Embed outputs sum of token and absolute positional embeddings h1 E B and Attn denotes the Sieuulllinerr-tttee operation that combines tokens such that output at 7th position is,
(Attn(hi,ee be,...,, hL)]; Wo
Σ (nn WohhhW Wo: WK. Wo, Wy E Raadd (2)
MLP denotes the 2-layer neural net h; + (2 for W2 E TRdx W, E Riddd and the GELU activation. LM is linear layer that maps the hidden state h; E ER' to logits v E ERVI. Note that all linear maps above implicitly include a bias term, and we use pre-LayerNorm so that before the Attn, MIP, and LM, LayerNorm operation is applied to the hidden states hG. For generating sequences, we use greedy decoding i.e. output token is determined by the maximum logit over the vocabulary (please see Appendix A for more implementation details). We use linear attention to avoid small gradicnt issues from softmax function being a contributing factor toward abrupt learning, as argued in [20). We also show similar results on softmax attention, multi-layes / multi-head models, and models with varying d in Appendix F
Training, The model is trained to minimize the standard next-token- next-token-predictinn cross-entropy loss for the MWS task. We evaluate accuracy averaged over these n positions. We use the Adam optimizer with a constant learning rate 10-4 and no weight decay. The training is conducted in online / single-epoch fashion, where a new batch of 256 training samples is drawn from the data distribution at each training step. Note that in this setup, the training and test losses essentially coincide. For completeness, we verify in Appendix G that our reported phenomena occur when using different optimization algorithms (SGD, Muon (25)) or optimization hyperparameters (learning rate, LR schedules, batch size, initialization scale, weight decay) well, and are not merely an artifact of a specific training setup. Code for experiments is available at github. www. http///llitgggllllttt---eee plateau
ye............. over the full sequence i.e. (I),.... ..i....,fa..... SEP, y...... yn)
over the output portion of the sequence, i.e., y.,.... yn.
Abrupt Learning. Following the training procedure described above will result in a characteristic abrupt learning curve, where the training/test loss is stuck at some sub-optimal value for significant number of steps, before suddenly and rapidly decreasing to its optimal value (FigurePat This drop in loss is accompanied by a similarly rapid increase in accuracy, indicating that the optimal solution is learned abruptly.
Attention Map. We analyze the attention map at different points during training. We find that the attention map shows sparse, interpretable pattern after the sudden loss drop, while no such pattern is shown before the sudden drop (Figure D}. For the MWS task, this optimal attention pattern corresponds to each output token y; attending only to the input tokens relevant to computation, i.e. attending to 2 for V1, and to 24, I; for yt, i 2 2. We further use Attention Progress Measure (APM) to record the progress of the attention map toward its optimal pattern during training, defined as
APM = EiNeE |Ayl
E(.,) |Ayl
where Aij denotes the attention score allocated the jth token when computing output at the ith position in the sequence, and f2 is the set of position pairs in the optimal attention map. This measure defined with absolute values due to our choice of linear attention that A, could be positive or negative. In experiments, we calculate APM averaged over a random batch of sequences. The APM sets Ω for all relevant tasks in this paper are defined in Table 2
Figure shows that the APM monotonically increases from near 0 to near 0.8 during training, and its increase is more gradual than the loss/accuracy dynamics. In particular, APM already increases to a nontrivial value during the loss plateau and before the sudden loss drop.
3 Implicit Biases in the Early Phase of Training
In this section, we characterize several key manifestations of the implicit biases in the early phase of Transformer training. These patterns robustly co-occur with the loss plateau, and provide intuitive indicators that the model is getting stuck at degenerate state during the plateau.

---

Partial Solution. During the loss plateau, the model often has already learned to implement partial solution to the task. This means it correctly predicts a subset of the output tokens. typically those corresponding to an intuitively simpler part of the problem, while failing on the more complex parts. For instance, in the MWS task, the model quickly learns to predict the first output token yı correctly (see Figure 2a for the first-token accuracy). as it is simply copy of the first input token 31. while the overall loss remains high and accuracy on subsequent tokens is poor. This ability to solve easier sub-components of the task early on is observed across various algorithmic problems (see Table I in Appendix(D) B
Repetition Bias. Concurrent with learning the partial solution, the model's outputs during the initial phase of training display a strong repetition bias, which refers to a tendency of the model to generate repetitive tokens of the form z, I, z ...... One way to quantify such repetitions is to simply count the frequency of output tokens that equal the next one: for output sequence y1, V2...... Yn. define its repetition frequency as,
Σ 1ly 11+
We observe that p increases rapidly during the early phase of training, when the optimal attention map has not been learned yet (Figure 2b). Note that this frequency is small at initialization, and grows rapidly in the first x 50 steps to Rs 0.8, indicating that it is implicit bias coming from gradient-based training.
Representation Collapse. Motivated by the frequent repetitions in model outputs, we further study the relation between the hidden representations at different output positions. We find a strong representation collapse phenomenon theee representations become nearly parallel in the early phase of training (except for the first output position which is correctly predicted in the partial solution). We measure the pairwise cosine similarity between hidden representations at positions i, j in the output,
(h. by)
COS.J
(5)
where h, E Rd is the hidden state at position in the sequence (this quantity is averaged over a random batch of sequences). We find that in the early phase of training, there is rapid increase in C
over all output positions i, except the first position, this quantity increases to A 0.95 (Figure 2b). We emphasize that similar to repetitions, representation collapse is not present at initialization and only appears after a few steps of training. This is in contrast to the rank collapse phenomenon for deep adfimax-attention Transformers (2) that occurs at initialization. Also, while we focus on the final-layer representation (before the LMM layer) in the main text, we show in Figure 42 that representation collapse happens in all intermediate layers to varying degrees.
4
The Role of Learning Attention
Observe that though the loss dynamics are abrupt, attention progress measure as well as repetitions and representation collapse are not (Figure 2b); that is, even when the loss is barely decreasing (between steps 50 and 150), attention progress measure notably increases, accompanied by a decrease in repetition frequency and representation collapse. Via training-time interventions, this section shows that learning the attention map plays crucial role in shaping the loss plateau as well as repetitions and representation collapse.
Representation Collapse Occurs After the Attention Layer. We start by verifying whether the attention layer is responsible for representation collapse during the early phase of training. To this end, we plot the cosine similarity of the residual stream for output tokens just before and after the attention layer. Formally, let the residual stream before attention layer (i.e., token + positional embeddings) be h IR and the residual stream after attention layer be h E Rdd measure the norm and pairwise cosine similarity for h, and h in Figure 41
We find that in the early phase of training, the cosine similarity between different positions in the post-attention residual stream representations approaches 1.0 rapidly, which is not the case for pre- attention. Furthermore, the norm of hy grows rapidly in this phase, while the norm of h, remains
5

---

Train/Test Loss
Cosine Similarity
Repetition Frequency
1.0
2.75
Similarity
0.8
Repetition Freiq
0.6
2.25
2.5
2.00
044
0.4
1.75
1.50
0.0
0.0
200
300
400
300
400
300
400
Figure 3: Biasing attention map by 1. We find that multiplicative biasing the attention map towards more weight to optimal positions leads to faster convergence, accompanied by less repetitions and average cosine similarity.
near-constant. Hence, in the residual stream, representation collapse occurs after the attention layer during the early phase of training.
Biasing the Attention Map. To study the role of attention map, we slightly modify the training process starting at different time points in training, biasing it towards (or away from) the optimal attention map to check if repetitions, representation collapse, and loss plateau are reduced (resp. am- plified). We do the following: For each step in training after la. for attention map A E E RL--1)x(L-) we use the mask M E (tt-)))(--I), M, = for (i.j) 22, M,, (since that a partial solution and converges early in training). Then, we use the modified attention map A @ M (Hadamard product) for training and inference. Hence, for c > 1, this implies biasing the model towards the final (optimal) attention map, whereas for < < 1, this implies biasing the model away from the optimal attention map.
otherwise, except for i =
We find that, for 1 and various values of lo; such scaling leads lower average cosine similarity between hidden states, lower frequency of repetitions, and faster convergence (Figures 3 and 43). Whereas, for c I, we find the opposite: the model in representation collapse state for a longer time and converges later compared to the non-scaled (c = 1) case, while the repetition frequency remains large throughout the plateau (Figure 4
For example, for tu = 0, c = 10, i.e. scaling 10x from the start of training, we find that the peak cosine similarity attained during training is 2 0.6, much smaller than the z 0.95 attained for 1, and further the peak for e= 10 is for negligible duration compared to that for c == 1. Later values of 25, 50, 75 show similar results wherein the cosine similarity drops immediately above biasing operation, followed by lower repetition frequency and convergence to optimal solution (Figure 43). On the other hand, for to and is in representation collapse/ / large repetition frequency state for much longer. This is in line with our expectation that lower attention map values for the optimal positions lead slower learning and prolonged representation collapse.
the
c 0.2, 0.5, the model takes much longer converge
Hence, learning the optimal attention map has direct effect shaping the loss dynamics well as repetitions and representation collapse.
Train/Test Loss
Cosine Similarity
Repetition Frequency
2.50
Simllanyy.
0.8
0.0
0.6
Col
Repotitinn
2.00
0,2
0.4
0.4
1.75
02-
1.50
1.25
0.0
200
400
600
800
1000
200
400
600
800
1000
0
200
400
600
800
1000
Steps
Steps
Steps
Figure 4: Biasing attention map by < 1. We find that biasing the attention map to have lesser weight optimal positions leads slower convergence, and more representation collapse and repetitions.

---

Train/Test Loss
Cosine Similarityy
Repetition Frequency
3.0
Similarity
0.88
0.5
Emoeicg...)
brie Ii00
2.5
0.4
0.2
0.1
0.00
100
200
3.0
400
1.00
200
300
400
1000
200
300
400
Steps
Steps
Steps
Figure 5: Different optimal initializations and effect on training. We find that fixing attention and embedding weights (i.e. attention map) to optimal value, and training other components leads to faster convergence and lesser representation collapse repetitions. Similar effect does not hold for fixing optimal MLP or Embeddings. (K, Q,O, V respectively denote the parameters WK, Wo. Wo. Wy.)
Training with Optimal Attention Map v. Other Components In the first part of this test, we initialize with the optimal attention map by fixing embeddings, LayerNorm for attention layer and attention layer weights to their final values at the end of normal training run, that at initialization, the correct attention map is already available to subsequent layers. We re-train the subsequent non-fixed layers starting from random initialization. For the attention layer, choose the set of parameters to initalize in 2 ways: (a) only Key, Query (WK, Wo) weights, and (b) All of Key, Query, Value, Output (WK. K, Wo; Wy, Wo) weights. We find that in both of these cases, learning only the subsequent layers (i.e. MIP, LM Head) take significantly shorter time than training the full model, without any significant representation collapse, repetitions or plateau in loss (Figure [5)- Further, between (a) and (b), we find that additionally having Wo. Wv layers initialized to optimal values slightly speeds up learning, and average cosine similarity goes up approx 0.15 instead of 0.45 when only initializing W K, Wq weights. This indicates that Wo, Wy layers also play a non-trivial role causing representation collapse.
On the other hand, fixing MLP or embeddings (together with LM head) to their final optimal values and re-training the other components does not qualitatively change the training dynamics from the full training case, i.e., significant loss plateau, repetition bias, and representation collapse still occur (Figure 5)
This result confirms that attention map is a major bottleneck that leads to early representation collapse and loss plateau, and that there is little benefit from having the optimal MLP or embeddings at initialization compared to attention map.
Our results are similar to 44 that shows how various sub-circuits (inel. attention map) for in-context learning affect the loss plateau via training-tinn time interventions, and (29) that shows the effectiveness of attention transfer from pretrained model for downstream task training in Vision Transformers
Tracking Progress in Hidden States
To further track model progress during training, we probe the output of Attention layer before it is added to the residual stream (Equation ( using a 2-layer MIP probe with ReI.U activation, and find that the accuracy for this probe increases earlier than the model accuracy. Please see Appendix for more details.
Searching for the Right Tokens The preceding results have demonstrated the important role of learning attention map (i.e., attending to the correct tokens) towards the plateau in training loss. Furthermore, once there is sufficient progress in the Attention map (visualized via APM in Figure 2b and output probing in Figure 20), the final model output converges rapidly to the correct solution. These observations conceptually align with abrupt learning observed when training 2-layer neural nets on specific class of target functions (T). where the loss plateau is attributed to 'search' phase for aligning the first layer with the support of target function. In our setup, we hypothesize that a similar 'search' phenomenon could be occurring, where gradual increase in attention map weights at the required input tokens underlies the loss plateau.
Similar progress in model weights during loss plateau for a sinusoidal neuron trained on sparse parity task was shown in (3) Fig. 3]. (38) discuss parallel search when training multi-head Transformers on sparse parity task, and how the number of attention heads affects training dynamics.
7

---

Train/Test Loss Accuracy
Attention Progress Cosine Similarit laiity
Sequence Entropy
1.0
2.75
Aotorbern Progress
2.50
Sog Entrnoy
2.25
Loss
Accurios
Cosine Sim.
Curput Entro
2.00
Portial
Accuracc
0.4
Attemion Prog
Data Entreay
16
0.2
1.50
0.1
1.2
1.00
203
100
400
500
0
100
203
300
4.00
500
100
200
100
500 4:00
Stops
Steps
Figure 6: Prefix sum task training dynamics. While the usual contiguous repetitions do not occur for this task, an alternate form of repetition occurs in terms of having only a few distinct tokens in the output sequence. 'Sequence entropy' quantifies this repetition by measuring the entropy of the empirical distribution of tokens in sequence, and averaging this entropy over a batch of sequences.
5 A Further Look at Repetition Bias
Having observed that Transformer models exhibit a strong repetition bias in the early phase of training. which co-occurs with the loss plateau, we now take a further look this repetition bias and study how it might be affected by the amount of repetitions in the training sequences. We show that such bias still exists even when there is almost repetition in the training data. Subsequently. also show that learning simple, repetitive sequences is easier for Transformers, i.c., no loss plateau, indicating why the model might be biased towards repetitions in the early phase of training.
Beyond Repetitions in Consecutive Tokens One hypothesis for the reason behind repetition bias that the training data may consist of some repetitions, and the model may pick up these patterns and amplify them in the early phase of training. To investigate this, we consider a task with low repetitions in the training data. In particular, we consider the prefix sum task, where the outputs y1.....2 Yo are defined as y/ (Ei= zy) )mod p. Our choice of the input distribution ensures that there is no repetition in consecutive output positions (i.e.. V/s * Yi+1 for all i). Indeed, training a Transformer on the prefix sum task does not result in significant increase in the repetition frequency at any point in training, unlike the MWS task. Nevertheless, in the early training phase, we still observe that only a few tokens appear repeatedly in the model output though not contiguously as in the MWS task. Therefore, we consider an alternative measure of repetitions based on entropy: for an output sequence y1, 12. ......... Un. we define
WI SeqEnt ..... Un) Σ
p. log(1/p)); p. Ky = v. E [n]]]
(6)
ie. simply the entropy of the empirical distribution tokens in the sequence. Intuitively, the entropy is lower if most probability mass is concentrated at few tokens, and larger if the tokens are more uniformly distributed. We find that the model output entropy quickly goes to quite low values early in training compared to the entropy of ground-truth data (Figure 6). indicating that the model still has a form of repetition bias. Further, representation collapse still happens in the early phase, with the average cosine similarity going to 0.8 during the plateau. Hence, we find that repetition bias might take different forms depending on the task, but still robustly occurs in the early phase of training.
Repetitive Sequences are Easier to Learn We study what happens when training our model on data that has a lot of repetitions. We consider a simple task REPEAT, of the form ...,
T1.222..... Ln, SEP, V1, 1/2,... ...,,... M/. where yi 21 Vi. Unlike other tasks, the loss curve for REPEAT, does not have any noticeable plateau, though the accuracy still shows a small plateau period (Figure 7. This observation indicates that such repetitive sequences are easier from an optimization perspective and hence likely "preferred" during the early stage of training. In fact, just one gradient step is sufficient to bring the average representation cosine similarity to 0.5. We show similar results for other task variants (REPEAT, REPEAT,) in Appendix
To further understand the early training phase model output, we define another metric that measures to what extent the model simply outputs the same token for all output positions: on = En - 1/y = y1l. Note that this is distinct from accuracy, in that the model might output the wrong
8

---

Trair/Tstt Accuracy
Cosine Simllrtty Repeliiin
2.0
2.4
222
2.0
Thal: Acourayy
Smh
Repar Bon Free FirtTsiun Free (a)
1.0
14
Figure REPEAT, training dynamics.
31, however repeats Y1 at .2..... Un. We find that a; rapidly increases to near perfect values (> 0.9) the early phase of training, showing that the model tends to repeat the first token identically at most positions, even though the output token itself might be incorrect. Hence, repetitive sequences appear be inherently easier for the Transformer to learn, and this likely the reason for repetition bias in the early phase of training.
6 Repetition Bias and Representation Collapse in LLMs
Having established that the degenerate patterns of repetition bias and representation collapse are prevalent in small Transformers trained on algorithmic tasks, crucial question is whether these phenomena also happen beyond toy settings, specifically during the early pre-training stages of LLMs. We verify that this is indeed the case, using early-stage checkpoints of open-source LLMs Pythia 66]
and OLMo-2 (37].
For Pythia models with 14M, IB, 1.4B, and 2.8B parameters, we find strong representation collapse in the early training steps in their last layer and repetition bias in the output sequence (Figure [8) Specifically. we randomly sample 100 questions from the test split of the A12 ARC-Easy dataset [12) For each question, we let the model generate 8 tokens, and compute the pairwise cosine similarity of the hidden states (see Appendix for more implementation details). Figure shows that at initialization, the average cosine similarity relatively low (0.4-0.65). but within a few steps of training for all models, it sharply increases to > 0.9. These results remain similar if we use random sampling instead of greedy decoding, and other datasets like GSMSK and ARC-Challenge (Appendix D Further, the outputs for many prompts in the greedy decoding case are trivial repetitions of the same token, c.g., newline "(n', a clear manifestation of repetition bias.
Similar representation collapse patterns as Pythia are observed for the OLMo-2 2 7B model. For its carliest available training checkpoint (step 150. OLMo 1B)4-7B), the average representation cosine similarity in the setup from Section 6 is R 0.93; for the next checkpoint at step 600, this value has already decreased to R 0.43 (similar for both greedy decoding and random sampling strategies). Hence, repetition bias and representation collapse occur in the early pre-training phase of LLMs, validating our findings beyond toy settings.
Pythia-14m
Pythia=
Pythia-1.40
Pythia-2.
1.0
1.0
Cosine Semiaray
0.0
3.5
0.0
0.8
066
23
0.4 3.6-
0.6
2.5
3.4
0.2 n.4
0.4
100
"10""
10+
102
100
201
10:
"10F
110
104
100
Pretraining
Pretraining
Pretraining
Figure 8: Representation Collapse, Repetitions during Pythia Pretraining. Representation collapse at different Pythia pretraining checkpoints evaluated on the ARC-Easy test dataset; tokens generated via greedy decoding. We find that similar to our toy experiments, average pairwise cosine similarity between hidden states and repetition frequency start at relatively low values, and increase to near 1.0 during the early phase of pretraining, followed by decrease around step 1000.

---

7 Related Work
Abrupt learning has been studied in various settings: study sudden drop in loss when training BERT, and show abrupt learning of the attention map (termed Syntactic Attention Structure / SAS) concurrent with the sudden drop in training loss. (20) study sudden drop in loss ('Eureka moments') when training Transformers on multi-step tasks. They show that ill-distributed softmax attention scores lead to small gradients for key, query weights causing slow learning of attention map, and that appropriately modifying softmax temperature can alleviate optimization issues. [3] studied abrupt learning in various neural net architectures trained on sparse parity tasks, demonstrating underlying hidden progress during training via measures such as Fourier gap and weight movement. [8] study abrupt learning when training a BERT model on matrix complction. while [32] analyze abrupt learning for grammar data setup through graph percolation. For abrupt learning in in-context learning [17]. there has been a line of recent works [10 15
47 54, 55] that proposed various theoretical and empirical explanations. [15] studied abrupt learning for Markov chain in-context learning task with Transformers, and demonstrate a partial solution in their setup in the form of a unigram estimator. [46] analyze partial solutions for in-context n-gram task with Transformers, corresponding to k-gram estimators (k- < n). For diagonal-attention Transformers with small initialization, ( analyze during training.
dynamics for theoretical understanding of loss plateaus
A line of recent work has focused on understanding 'grokking' (39), i.c. abrupt generalization after an extended phase of memorization of training data by the model. Subsequent works have representation learning (3T), delay in feature learning [27 33 52 40) further study 'naive loss minimization' and numerical stability issues for understanding grokking. Note that grokking is fixed-datase phenomenon concerning memorization of the training dataset, and hence somewhat distinct from abrupt learning studied in this paper.
studied grokking through circuit formation 34
36 45
Rank collapse is related phenomenon for deep softmax transformers at initialization that might hinder training (2): however, our representation collapse phenomenon is different in that (i) we use shallow (1 or 2 Tayers) Transformers instead of deep ones; (ii) we use linear attention instead of softmax; (iii) our observed representation collapse occurs only after a few steps of training, not at initialization. (4) show form of representation collapse that for 2 sequences (U1, U2,..., vn) and (v), U....., Un, Un), as n grows large, the pretrained model's hidden state representation for the last token becomes identical for both sequences (Theorem 4.2, (4). Note that study evolution of representation collapse in the early phase of Transformer training, distinct from the notion of representation collapse at initialization, or in final pretrained models.
Repetition language model outputs
well studied problem [16
19
30 48 51,
53)
However most works focus not on the early phase of training, but on how repetition may arise in trained language models, and how to mitigate them. Towards understanding learning dynamics, [8]
IT] report token repetitions in the early phase of training language models. Please see Appendix for additional discussion on related work.
Discussion
We identified repetition bias and representation collapse as key characteristics of the early-phase implicit biases of Transformer training, which are closely connected to the commonly observed loss plateau. We further discussed 'search' over sequence tokens as a possible cause for loss plateau. The question of why behaviors such as representation collapse exist during early time training is an important question for future work. Furthermore, while we hypothesize a search-like phenomenon to be the reason behind slow learning of attention map, a rigorous theoretical validation of this hypothesis, and how possibly connects to the intuitive "complexity" of the task, are important questions for further research.
This work was supported by DARPA. The authors thank Misha Belkin, Ekdeep Singh Lubana, Yongyi Yang, Zhiwei Xu, Vaibhav Balloli, and anonymous reviewers for helpful comments at various stages of this work. Part of this work was done when the authors were visiting the Modern Paradigms in Generalization Program at the Simons Institute for the Theory of Computing.
Acknowledgments
10

---

References
[1] Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. SGD learning on neural net-
works; leap complexity and saddle-to-sadeee dynamics. 2023. arXiv: 2302. 11055 [Cs.LG]
URL: hwt.:httss//aaae
org/orgga/aoo 22022 11055
[2] Sotiris Anagnostidis et al. "Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse". In: Advances in Neural Information Processing Systems. Ed. by Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho. 2022. URL:
https://openreviart net/ wtw.ooruumnddPP/111111-101
[3] Boaz Barak et al. Hidden Progress in Deep Leurning: SGD Learns Parities Near the Compu- tational Limit. 2023. arXiv: 2207 08799 [cs LG]
URL:
08799
[4]
Federico Barbero et al. Transformers need glasses! Information over-squassinnngggggggggg. in language tasks. 2024. arxiv: 2406. 04267 [c. URL: hal::httts://arree htttta/://aziii.rr/abbaaa/////// wwwpg//aa///40.../ 042667
[5] Nora Belrose et al. Neural Networks Learn Statistics of Increasing Complexity. 2024, arXiv:
2402. 04362 [cs. LG] URL:: http://ttss:///rrrttee org/ abs /b0/ 2ebpb///40000322222
[6] Stella Biderman al. "Pythia: A suite for analyzing large language models across training and scaling". In: International Conference on Machine Learning. PMLR. 2023, pp. 2397-2430.
[7] Enric Boix- Adsera et al. Transformers learn through gradual rank increase. 2023. arXiv:
2306. 07042 [cs .G]] URL: httl::htpss/// //tpx//aavvvvv// http//bbb/33.... 07042
[8] Tyler A. Chang, Zhuowen Tu, and Benjamin K. Bergen. Characterizing Learning Curves During Language Model Pre- Tre-Trainnng: Learning, Forgetting, and Stability. 2024. arXiv:
2308. 15419 [cs. CL] URL: http:::/t////t arxiv org aataaa//.../...ee
15419
[9] Angelica Chen et al. Sudden Drops in the Loss: Syntax Acquisition, Phase Transitions, and Simplicity Bias in MLMs. 2025, arXiv: 2309. 07311 [ca. CL]
URL: https://arxii: org/
(bbs/230900730
[10] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training Dynamics of Multi- Head Softmax Attention for In- Context Learning: Emergence, Convergence, and Optimality.
2024. arXiv: 2402. 19442 [cs LG] URL: htttps://axxii wwte/orggggoee ht///bb//2200 0aabs//440..1900 19442
[11] Leshem Choshen, Guy Hacohen, Daphna Weinshall, and Omri Ahend. The Grammar-Leane
Trajectories of Neural Language Models. 2022. arXiv: 2109.06096
(cs.CL)
URL: https
//arxiv. org/ abs abs/ /2109. .00096
[12] Peter Clark al. "Think you have Solved Question Answering? Try ARC, the A12 Reasoning Challenge". In: arXiv:1803.054533 8033055 (2018).
[13] Karl Cobbe et al. Training Verifiers to Solve Math Word Problems. 2021. arXiv: 2110. 14168
[cS.LG] URL: httpt: wwtp://...r/bbbbbbbbbbbbbb ab1/1110 14168.
LG]
[14] Hugo Cui, Freya Behrens, Florent Krzakala, and Lenka Zdeborová. A phase transition between positional and semantic learning a solvable model of dot-product attention. 2024. arXiv:
03902 [cs LG] URL: httpt:ttta/saaass
org/ 2lbs/2402.00 03902
[15] Ezra Edelman et al. "The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains". In: The Thirty-eighth Annual Conference on Neural Information Processing net/forum? d=qaRT6QTIqq]
Systems. 2024. URL: https::/ http:::s///p///vvviiieee
[16] Zihao Fu, Wai Lam, Anthony Man-Cho So, and Bei Shi. Theoretical Analysis of the Repetition Problem in Text Generation. 2021. arXiv: 2012. 14660 [cs. CL] URL: https
//arxiv. org/ 2012 14660
[17) Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What Can Transformers Learn On-Context? A Case Study of Simple Function Classes. 2023. arXiv: 2208 01066
[cs. CL] URl:: wtt.ttttss://aaret http://wxvvvrg../e//// abs / 2208. 01066,
[18] Pulkit Gopalani, Ekdeep Singh Lubana, and Wei Hu. "Abrupt Learning in Transformers: A Case Study Matrix Completion", In: The Thirty-eighth Annual Conference Neural Information Processing Systems. 2024. URL: https openreview net / forum ? id= D9RZAEP341
[19] Tatsuya Hiraoka and Kentaro Inui. Repetition Neurons: How Do Language Models Produce Repetitions? 2025. arXiv: 2410. 13497 [cs. CL] URL: https://arxiv, arw.v. org/ abs/ 2410
13497

---

[20] David T. Hoffmann et al. Eureka-Moment in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems. 2024, arXiv: 2310 12956 [cs. LG]
URL: https:://aaaa
org/ abs/2310 12956
[21] Ari Holtzman et al. The Curious Case of Neural Text Degeneration. 2020. arXiv: 1904.09751
[cs. CL] URL: pettss//arriize org/ abs/ hta.s/999...997 09751
[22] Jesse Hoogland et al. Loss Landscape Degeneracy Drives Stagewise Development in Transform-
ers. 2025. arXiv: 2402. 02364 [c3 .G]] U'RL: https: http://xxvvv.o//i http://wwwrr//b//////e//// /t////////2///// 2202/2402..22 02364
[23] M. Emrullah Ildiz et al. From Self-Attention to Markov Models: Unveiling the Dynamics of
Generative Transformers. 2024. arXiv: 2402. 13512
[G]
URL: httts:///rxiv org/
abs/2402. 13512
[24] Masato Inoue, Hyeyoung Park, and Masato Okada. "On-Line Learning Theory of Soft Commit- tee Machines with Correlated Hidden Units -Steepest Gradient Descent and Natural Gradient Descent-". In: Journal of the Physical Society of Japan 72.4 (Apr. 2003), pp. 805-810. ISSN;
1347-4073, DOI: 10. 10.11143jja8 jpsj 72.805
URL: http: //xx doi . org w0w.//..113////...... 114.:/143/JJJJ...
805
[25]
Keller Jordan et al. Muon: An optimizer for hidden layers in neural networks. 2024. URL:
https://kellerjorrt github. www.gi.i..o.st...taaeeeoooo ipppss/mmuooi
[26]
Andrej Karpathy. KarpathyminnGPP: A minimal pytorch re-implementtation of the openai GPT (generative pretrained transformer) training. 2022. URL.:
MinGPT
[27]
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the Transition from Lazy to Rich Training Dynamics. 2024, arXiv: 2310.06110 [stat.ML. URL:
https://arxiis org/ abs 2310. 231./060100000000001001111111100
[28]
Nayoung Lee et al. Teaching Arithmetic to Small Transformers. 2023, arXiv: 2307 03381
[G] URL: https: httppp://aaaiv.orr//bbbbbb
org/ absp//b///2/3000 03381
[29]
Alexander C. Li et al. On the Surprising Effectiveness of Attention Transfer for Vision Trans-
formers. 2024. arXiv: 2411 09702 [cs. LG]
URL: https: arxiv. org/ abs 2411
09702
[30]
Huayang Li et al. "Repetition In Repetition Out: Towards Understanding Neural Text De- generation from the Data Perspective". In: seirty- Conference on Neural Information Processing Systems. 2023. URL: ptcttp:://openreviee wwwpiarvviiieeeeeeeeeee
[31]
Ziming Liu et al. Towards Understanding Grokking: An Effective Theory of Representation Learning. 2022. arXiv: 2205. 10343 [cs.
URL: https: //axxiva.vvv////////// org/ abs /aba//25 2205
hbgg/bb///
10343
[32]
Ekdeep Singh Lubana, Kyogo Kawaguchi, Robert P. Dick, and Hidenori Tanaka. A Percolation Model of Emergence: Analyzing Transformers Trained on a Formal Language. 2024. Xiviv:
2408. 12578 [ca LG] URL: https:: arxiv. org/ 2abb/2408 12578
[33]
Kaifeng Lyu et al. Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking. 2024. arXiv: 2311 18817 [C]
URL: https: //axxiv....//////////// org/ abs/ 2311
arg/abs/
18817
[34] William Merrill, Nikolaos Tsilivis, and Aman Shukla. A Tale of Two Circuits: Grokking as Competition of Sparse and Dense Subnerworks. 2023. arXiv: 2303. 11873 [cs.LG] URL:
https://arxiv org/ abs////b//23.3.. 11873
[35]
MLPClassifier scikit-learn. URL: https::/sctii http:::://ciitteieee learn. otgg/ttabee/oouules... aonerated/sklearm neural nww.wrk. _ewral_network..eeee wtw.:ott........lssssssi........ html
[36]
Neel Nanda et al. "Progress measures for grokking via mechanistic interpretabilityy. In The Eleventh International Conference on Learning Representations. 2023. URL: https
//openreview
[37]
Team OLMo et al. OLMo 2 Furious. 2024. arXiv: 2501 00656 [cs. CL] URL: https
/arxxiv.orge abs/ 11.5010005555
[38]
Abhishek Panigrahi et al. Progressive distillation induces an implicit curriculum. 2024. arXiv:
2410. 05464 [cs. LG]] URL: htt.:ttttt:///ariiv org/ abs/2410...... 054664
[39] Alethea Power et al. "Grokking: Generalization beyond overfitting on small algorithmic
datasets". In: arXiv preprint 2011
12

---

[40] Lucas Prieto, Melih Barsbey. Pedro A. M. Mediano, and Tolga Birdal. Grokking at the Edge of Numerical Stability. 2025, arXiv: 15501.04699 [cs.LG]
URL: httpps://axxxi..orttttt abs/
2501. 046997
[41]
Gautam Reddy. The mechanistic basis of data dependence and abrupt in in-context classification task. 2023. arXiv: 2312. 03002 [cs.LG]
URL: /htps:://aaaat /ttp///axxv..v/. org/ abs/
2312. 03002
[42]
Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. distributional simplicity bias in the learning dynamics of transformers, 2025, arXiv: 2410. 19637 [cs.CL]
URL: https::/// / arrivv..vv.///////////ar///i/iv/ org/ abs 2410. 19637
[43]
David Saad and Sara A. Solla. "On-line learning in soft committee machines", In: Phys. Rev. E 52 (4 Oct. 1995), pp. 4225-4243. DO1: 10 1103/ / PhysRevE. 52 4225 URL: https
//link. aps otg//oi/o/// /o///oi////
S..0//hysseee...... 52. 4225
[44]
Aaditya K. Singh et al. What needs to go right for an induction head? A mechanistic study of in-comtext learning circuits and their formation. 2024, arXiv: 2404. 2400.07129
[G] URL:
https://arxi>> org/ abs/ 2404. 07129
[45]
Vikrant Varma et al. Explaining grokking through circuit efficiency. 2023. arXiv: 2309. 02390
[cs [cS.LG] URL: httpss//tattti......// org hwtp:/orgggabb///aeeeees 02390
[46] Aditya Varre, Gizem Yüce, and Nicolas Flammarion. "Leeaming In-context SnS-grams with Transformers: Sub-SnS-grams Are Near-Stationary Points", In: Forty-second International Conference on Machine Learning. 2025. URL: https: htt:://pppereiiieeee net met/foruu fort/ffoum?e0 idd=
DMLdvGDeHL.
[47] Mingze Wang, Ruoxi Yu, Weinan E, and Lei Wu. How Transformers Get Rich: Approximatiin and Dynamics Analysis. 2025. arXiv: 2410. 11474 [cs.LG] URL: https://arxiv org/
abs abs/2410, 11474
[48] Weichuan Wang et al. Mitigating the Language Mismatch and Repetition Issues in LLM- based Machine Translation via Model Editing. 2024. arXiv: 2410 07054 [cs.CL) CL] URL:
https:// //t///axvvv.// wrrii.orgg abs/2410 07054
[49] Jason Wei et al. "Emergent Abilities of Large Language Models". In: Transactions on Ma- chine Learning Research (2022). Survey Certification. ISSN: 2835-8856.
URL: https:
openreview.oo wett:foorunndd-yyak55uua..
[50] Thomas Wolf al. HuggingFacess 's Transformers: State-of-the-art Natural Language Process-
ing. 2020. arXiv: 1910 03771 [cs. CL]
URL: |tt:::t:::///rria http://///xxii.oo/aabbba http:///og/bbbss////// a9teaa///11...1 1910 03771
[51] Jin Xu et al. Learning to Break the Loop: Analyzing and Mitigating Repetitions for Neural Text
Generation, 2022. arXiv: 2206 02369
URL: hittp:://aaae
org/ abs 2206
02369
[52]
Zhiwei Xu et al. Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Data. 2023, arxiv: 2310. 02544 [cs LG] URL: pottps://ost htttps///axiiv arxiv. www./bs////eoeooooooooooeee 2310 02541
[53]
Junchi Yao et al. Understanding the Repeat Curse in Large Language Models from a Feature
Perspective. 2025. arXiv: 2504. 14218
[Cs
URL.: https: aaxxvv.vvvv. org/ abs 2504
14218
[54]
Yedi Zhang, Aaditya K. Singh, Peter E. Latham, and Andrew Saxe. Training Dynamics of In-Context Learning in Linear Attention. 2025, arXiv: 2501 16265
URL: https
farxiv.org abs///as...1 16265
[55] Nicolas Zucchet, Francesco d' Angelo, Andrew K. Lampinen, and Stephanic C. Y. Chan. The emergence of sparse attention: impact of data distribution and benefits of repetition. 2025.
arXiv: 2505. 17863 [cs. [C] URL: https:////// pipps///arxivv .olvv.org///o. abs/ 17863
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All claims made in abstract and introduction are supported by evidence in subsequent sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed limitations of this work in Section g
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
14

---

Justification:
There are no theoretical results in this work.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proof' can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We have described the model and data generation process in detail for reproducibility in Section [ and Appendix We have also released code at
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
15

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We have released code at githhb..oom/lllliieoppaaiiflloas also
specified in Section 2
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
Cotee//ddeeeebmmmcia ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Described in Section 2 Appendix A and at other relevant places wherever applicable.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We have not included error bars due to computational constraints.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We have described compute resources used in this work in Appendix A
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The research in this paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss societal impact in Appendix A
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI...
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our data and models specifically for toy setups designed to understand training dynamics of Transformers, and hence there no risk of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have used existing GPT implementation (6).. Pythia models and OLMo- model (37) via HuggingFace Transformers (50), ARC- Caall datasets [2] and
GSM8K dataset (13) for which we have added relevant citations wherever necessary. Licence
information for all these is given in Appendix
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, haperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
18

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We have released code at gitthh.com//lkigggggaaaittttttos details provided in README.md.
with
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
There is no research involving human subjects in this paper.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: There is research involving human subjects in this paper.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
19

---

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non standaad components.
Please refer to our LLM policy for what should or should not be described.
20

---

A Implementation Details
Compute Resources. All experiments were conducted single GPU (NVIDIA A100 L40S) on an academic computing cluster. Most training runs in this paper complete within a few hours.
Causal Linear Attention. Linear attention transformer is obtained simply by removing the softmax activation function when computing the attention map, and setting the causal mask to 0 instead of -oo. We use the existing minGPT implementation [26] (MIT licence) for our experiments, modifying the code as above and wherever required.
LLM Experiments. We use Pythia [6] OLMo-2 37 pretrained models (Apache 2.0 Licence) hosted on Huggingface Transformers (50) and evaluate them on the ARC-Easy Challenge datasets 12] (CC-BY-SA 4.0 Licence), and GSM8K [13 (MIT Licence). We set the use caee_ccch=Faaseee the generate function, and use the hidden state used for predicting each of the 8 output to- kens. For random sampling, we use sassmmll-eTuu (using default temperature value), using do do_sample=Fals for our greedy decoding results.
B Results for Other Algorithmic Tasks
This section presents results on a suite algorithmic tasks, verifying generality of our identified phenomena.
Table I: Algorithmic tasks that show abrupt learning and partial solution during plateau
Task
Description
Partial Solution
Moving Window Sum (MWS)
Sum over moving window of elements, copy 1st element
First input element
Prefix Sum (PRE)
Compute prefix sum of a given n-length sequence
First input element
Permutation (PER)
Permute an n- sequence by given permutation
Incorrect permutation of input sequence
Multi-Digit Addition (ADD)
Add atmost-n-digi numbers
First digit (0 or 1) i.e. total carry-over from digits
Histogram (HIST)
Compute counts of each element in n-llngth sequence
100% Repetitive sequences Repetitive sequences
Reverse (REV)
Reverse n-length input sequence
Copy (COPY)
Copy n-length input sequence
Repetitive sequences
Copy MWS MWS;
Concatenation sequences from COPY, MWS, and MWS; (moving window of elements)
During 1st plateau: COPY part correct; 2nd
plateau: COPY and MWS parts correct (Appendix B.7)
The loss plateau is very brief, hence a partial solution like other cases is not applicable.)
In the following table we describe APM sets J for different algorithmic tasks used this work. We assume the attention map of the shape A E (14-1)x(L-1) in the next token prediction setup on the full input sequence. Hence, for MWS, Prefix sum, Histogram, Reverse, Copy, 7. = 33, for permutation = 50, and for Multi-digit addition I, 15. The APM sets are denoted by sets of tuples (i.j) LL- 1] x indicating the row/column indices in the attention map A.
21

---

Table Attention Progress Measure (APM) sets (52) for all tasks
Task
APM set S2
MWS
((16,)))(( (1000))000000+----)))44008000000 [ [1,
Prefix Sum
[5]
Multi-digit
{(i. 12 - |e -) ie 99. 12
Addition
13 - ( (
[3]
Permutation
Layer 1: (17+4mm-11)]i 15]}
Layer 2: (33+1,17+)
15]}
Histogram
Layer 1: ((16+1,4) | i [0)
Copy
{(16 + i,d) | 15
Reverse
{(16 +5, i [0, 15]}
B.1 Multi-Digit Addition
This task involves adding 2 atmost 4 4-digit numbers; if the numbers are represented as a d10203044 = b1b26zb4 and their sum a = e26r44 then the training sequences for ADD are of the form
01 03 +. b1, 62, b3, by. =, c4. c3, C1.
Note that the output sequence is reversed, following the observations from 28 We find similar abrupt learning characteristics (Figure()) partial solution in this case being Co L.e total carry- from single digit add operations. An interpretable attention map learnt for the output sequence shown in Figure TO)
Train/Test Loss. Accuracy
Cosine Similarity & Repetitions
Attention Progress Measure
1.0
Resetit on Freg
Codne Simiadtty
2.2
0.8
All Progress
2.0
33 3.0
Loes Acorracy
Fortill Sol. ACC.
3.6
2.4
022
3.2
5.2
0.0
Steps
Figure 9: Training dynamics for Add task. (left) Train/Test Loss, Accuracy and Partial solution progress (Co accuracy); (middle) Repetition frequency and representation collapse; (right) Attention progress measure.
Figure 10: Attention map for add task, note that the model attends to the relevant digits in the input numbers, and to somewhat lesser extent to the preceding digits as well (highlighted positions show entries with larger magnitude).
22

---

B.2
Prefix sum
This task involves computing the cumulative (prefix) sum of an n R-length sequence integers, so that the training sequences in PRE are of the form (n 16, SEP = 17),
T1, D2,
Iny SEP, y1, U2, Un
3/
Σ zi
mod 17 V [n]
Training dynamics for this task shown in Fig. which show similar abrupt learning behavior as MWS and partial solution learning for y1- The interpretable attention map learnt for this task is shown in Figure 12
Triin/Tsst Loss & Accuracy
Attention Progress Cosine Simi larilarity
Sequence Entropy
1.0
0.0
2.75
0.7
2.50
Attention Progress
Seo. Entroy
2.00
2.25
Portial Sol. Acc
Acouracy
Accuraay
Actertion Prng
5m
Curput Entray
Data Entraay
0.2
1.50
1.2
300
180
200
300
450
5.0
300
4.00
500
Steps
Steps
Figure 11: Training dynamics for Prefix sum task. (left) Train/Test Loss, Accuracy and Partial solution progress (y1 accuracy); (middle) Attention progress and representation collapse; (right) SeqEnt for data and model output sequences.
Figure 12: Attention map for Prefix sum task, that uses the relevant token in the input, as well as the previous token in the output to track prefix sum (highlighted positions show entries with larger magnitude).
B.3 Permutation
This task involves training a 2-layer, 1-head Transformer on permuting a length sequence using the permutation n, which generated random and distinct for each training sequence. Formally, for a sequence of positive integers (II,..., .n) and a permutation (Tı,..., nn) over [n)], training sequences for PER, k 1, 2, are given by
21. ................................ J.,. SEP, M...... NT. SEP, I "
where Unif [7,, 18, 32). 16, SEP O. The partial solution in this case is the output sequence being an permutation of the input sequence I, i.c., learns copy the tokens correctly, but in wrong order (Figure 13). The interpretable attention maps in this case show that the model learns copy the correct tokens based on the permutation provided (Figure 14a) and then uses them for the final output sequence (Figure 14b).
23

---

TrairyTest
Cosine Similarity Repetition
Attention
1.0
Cosine Stn,, RRepetetion Reo
Cheine simiarky
tepettion
6:25
0.8
Alvesiion Smogeee
6.22
:30 Accurdyy
Stepss
Stmps
Figure 13: Training dynamics for Permutation task. (left) Train/Test Loss, Accuracy and Partial solution progress; (middle) Repetition frequency and representation collapse; (right) Attention progress measure. Note that the repetition frequency decreases by step 100, which followed by the partial solution.
(a) Attention map Layer where rows are attention weights over the input part
(b) Attention map Layer the rows (output part of the sequence) are attention scores over the part of
T1, ,,, In of the sequence. The highlighted positions attending to T2,,2...., 15,4, 1,,3, 6,, 10, 11, 9, 16, 12,. index [n].
sequence to which Layer attention map copies the correctly permuted tokens. This implies that this at- tention map simply copies the correct token from the residual stream after Layer
for
Figure 14: Attention maps for the 2 layer Transformer used for Permutation task; highlighted positions show entries with larger magnitude.
B.4 Histogram
This task [4] involves computing the counts of elements in the input sequence, and training sequences of the form
31 I2, Fa, SEP, 11.92,......
11 Σ
1 xx,]
where U1if(1,2..... n
n 16, SEP O. We train 2-layer, 1-head transformer for this task, with gradient clipping (1.0) to avoid loss spikes (Figure 15). We note that the repetition bias in this case is quite strong which leads s 100% repetitions in the early phase of training, and which we characterize as partial solution for this task. Further we only consider the attention map from layer 1 (Figure T6) since this is the most consistent and clearly interpretable across runs, and indicates an ceentit-mapp-lae function.
24

---

Similarity Repetitions
7.0
2.50
0.5
Proqrass
2.25
0.8
Cossmee Sim, Repetiton
0.3
0.7
Arcunics
1.75
Accurscy
0.0
Similatty
0.
Layer
1.30
9.4
0.1
255
1200
750
10.00
12.10
15.00
0
250
500
750
1903
1253
1503
255
500
750
1000
1250
1500
Stops
Steps
Stops
Figure 15: Training dynamics for Histogram task. (left) Train/Test Loss and Accuracy: (middle) Repetition frequency and representation collapse: (right) Attention progress measure. We only measure attention progress for the 1st layer, since that the one that consistently and clearly shows an interpretable pattern (Figure 16).
Figure 16: Attention map in layer 1 for histogram task, where rows for the latter half of the sequence compute attention weights over the input tokens 24, similar to an identity map (highlighted positions show entries with larger magnitude).
B,5 Reverse
This is the task of reversing the input sequence, so that the training sequences for reverse task REV are given as,
21, 22.......... Tn, SEP, Tny Fn-1n....,
for 1 16), 16, SEP 0. The training dynamics are shown in Figure 17a and the interpretable attention map is shown in Figure 17b
Train//estt Loss. Accuracy
Repetition, Cosine Sim, and Attention Prog.
1.5
Ropotttin Fea./Co Corhe Sin
Accurry
4.30
Cedine Sindartt Abartie Frog. Maasue
0.00
(b) Attention map for re- verse task (highlighted po- sitions show entries with larger magnitude).
107
150
(a) Training dynamics for reverse task
Figure 17: Training dynamics for Reverse task. We see Abrupt Learning, Representation Collapse and Repetitions, though to lesser extent than MWS task. Note that the plateau is much shorter compared to MWS, possibly explained by the fact that reversing a sequence is 'easier" than computing the moving window sum.
B.6 Copy
This is the trivial task of copying the input sequence as is, so that the training sequences for copy task COPY are given as,
Ø1, D2 #2, ..., Ln, SEP, 01, ...
25

---

for ~ Uni 2.... 16), n 16, SEP = 0. The training dynamics are shown in Figure 18a and the interpretable attention map shown in Figure 18b
Accuracy
Cosine Sim,
25
5mh
26
24
Repettin1 Frng
8 14 Progress
Use Acouscyy
Cecre arty
Alurtion
16
14
(b) Attention map copy task (highlighted positions show entries with larger magnitude).
Staps
Ste90
(a) Training dynamics for copy task
Figure 18: Training dynamics for Copy task. Similar to reverse task, observe Abrupt Learning, Representation Collapse and Repetitions for Copy task, but this time to an even lesser extent than reverse task itself.
B.7 Concatenation of Copy, MWS, MWS;
We train the model on following sequences.
Il, T1, , SEP, y1, ,), .2, 9/13
i [1, 4]
= (x; + x(41) mod 17
i [,
(r++1 + r442 + 1443) mod 17 E [10, 13]
with x; [1, 2,....... 16), SEP 17, and find that the train loss follows a 'staircase' -laircase'-like evolu-
tion, learning the 3 distinct parts ([v1..... Val, [us...., yol. [810..... yıs)) sequentially (Figure 19p.
Similar sequential trend is seen for accuracy and average cosine similarity, measured separately over these parts (labeled Part 1,2,3).
Train/lest Loss and Accuracy
Part-wise Accuracy
Part-wise Cosine Similarity
0.8
Coline Similtarity
2.2
lass Accurity
2.6
Aecurrccy
0.6
Tain Acc
Tain Act.
0.6
Cosine 5im.
Coslne
0.4
Tain Acc
0.1
Cosine 5m.
0.2
0.2
1.4
0.0
520
1000
1900 Steps
7030
25.00
3000
530
1000
1500
2000
2500
3000
500
1200
1505
2000
2500
3000
Stops
Steps
Figure 19: Partial Solution for Concatenated Target Sequence.
C Probing Attention Head Output
We probe the output hidden states of the Attention block, before they are added to the residual stream (Equation (2). Specifically,
2-layer MI.P probe with hidden layer width 256 and ReLU activation function to map these hidden states to the ground truth output labels Formally, for hidden states h 256 the probe defined Probe(h) == We2 ph, Wp2 17x226 WP1 E 255 implemented using scikit-leam MLPCI aLPClassifier (35). We probe at every 10 training steps of the Transformer model, training the probe on 1024 samples and evaluating using 1024 held-out samples, both different from training data for the Transformer.
We find that the probe test accuracy increases earlier than the sudden jump in model accuracy (Figure 20,, demonstrating 'hidden' progress in Attention block outputs.
26

---

Train Loss & Attention Progress
Model and Probe Accuracy
1.0
0.7
0.6
Accuracy
22
Train Loss Attendion
APM
0.6
Model Trallyesst Acc.
2.0
0.4
Probe
18-
0.3
1.6
0.2
0.2
0.1
Steps
Steps
Figure 20: Probing Attention Outputs. We find that the probe accuracy when probing attention block outputs (Equation ( starts increase earlier than the model aceruacy during training, highlighting 'hidden' gradual progress in the hidden states in addition to attention map (via attention progress measure).
D Results for REPEAT,, REPEAT,
REPEAT
This task defined as,
Ø (x1 + 1) mod 17 9 i 16
15i58
yi
The training dynamics for REPEAT, are given in Figure 21a We find that similar to REPEAT,, the training loss does not exhibit any plateau. Moreover, the repetition frequency P and metric a1 increase rapidly to = 1.0 early on in training. We measure the average cosine similarity for hidden states for repetitive blocks of the output sequence (see Figure 21b
Accuracy
Repetition Frequency
26
Vung Mamg
0.E
22
Tranffiad
18
03
1.0
Castre Simlartty
Recerition Free (nt
1.4
Fiuit Tooen frog (a)
40
120
10c
(b) Pairwise hidden state cosine similarity for REPEAT, task step 20.
(a) Training dynamics REPEAT, task
Figure 21: REPEAT, training dynamics. Note that there no plateau in loss, similar to REPEAT, task. Further, the pairwise cosine similarity for hidden states take specific form indicating the blocks of repeated tokens in the output.
REPEATa This task defined as
I1
15i54
(w) + 1) mod 17
5gi58
(x) + 2) mod 17 9 0 12
(w) + 3) mod 17 13 < i 16
Similar results are observed for REPEAT, as well (Figure 22a for this case we measure the average cosine similarity hidden states for repetitive blocks of the output sequence (see Figure 22b
27

---

Trait/Test Accuracy
Cosine Similarity Reyetition
26
18
2.2
Than
Sm.
fooe thn Tooer (R)
1.0
1.4
(b) Pairwise hidden state cosine similarity for REPEAT, task step 20.
Steps
(a) Training dynamics for REPEAT, task
Figure 22: REPEAT, training dynamics. Similar to REPEAT,, there is plateau in loss. The pairwise cosine similarity for hidden states takes a form indicating the blocks of repeated tokens in the output.
E Additional LLM Results
Pythia-14m
Pythia-10 1b
Pythia-1-48
Pythia-2.80
1.0
1.0
110
110
Cosine Smilaty
0.0
0..
0.0
0.8
0.5
0.6
0.6
0.4
0.4
0.6
0.4
0.6
Ropiatttinn
0.2
0.2
0.6
0.4
107
107
15"
1."
10'
Pretraining
Pretraining
Pretraining
Pretraining Step
Figure 23: Representation Collapse, Repetitions during Pythia Pretraining (GSM-8K
13
Representation collapse at different Pythia pretraining checkpoints evaluated on the GSM-8K test dataset; tokens generated via greedy decoding.
Pythia-14m
Pythia-1b
Pythia-1 4b
Pythia-2.88 2.8b
1.0
1.
Hipal mi..
Casine Simiarky
0.0
0.8
0.8
0.6
0.4
D.1
0.6
0.4
0.6-
0.5
0.
2.4
0.2
0.4
10r
10
10
101
30P
10
30
307
10:
10"
10t
10r
10"
10
Pretraining Step
Pretraining Step
Pretraining Step
Pretraining Step
Figure 24: Representation Collapse, Repetitions during Pythia Pretraining (ARC-Challenge [1].. Representation collapse at different Pythia pretraining checkpoints evaluated on the ARC- Challenge test dataset; tokens generated via greedy decoding.
Pythia
Pythia
Pythia-1.44
Pythia
0.9
Cosine Smilsnit
0.9
0.8
0.8
0.0
0.7
0.7
0.6
0.6
0.5
0.5
0.5
0.4
0.4
100
100
102
102
10"
10P
102
102
102
10
10P
10t
10a
102
10a
0.3
10c
104
102
162
10"
Pretraining Step
Pretraining Step
Pretraining Step
Pretraining Step
Figure 25: Representation collapse during Pythia pretraining (inference with random sampling).
28

---

Varying Model Configurations
Loss
Accuracy
Casine Similarity
Repetition Frequency
1.0
1.0
2.75
2.50
Cosine Sinilarity
0.8
Repotition Prea.
0.6
0.6
2.00
0.0
1.75
0.2
0.2
0.2
1,25
0.0
0.0
200
Steps
Steps
Steps
Stops
Figure 26: Number of Layers (L). Abrupt learning, representation collapse in the last layer and repetitions for multi-layer models.
Loss
Accuracy
Cosine Similarity
Repetition Frequency
1.0
2.75
0.8
0.8
Cosine Similarity
0.8
Accuracy
0.6
2.25
Repetition
2.00
0.4-
0.0
1.5
0.2
0.2
0.2
1.25
0.0
100
200
300
400
100
200
300
400
100
200
300
400
100
200
300
400
Stopp
Stops
Stops
Stops
Figure 27: Number of attention heads (H). Representation collapse and repetition bias in 1-layer multi-attention head models (embedding dimension is fixed at d = 256 for all cases).
Loss
Accuracy
Cosine Similarity
Repetition Frequency
1.0
2.5
0.8
Smillarity
Repetition From.
d 256
128
Accuraiy
0.6
0.8
0.4
0.4
0.4
15
0
500
10000
1500
500
1000
1500
500
1000
1500
500
1000
1500
Steps
Steps
Steps
Steps
Figure 28: Embedding dimension (d). Abrupt learning with representation collapse and repetition bias in I-layer 1-head models with different embedding dimensions.
2-layer
laaeer
6-layer
Cosine Similanty
0.8
Casine Similannty
0.8
Similantt
0.8
Layer
0.6
Layer
0.6
Layer
Layer
0.6
Layer 2
Layer 3
0.4
Layer
Layer
0.2
Layer
0.0
0.0
0.0
150
200
50
150
2000
250
150
250
300
Steps
Steps
Steps
Figure 29: Extent of Representation collapse at various intermediate layers. Cosine similarity values showing the extent of representation collapse after each intermediate layer multi-layer models. Note that the representation collapse is not so severe in the early layers of multi-layer models, but the cosine similarity becomes close to 1.0 as we progress to the final layer.
29

---

Train/Test Loss & Accuracy
Attention Progress and Cosine Similarity
1.0
Repetiiio Freq.
0.8
0.8
0.7
2.4
0.6
Miconson Progres
2.2
Accuracy
Cosine Accurac)
Repelition freq.
Cosine
0.5
2.0
Partial Sol. Acc.
0.4
Altn, Prog. Messure
0.2
0.2
1.4
0.0
0.0
0.1
200
400
600
800
1000
1200
1400
0
200
400
6.00
800
1000
1200
1400
Steps
Figure 30: Softmax Attention. For completeness show that repetition bias and early-phase representation collapse are not limited to linear transformers but are observed softmax attention transformers as well. Note that the loss plateau is longer than that for linear attention.
G Varying Optimization Setups
Train/Test Loss & Accuracy
Attention Progress and Cosine Similarity
2.8
1.0
Repetition Frest
0.7
Progress
Accuracy
Costne
Acourady
0.4
Repetttion Freq.
Partial Sol. Acc.
Aztn. Prog. Measure
1.8
0.4
0.3
0.3
1.6
0.2
0.2
1.4
M
0.1
50
100
150
200
250
300
0
100
150
200
250
300
Steps
Figure 31: Using SGD for training. We show that abrupt learning is not limited to Adam optimizer, and occurs with SGD (1) = 0.1) as well. We chose this value of 7) since smaller values typically lead to much longer periods of little decrease in loss, without increase in accuracy.
Train/Test Loss, Accuracy
Attention Progress and Cosine Similarity
Cosine Sim. Repesiioo Freq.
0.8
Lass
0.6
Accuracy
0.4
Cosine Simlarity
Repetition Frequency
Progress
0.2
100
125
50
200
Steps
Steps
Figure 32: Using Muon for training, with Adam for embeddings, LM head and bias parameters. We use learning rate 0.02 for Muon, and 1e-4 for Adam. Note that while the plateau duration decreases significantly compared using only Adam (Figure 2a). plateau of duration steps still occurs.
30

---

Train/Test Loss, Accuracy
Attention Progress Cosine Similarity
2.8
Cosine Sim, Repelition Freq.
2.6
0.8
0.8
2.4
0.6
2.2
0.6
Accuracy
Repetition Frequency
Simlartty
2.0
Accuracy
Attention Progress Measure
1.84
0.4
1.6
0.2
1.4
0.0
50
75
100
125
150
200
25
50
100
150
175
200
Steps
Steps
Figure 33: Using Muon for training, with Adam W for embeddings, LM head and bias parame- ters. We use learning rate 0.02 for Muon, and 1e-4 for AdamW (weight decay= 0.01).
Cosine Similarity
Repatition Frequency
3.0
1.a
0.8
Rupictiinn Proc.
0.6
1.5
0.2
0.0
Stege
Figure 34: Training Dynamics with Learning Rate Warmup. The learning rate increases linearly from 0 to le-4 during training steps [0, t].
Loss
Cosine Similarity
Repetition Frequency
3.0
1.0
2.0.8-
Repettion Frez.
0.2
150
230
300
100
200
400
3.00
400
2000 Stops
5:30
4000
Steps
Stops
Sto.ps
Figure 35: Training Dynamics with Learning Rate Warmup + Cosine Decay, The learning rate increases linearly from 0 to 1c-4 during training steps [0, t), and decreases cosine function during steps [1. + 1, 4,00
Loss
Accuracy
Cosine Similarity
Regetition Frequency
1.0
2-
1.0
1.0
8.03
8.1
Repetttin Fren:
2-
0.4
0.2
0.2
0.0
0.0
Figure 36: Training Dynamics with different Initialization Scales. We find that smaller initializa- tion scale (a < 1) leads to a longer plateau, and cosine similarity being very close to 1.0 throughout the plateau. With 10, the representation collapse is not strong, reaching peak value of zs 0.6. We smoothen the repetition frequency curve with window of size 20 for presentation clarity, using np. convolve crepeat_freg np. ones (20) /20, mode= same')
31

---

Accuracy
Similarity
Repetition Fraquency
3.0
0.8
R 1ed
Rep-titon Free.
LO0N
0.8
0.4
1.5
0.0
153
10'
0.0
10e
Steps
Figure 37: Training Dynamics with different Learning Rates.
Accuracy
Similarity
Repetition Frequency
1.0
0.8
0.8
2.25
Cosine 5im
0.6
2.00
0.4
1.79
1.50
1.25
0.0
0.0
Steps
Figure 38: Training Dynamics with varying Batch Size.
Accuracy
Similarity
Repetition Frequency
1.0
2.75
Cosine amiarrty
0.0
2.25
0.6
2.00
0.4
1.79
1.50
1.25
Figure 39: Training Dynamics with Adam W, varying weight decay parameter (A).
Acouracy
Similarity
Repetition Frequency
1.0
0.6
Cosine Smiaaryy
0.8
2.25
6
2.00
1.50
1.25
Figure 40: Training Dynamics with &g regularization, varying regularization parameter (X).
32

---

H Additional Figures
Train Lcss Cosine Similarity
Residual Norm
Resioual Nor
0.2
Costre Sm 1a1
Oodine Sim
2.5-
Coi.n Simlarky
2.50
2.0
Losh
Narm
2.0
0.42
4/
100
200 Stops
2.0
302
402
0.1
230
3.0
Stops
Stops
Figure 41: Norm and representation collapse dynamics for (a) pre- and (b) post-attention residual streams for all positions i, except the first position.
Cosine Similarity
Taken+Passit onal Embed
Attm Output
Cosine Similarity
Post-Alt: Res. Stream
MLP/Pre
0.6
0.4
0.2
0.0
0
50
100
150
200 Steps
250
300
350
400
Figure 42: Cosine similarity at various points in residual stream for 1-layer, 1-head Transformer trained on MWS task.
Train/Test Loss
Cosine Similarity
Repetition Frequency
2.75
0.8
Similantyy
0.8
2.50
to = 75
0.6
Repotiton
2.00
1.75
m
0.2
0.2
1.50
1.25
0.0
0
100
200
300
400
100
200 Steps
300
400
100
200 Steps
300
400
Steps
Figure 43: Biasing attention map by c 10 at different ly during training.
I Additional Related Work
Techniques from statistical physics have also been used towards understanding initial loss plateaus in neural training [3): they work in a 2-layer teacher-student setup, where the second layer is fixed during training, and use order parameters to study training. They show that there is permutation symmetry in the weight vectors of the first layer during the early plateau stage, and exiting this state leads to drop in loss. Singular learning theory (22) has also been used to explain stagewise learning dynamics in Transformers; they estimate the 'Local Learning Coefficient" (LLC) during training to quantify degeneracy in the loss landscape, and consequently explain the learning dynamics. The interplay of simplicity bias and Transformer learning dynamics has also been studied recently in (5) show that Transformers progressively learn higher-order ("many-body interactions among tokens sequence. In (5) authors show that neural nets learn use lower-order moments of data early in training via interventions on test data, and show equivalence between n- gram statistics and embedding moments for discrete domains.
42
33


Reviews:
[{"review_id": "J7vMOxijHg", "reviewer": "Reviewer_u6ru", "summary": "This 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.", "questions": "1. 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.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "While the paper compellingly identifies important phenomena such as repetition bias, representation collapse, and attention as a potential bottleneck, it stops short of explaining why this bottleneck arises. I understand the authors acknowledge this as future work and provided some hypotheses in the rebuttal upon my request, which I appreciate. However, the lack of theoretical framing or concrete mechanistic explanation in the main paper limits the depth of the current contribution. The work would be stronger with a simple toy model or analysis to support these observations. That said, I find the empirical findings interesting and valuable, and I am raising my score to a weak acceptance.", "rebuttal": "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. \n\n\n*Weaknesses*:\n\n1. **Hypotheses for observed phenomena**\n\n*Slow learning of attention map*: \n\nOur 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. \n\n[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.  \n\nWe 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.\n\n*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). \n \n\n2. **Robustness to hyperparameters**: We show that abrupt learning occurs across a range of optimization hyperparameter choices as below. \n\nTo 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.\n\n- *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 \n\na) AdamW (lr=1e-4, wd={1e-3, 1e-2, 5e-2, 1e-1}) and \n\nb) Muon (lr=0.02) with {Adam (lr=1e-4), AdamW (lr=1e-4, wd=1e-2)} for embedding, LM head and all bias parameters. \n\nand 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. \n\n- *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. \n\nAdditionally, 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). \n\n- *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)\n\n- *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.\n\n- *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. \n\n- *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.\n\n| Sweep         \t\t | Value            \t\t   | Abrupt Learning and RC | \n|------------------------|-----------------------------|------------------------|\n| Optimizer    \t\t\t | AdamW \t\t\t \t\t   | Yes, similar \t\t\t\t\t|\n|         \t   \t\t\t | Muon + Adam     \t \t\t   | Yes, shorter plateau \t|\t\n|              \t\t\t | Muon + AdamW     \t\t   | Yes, shorter plateau \t|\n| Learning rate schedule | Linear Warmup \t\t\t   | Yes, similar \t\t\t\t\t|\n|  \t\t\t\t\t\t | Linear Warmup + Cosine Decay| Yes, similar \t\t\t\t\t|\n|  \t\t\t\t\t\t | [Gu 25] (suggested by reviewer, adapted to our task)       \t\t\t   | Yes, similar \t\t\t\t\t|\n| L2 regularization \t | 1e-3, 1e-2 \t\t\t\t   | Yes, similar           \t\t\t|\n|  \t\t\t\t\t\t | 5e-2       \t\t\t\t   | Yes, slightly longer plateau, cosine similarity peak ~0.8 |\n| Learning Rate\t\t \t | 5e-7, 1e-6, 1e-5 \t\t   | Yes, longer plateau  \t\t\t|\n|  \t\t\t\t\t\t | 1e-3       \t\t\t\t   | Yes, similar           \t\t\t|\n| Batch Size\t\t \t | 8, 32 \t\t\t\t\t   | Yes, longer plateau  \t|\n|  \t\t\t\t\t\t | 2048, 4096       \t\t   | Yes, similar           \t\t\t|\n| Initialization Scale\t | 1e-3, 1e-2, 1e-1\t\t\t   | Yes, longer plateau, stronger RC |\n|  \t\t\t\t\t\t | 1e1       \t\t   \t\t   | Yes, longer plateau, weaker RC (peak cosine sim ~0.6) |\n\n(RC= Representation collapse)\n\n*Questions*:\n\n1. Please see response to Weakness 1 above.\n\n2. \n\n- Hyperparameter / optimizer choice: please see response to Weakness 2. \n\n\n- 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.\n\nWe 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.  \n\n\n[Abbe 23]: Abbe et al., 2023. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. COLT 2023.\n\n[Chen 24]: Chen et al., 2024. Sudden Drops in the Loss: Syntax Acquisition, Phase Transitions, and Simplicity Bias in MLMs. ICLR 2024.\n\n[Gu 25] Gu et al., 2025. When Attention Sink Emerges in Language Models: An Empirical View. ICLR 2025.\n\n[Loshchilov 19]: Loshchilov et al. 2019. Decoupled Weight Decay Regularization. ICLR 2019"}, {"review_id": "W8RXuvU78R", "reviewer": "Reviewer_ARVG", "summary": "The 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.", "questions": "**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.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "I maintain my positive score.", "rebuttal": "We thank the reviewer for their positive evaluation of our work! We address their concerns (about partial solution, attention progress measure, etc.) below.\n\n*Weaknesses*:\n\n1. 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.\n\n2. “...*APM metric is not readily available in most tasks*…”: We provide the optimal $\\Omega$ for all tasks considered in our paper below.   \n\n    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$. \n\n\n\n| Algorithmic Task\t | APM set $(\\Omega \\subset [L-1] \\times [L-1])$ | \n|--------------------|----------|\n| MWS    \t\t | $\\\\{(16,0)\\\\} \\cup \\\\{(16+i, i-1) \\mid i\\in[1,15]\\\\} \\cup \\\\{(16+i, i) \\mid i\\in[1,15]\\\\}$|\n| Prefix Sum    \t\t | $\\\\{(16,0)\\\\} \\cup \\\\{(16+i, 16+i) \\mid i\\in[1,15]\\\\} \\cup \\\\{(16+i, i) \\mid i\\in[1,15]\\\\}$|\n| Multi-digit Addition   \t\t | $\\\\{(i, 12-i) \\mid i\\in[9,12]\\\\} \\cup \\\\{(i, 17-i) \\mid i\\in[9,12]\\\\} \\cup \\\\{(i, 13-i) \\mid i\\in[10,13]\\\\} \\cup \\\\{(i, 18-i) \\mid i\\in[10,13]\\\\}$|\n| Permutation | Layer 1: $\\\\{(17+i, \\pi_{i+1}-1) \\mid i\\in[0,15]\\\\}$|\n|| Layer 2: $\\\\{(33+i, 17+i) \\mid i\\in[0,15]\\\\}$|\n| Histogram | Layer 1: $\\\\{(16+i, i) \\mid i\\in[0,15]\\\\}$|\n| Copy | $\\\\{(16+i, i) \\mid i\\in[0,15]\\\\}$|\n| Reverse | $\\\\{(16+i, 15-i) \\mid i\\in[0,15]\\\\}$|\n\n\n\n- “...*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.\n\n\nWe 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.\n\n*Questions on APM:*\n\n1. 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.\n\n2. Please see our response to weakness regarding APM above. \n\n3. 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.\n\nMiscellaneous:\n\n4. We agree and will replace degeneracy by a more suitable word as suggested.\n5. and 6. We will make the required changes in plot styling / color palette in the updated version.\n\nWe hope our response has addressed reviewer concerns."}, {"review_id": "ljWGN4izRx", "reviewer": "Reviewer_ybmK", "summary": "This 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).", "questions": "* 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}$?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## 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).", "was_revised": true, "final_justification": "As stated in my comment below, the authors sufficiently addressed my concerns. I've read other reviewers' comments and the authors' responses. Besides some mentioned criticisms, I remain positive of the present work. Specifically, I addresses a typically overlooked problem (early-training loss plateau) and through controlled studies identifies the core mechanisms. While I agree with reviewer dDr4 that the connection to LLMs is a bit ad-hoc and handwavy, I think it sufficiently suggests the presence of the observed phenomenon also in real LLMs. I believe the paper has also enough merit without this part and, thus, remain with my initial rating \"accept\".", "rebuttal": "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.\n \n*Weaknesses*:\n\n1. 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.\n \n    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,\n\n$$x_1, x_2, \\ldots, x_{16}, {\\rm SEP}, y_1, y_2, \\ldots, y_{13}$$\n\n$$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}$$\n\nI.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.\n\nWe 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.\n\n2. 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.\n\n3. 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.\n\n4. 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.\n\n5. We will provide code to replicate all experiments at a later stage.\n\n\n*Comments*:\n\n1. , 3., 4. We will make the necessary changes as suggested by the reviewer in the updated version of the manuscript.\n\n2. 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.\n\n\n\n*Questions*:\n\n1. 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).\n\n2. 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.\n\n3. 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.\n\n\n[Abbe 23]: Abbe et al., 2023. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. COLT 2023.\n\n[Varre 25]: Varre et al., 2025. Learning In-context n-grams with Transformers: Sub-n-grams Are Near-Stationary Points. ICML 2025."}, {"review_id": "SCzAfEr1K6", "reviewer": "Reviewer_dDr4", "summary": "In 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.", "questions": "1. 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.", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 1, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors clarified some of my initial concerns with additional experimental results and clarifications. However, still cannot recommend a full acceptance of the paper, as I still believe that the paper could have been made stronger by including more insights on why the observed phenomena emerge, to what extent they represent a real problem for language models, and how to alleviate them in practice. Also, I do not know to which extent the authors will address my concerns regarding issues in the manuscript itself; I sincerely hope that, in case of acceptance, all of the additional rebuttal results/discussion will be included in the paper.", "rebuttal": "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.\n\n*Weaknesses*:\n\n1. 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.\n  \na) “*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:\n\n**Hypothesis for slow learning of attention map**: \n\nOur 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. \n\n*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.  \n\nWe 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.\n\n\n**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).\n\n\nb) “*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.\n \nc) \"*…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.\n\n2. “*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.\n\n    “*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.\n\n3. “...*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. \n\n    To further illustrate this point, we train the Transformer on concatenated sequences of increasing complexity, as below,\n\n    $$x_1, x_2, \\ldots, x_{16}, {\\rm SEP}, y_1, y_2, \\ldots, y_{13}$$\n\n    $$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}$$\n\n    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.\n\n4. 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.\n\n|      |   0  |  1   |  2   | 4   | 8    | 16    | 32   |  64  |  128   |  256 |  512  |  1000   |  2000   | 4000  | 8000  | 10000  |\n|--|--     |--     |--     |--    |--     |--     |--    |--    |--     |--  |--    |--     |--     |--   |--    |--     |\n14M |  0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.97 | 0.93 | 0.98 | 0.39 | 0.13 | 0.17 | 0.12 | 0.13 | 0.11  | \n1B |  0.03 | 0.03 | 0.03 | 0.03 | 0.99 | 0.15 | 0.81 | 0.91 | 0.83 | 0.65 | 0.19 | 0.14 | 0.13 | 0.14 | 0.17 | 0.17  | \n1.4B |  0.16 | 0.16 | 0.16 | 0.2 | 0.99 | 0.73 | 0.89 | 0.86 | 0.8 | 0.13 | 0.15 | 0.15 | 0.13 | 0.14 | 0.15 | 0.17  | \n2.8B |  0.04 | 0.04 | 0.05 | 1.0 | 0.82 | 0.88 | 0.97 | 0.82 | 0.79 | 0.76 | 0.14 | 0.12 | 0.16 | 0.15 | 0.17 | 0.17  | \n\n\n\n\nWe will add these results through a plot similar to Fig. 8 in the updated version of the manuscript.\n\n5. 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. \n\n6. We will provide code to replicate all experiments at a later stage.\n\n\n\n*Questions:*\n\n1. 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$).\n\n2. 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.\n\n3. We will make necessary changes in plot styling / color palette in the updated version. \n\n4. 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. \n\nWe hope that our response has clarified the reviewer’s concerns about our work.\n\n[Abbe 23]: Abbe et al., 2023. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. COLT 2023.\n\n[Varre 25]: Varre et al., 2025. Learning In-context n-grams with Transformers: Sub-n-grams Are Near-Stationary Points. ICML 2025."}]

Author rebuttal:

Meta-review: This paper studies the training dynamics of transformers for simple algorithmic tasks.
In particular, the focus is about what happens during the long loss plateau phase of training, which is often followed by a sudden drop in loss.
It has been shown that during the loss plateau, the model will learn partial solutions while suffering from repetition bias and representation collapse, due to slow key-query alignment.
These phenomena can also be found in early training of large language models.

Overall the paper is well-written, and the authors have included extensive experimental details to illustrate the loss plateau phase.
The findings are valuable for further understanding of the training dynamics of transformers.
However, as noted by the reviewers, the current paper falls short in the investigation of the mechanisms behind the observed phenomena.
Nonetheless, the empirical observations would provide insights for future work.

For paper revision, the authors should include ablation studies with respect to the training configurations as provided during rebuttal.
Other comments and discussions should also be incorporated.
Decision: Accept (poster)