Loss Landscape Degeneracy and Stagewise Development of Transformers

The computational cost per LLC estimate is proportional to a standard training step, denoted $S$. We expect to require a constant number $CT_{\text{SGLD}}$ of samples (on the order of $10^{2}$–$10^{4}$) to yield robust estimates, independent of model size. Using logarithmically spaced checkpoints, the total computational complexity for generating an LLC curve over the entire training process scales as $O(SCT_{\text{SGLD}}\log N)$, where $N$ is the total number of training steps. The space complexity incurs a modest linear overhead compared to standard SGD, requiring storage for one additional copy of the weights to enable localization.

For language models, we use SGLD to sample 20 independent chains with 200 steps per chain and 1 sample per step. For the one-layer model, we used $\epsilon=0.003,\gamma=300$, and for the two-layer model we used $\epsilon=0.001,\gamma=100$. Estimating the LLC across all checkpoints took around 200 GPU hours for the two-layer model on a single A100 and around 125 GPU hours for the one-layer model. For additional runs of the two-layer model, we ran fewer chains, bringing the time down to about 2 TPU hours per training run.

We sampled a separate set of 1 million lines (lines 10m-11m) from the DSIR filtered Pile, denoted as $D_{\text{sgld}}$. The first 100,000 lines from this SGLD set (lines 10m-10.1m) were used as a validation set. The sampling of batches for SGLD mirrored the approach taken during the primary training phase. Each SGLD estimation pass was seeded analogously, so, at different checkpoints, the SGLD chains encounter the same selection of batches and injected Gaussian noise.

For in-context linear regression models, we generate a fixed dataset of $2^{20}$ samples. Using SGLD, we sample $10$ independent chains with 5,000 steps per chain, of which the first 1,000 are discarded as a burn-in, after which we draw observations once per step, at a temperature $n\beta=66.7$, $\epsilon=0.0003$, and $\gamma=13.3$, over batches of size $m=1024$. LLC estimation takes up to 72 CPU-hours per training run.

In Lau et al. (2025), the inverse temperature $\beta$ is set to a fixed “optimal” value $\beta^{*}=1/\log n$, where $n$ is the number of training samples. In practice, we find that it can be advantageous to sample at a higher temperature.

Since $\beta$ always shows up in a product with $n$ (in ˜8 for the SGLD step and in ˜3 for the LLC), we can view the inverse temperature as a multiplier that adjusts the effective size of your dataset. In a Bayesian setting, $\beta=2$ would mean updating twice on each of the samples in your dataset.

The problem with the default choice of $\beta^{*}$ is that as we increase $n$ we have to decrease the SGLD step size $\epsilon$ to prevent the update from becoming ill-conditioned, and this eventually causes the gradient term to suppress the noise term. This, in turn, leads to requiring larger batches to suppress the gradient noise and requiring longer chains to sufficiently explore the local posterior (Section˜A.5.3).

Instead of $n\beta=n/\log n$, we perform LLC estimation at $n\beta=m/\log m$, where $m$ is the SGLD batch size.

To smooth out the $\hat{\lambda}_{t}$ curves, we reset the random seed before LLC estimation run at each checkpoint. This means the sequence of injected Gaussian noise is the same for LLC estimation runs at different checkpoints. Additionally, if the batch size is held constant, the batch schedule will also be constant across different estimation runs. Figure˜A.3 shows that this does not affect the overall shape of the learning coefficient curves; it simply smooths it out.

As a rule of thumb, $\epsilon$ should be large enough that the $\hat{\lambda}$ estimate converges within the $T_{\mathrm{SGLD}}$ steps of each chain but not too large that you run into issues with numerical stability and divergent estimates. Subject to this constraint, $\gamma$ should be as small as possible to encourage exploration without enabling the chains to “escape” to nearby better optima, and $\beta$ should be as large as possible (but no greater than $1/\log n$).

To determine the optimal SGLD hyperparameters, we perform a grid sweep over a reparametrization of the SGLD steps in terms of $\tilde{\beta},\tilde{\gamma},\varepsilon$:

where $\tilde{\beta}=\varepsilon\beta n/2$, $\tilde{\gamma}=\varepsilon\gamma/4$.

The results of this hyperparameter sweep are illustrated in Figure˜A.4 for final checkpoints. Separately (not pictured), we check the resulting hyperparameters for a subset of earlier checkpoints. This is needed since, for example, a well-behaved set of hyperparameters at the end of training may lead to failures like divergent estimates (Figure˜A.5) earlier in training when the geometry is more complex and thus the chains less stable.

As a useful diagnostic when calibrating the LLC estimates, we propose an online variant for learning coefficient estimation. When overlaid on top of individual-chain LLC traces, this helps reveal common failure modes like divergent estimates, non-converged estimates, and escapes (Figure˜A.5). These traces display the running estimate of $\hat{\lambda}$ as a function of the number of steps taken in a chain (with the estimate averaged across independent chains).

Define $\hat{\lambda}_{\tau}(w_{0})$, the LLC at $w_{0}$ after $\tau$ time-steps for a single SGLD chain as follows (Lau et al., 2025):

Moving terms around, we get,

where

This can be easily extended to an online estimate over chains by averaging the update $n\beta\left(\ell_{n}(w_{\tau})-\ell_{n}(w_{0})\right)$ over multiple chains.

We empirically estimate the conditional bigram distribution by counting instances of bigrams over the training data. From this, we obtain the conditional distribution $\tilde{q}(t^{\prime}|t)$, the likelihood that a token $t^{\prime}$ follows $t$. The bigram score $B_{k}^{S}$ at index $k$ of an input context $S$ is the cross entropy between the model’s predictions $p(t_{k+1}|t_{k})$ at that position and the empirical bigram distribution,

where the $t_{k+1}^{(i)}$ range over the possible second tokens from the tokenizer vocabulary. From this we obtain the average bigram score

where we take fixed random sequences of $k_{i}$ and $S_{i}$ for $1\leq i\leq n=5,000$, which is displayed over training in Figure˜4(a). This is compared against the best-achievable bigram score, which is the bigram distribution entropy itself, averaged over the validation set.

In stage LM2 we consider $n$-grams, which are sequences of $n$ consecutive tokens, meaning $2$-grams and bigrams are the same. Specifically, we consider common $n$-grams, which is defined heuristically by comparing our 5,000 vocab size tokenizer with the full GPT-2 tokenizer. We use the GPT-2 tokenizer as our heuristic because its vocabulary is constructed iteratively by merging the most frequent pairs of tokens.

We first tokenize the tokens in the full GPT-2 vocabulary to get a list of 50,257 $n$-grams for various $n$. The first 5,000 such $n$-grams are all $1$-grams, after which $2$-grams begin appearing, then $3$-grams, $4$-grams, and so on (where $2$-grams and $3$-grams may still continue to appear later in the vocabulary). We then define the set of common $n$-grams as the first 1,000 $n$-grams that appear in this list for a fixed $n$, $n\geq 2$.

If we track the performance on $n$-grams and see it improve, we may ask whether this is simply a function of the model learning to use more context in general, rather than specifically improving on the set of $n$-grams being tracked. We measure performance against this baseline by defining an $n$-gram score. For a fixed $n$, we obtain the average loss $\ell_{\textrm{gram}}^{n}$ of the model on predicting the final tokens of our set of 1,000 $n$-grams and also obtain the average loss $\ell_{\textrm{test}}^{n}$ of the model on a validation set at position $n$ of each validation sequence. The $n$-gram score is then defined to be $\ell_{\textrm{test}}^{n}/\ell_{\textrm{gram}}^{n}$.

The in-context learning score is a behavioral measure of the relative performance of a model later in a sequence versus earlier in the sequence. We define $\text{ICL}_{k_{1}:k_{2}}$ to be the loss on token $k_{2}$ minus the loss on token $k_{1}$, so a more negative score indicates better relative performance later in the sequence. A more negative ICL score does not, however, mean that a model is achieving better overall loss on later tokens; it is only about the relative improvement. For the language model, we follow a similar construction as Olsson et al. (2022), where we take $k_{2}$ to be the 500th token and $k_{1}$ to be the 50th token. This is then averaged over a 100k-row validation dataset. The performance of the language model over the course of training can be seen at the bottom of Figure˜4(f).

In Figure˜C.1, we visualize changes in the model’s input/output behavior by comparing model predictions before and after developmental stages and highlighting tokens with the greatest differences.

In Figure˜C.2, we measure the effect of the positional embedding on model performance by comparing the model’s performance at particular context positions on a validation set over the course of training against performance on the same validation set but with the positional embedding zero-ablated. The full context length is 1024, and we measure test loss at positions 1, 2, 3, 5, 10, 20, 30, 50, 100, 200, 300, 500, and 1000. In the transition from stage LM1 to LM2, the model begins using the learnable positional embedding to improve performance. The difference between test loss with and without the positional ablation is negligible at all measured positions until the LM1–LM2 boundary.

Structurally, we might predict that the positional embeddings should organize themselves in a particular way: in order to understand relative positions, adjacent positions should be embedded close to each other, and far-away positions should be embedded far apart.

In Figure˜C.3, we examine the development of the positional embedding itself over time from two angles. The first is to take the embeddings of each position in the context and to run PCA on those embeddings. The result is that as training progresses, the positional embedding PCAs gradually resolve into Lissajous curves, suggesting that the positional embeddings might look like a random walk (Antognini & Sohl-Dickstein, 2018; Shinn, 2023). However, if we look to the explained variance, we see that it grows very large for PC1, reaching $94.2\%$ at training step 6400. This is much higher than we would expect for Brownian motion, where we expect to see about $61\%$ explained variance in PC1 (Antognini & Sohl-Dickstein, 2018).

The second perspective we use is to look at how the magnitudes of positional embeddings over the context length develop. In this case, we observe that the magnitudes seem to have a fairly regular structure. In conjunction with the PCAs and explained variance, we might infer that the positional embeddings look approximately like a (possibly curved) line in $d_{\textrm{model}}=256$ dimensional space. A positional embedding organized in this way would make it easier for an attention head to attend to multiple recent tokens, which is necessary if a single head is to learn $n$-grams.

Let $W_{Q}^{h},W_{K}^{h},W_{V}^{h}$ be the query, key, and value weights of attention head $h$ respectively. There are three types of composition between attention heads in transformer models in Elhage et al. (2021):

• Q-Composition: the query matrix $W_{Q}^{h}$ of an attention head reads in a subspace affected by a previous head

• K-Composition: the key matrix $W_{K}^{h}$ of an attention head reads in a subspace affected by a previous head

• V-Composition: the value matrix $W_{V}^{h}$ of an attention head reads in a subspace affected by a previous head

If $W_{O}^{h}$ is the output matrix of an attention head, then $W_{QK}^{h}=W_{Q}^{h\ T}W_{K}^{h}$ and $W_{OV}^{h}=W_{O}^{h}W_{V}^{h}$. The composition scores are

Where $M=W_{QK}^{h_{2}\ T}$, $M=W_{QK}^{h_{2}}$, and $M=W_{OV}^{h_{2}}$ for Q-, K-, and V-Composition respectively. See Figure˜C.4 for K-composition scores over time between attention heads in the induction circuits.

The previous-token matching score is a structural measure of induction head attention. It is the attention score given to $[A]$ by an attention head at $[B]$ in the sequence $\ldots[A][B]$ (i.e., how much the head attends to the immediately preceding token).

We compute this score using a synthetic data generating process, generating 10k fixed random sequences with length between 16 and 64. The first token is a special “beginning of string" token, and the remaining tokens are uniformly randomly sampled from other tokens in the vocabulary.

For each sample in this synthetic dataset, we measure the attention score that an attention head gives to the previous token when at the last token in the sequence. These scores are averaged across the dataset to produce the previous-token matching score for that attention head at a given checkpoint. The progression of previous-token matching scores over time can be seen in Figure˜4(d).

The prefix matching score from Olsson et al. (2022) is defined similarly to the previous-token matching score. Given a sequence $[A][B]\ldots[A]$, the prefix matching score of a particular attention head is how much the attention head attends back to the first instance of $[A]$ when at the second instance of $[A]$.

We compute this score using a synthetic data-generating process. We first generate 10k fixed random sequences of length 128. The first token is always a special “beginning of string" token and the $[A]$ and $[B]$ tokens are selected and placed randomly. One $[A]$ token is placed in the first half of the sequence, the other is placed in the second half, and the $[B]$ token is placed directly after the first $[A]$ token. The remaining tokens are randomly sampled from the tokenizer vocabulary, excluding the $[A]$, $[B]$, and beginning of string tokens.

For each sample in this synthetic dataset, we measure the attention score that each attention head assigns to the earlier instance of $[A]$ from the latter instance of $[A]$. These scores are averaged across the dataset to produce the prefix matching score for that attention head at a given checkpoint. The progression of prefix matching scores over time can be seen in Figure˜4(e).

In addition to training models on a data distribution in which tasks $\mathbf{t}$ are generated on-the-fly, we examine the setting of Raventós et al. (2023), in which a finite set of $M$ tasks is generated ahead of time, and training samples involve randomly selected tasks from this set.

Figure˜D.1 depicts (a) the mean square distance between the model’s predictions and the zero prediction in addition to (b) the mean square distance between the model’s predictions and the “task prior” prediction, using the component-wise averaged $\overline{\mathbf{t}missing}$ over the set of tasks encountered during training. For all models, the minimum distance to the task prior prediction is lower than the minimum distance to the zero prediction. Hence, we call stage LR1 “learning the task prior” rather than simply learning the zero prediction.

We consider two variants of the ICL score: $\operatorname{ICL}_{1:D}$, and $\operatorname{ICL}_{D:K}$.

If the noise term $\sigma^{2}$ equals zero and both tasks $\mathbf{t}$ and inputs $x_{k}$ are normalized (i.e., $\mathbf{t}\in S^{D-1}$), then $D-1$ observations of input/output pairs are enough to precisely identify $\mathbf{t}$. Therefore, $\operatorname{ICL}_{1:D}$ measures how successful the model is at initially locating the task. The fact that the tasks and inputs are not normalized changes this only slightly: the task will still sit near $S^{D-1}$ within a shell of vanishing thickness as $D\to\infty$.

Once localized, $\operatorname{ICL}_{D:K}$ measures how successfully the model refines its internal estimate of $\mathbf{t}$ with additional examples, which it can use to reduce the error due to noise.

In terms of implementation, it’s not necessary for the model to internally make a distinction between locating and refining its estimate of the task. For example, ridge regression makes no distinction. Still, we find it useful for reasoning about the progression of the model. In particular, we note that early in stage LR2, while the model begins to develop ICL for early tokens, it becomes worse at ICL over tokens late in the context. Later, at around 23k steps, $\operatorname{ICL}_{D:K}$ stabilizes, while $\operatorname{ICL}_{1:D}$ continues improving over the entire training run.

To further investigate behavior in stages LR2 and LR3, we probe the model on data sampled from different distributions than encountered during training.¹¹1Cf. Raventós et al. (2023) evaluating models trained on a set of discrete tasks on the “true” distribution consisting of novel tasks. We evaluate behavior on two families of perturbations: “OOD inputs” $x_{k}$, sampled according to a different scale

for some gain parameter $g$, and “OOD tasks”

Note that these inputs and tasks are not out-of-distribution in the sense of coming from a distribution with a different support than the training distribution. However, the samples drawn from these “extreme” distributions are exponentially suppressed by the original training distribution. Figure˜D.3 plots the normalized MSE for these two distributions over training time.

Between $t=1k$ and $t=4\mathrm{k}$ the model’s outputs rapidly diminish in scale for out-of-distribution samples, both for $g>1$ and $g<1$, especially for out-of-distribution inputs. While the model is moving away from predicting with the task prior for in-distribution samples, it moves closer to predicting with the task prior for-in-distribution samples.

Between $t=4\mathrm{k}$ and $t=23\mathrm{k}$, the model recovers on moderately out-of-distribution inputs $g<10^{1.5}$ with performance remaining close to constant beyond this range. Past this stage, performance improves constantly for out-of-distribution tasks.

For out-of-distribution inputs, performance eventually worsens for some ranges of $g$. Between $t=23\mathrm{k}$ and $t=80\mathrm{k}$ the model further approaches the task prior prediction for extreme out-of-distribution inputs $g>10^{1.5}$ . Subsequently, between $t=75\mathrm{k}$ and $t=130\mathrm{k}$ the model moves away from the task prior prediction for extreme inputs, and performance deteriorates for inputs with $g>10^{0.5}$. As of LR5, performance is roughly constant.

The embedding matrix $W_{E}$ is a linear transformation from $\mathbb{R}^{D+1}\to\mathbb{R}^{{d_{\text{embed}}}}$. Plotting the $D+1$ singular values of this matrix, we notice that the embedding partially loses one of its components starting at the end of LR2 (Figure˜D.4a).

The input “tokens” $x_{k}$ span a $D$-dimensional subspace of the $(D+1)$-dimensional “token space.” The target tokens $y_{k}$ span an orthogonal $1$-dimensional subspace. The collapse of one of the embedding matrix’s singular values means that the model learns to redundantly encode the inputs and targets in the same $D$-dimensional subspace of the space of residual stream activations. The almost order of magnitude separation in the magnitudes of the square singular value means that the $(D+1)$^th component of the token embedding explains only 2.9% of the variance in activations of the residual stream immediately after the embedding, whereas the dominant components explain roughly 24% each.

The positional encoding goes through a similar collapse to the unembedding starting during the second part of LR2 and continuing into LR3 (Figure˜D.4b). Additionally, throughout these stages, the subspace spanned by the embedding becomes more aligned with the subspace spanned by the positional encoding (Figure˜D.4c).

Over the course of training, we observe that some attention heads learn to attend solely (soft attention becomes hard attention) and consistently to certain positions (the attention pattern becomes content-independent). We call this phenomenon attention collapse in parallel with the other observed forms of collapse. Not only does this potentially contribute to a decrease in the LLC, but it also makes the attention heads identifiable: we find a self-attention head, previous-attention heads, previous-$x$-attention heads, and previous-$y$-attention heads.