Stagewise Development in Transformers and the Geometry of the Loss Landscape

Thank you for your review and for recommending that the paper should be accepted. We appreciate your positive feedback and you raising your concern about generalizability.

On generalizability to larger models: We appreciate your concern regarding the scalability of our methodology. Please see our top-level comment (on the scale of experiments) for a detailed response. Importantly, we agree that the specific stages that we identify may not generalize to larger models. In fact we do not claim that these stages will generalize to other models or with other hyperparameters. The thing that we claim is that we show evidence for the existence of a fundamental link between loss landscape geometry and neural network development. If this link truly exists, this is the thing that we think will generalize to larger models, which would allow future more sophisticated geometry-based methodologies to better understand development in larger models.

Thanks again! Given the new information presented in our top-level rebuttal clarifying the aims of our work, we wanted to ask if this addresses your concern about generalizability and whether you would consider strengthening your recommendation to accept the paper by improving your score? Either way, we welcome further discussion.

Dear Reviewer 3gci,

Thank you once again for your encouraging support of our paper. While there is still time in the discussion period, we wanted to check in with you.

We have taken steps to clarify the contributions and broader implications of our work. In particular we noted that our primary claim is not that the stage-detection methodology will necessarily generalize to different or larger models, but that we have contributed evidence of a deep link between loss landscape geometry and neural network development that has the potential to generalize (among other implications).

We wanted to ask whether you have considered the revised paper's new introduction and discussion sections, where we have made a clearer and more explicit case for the contribution of the paper than in the original submission. Has this new perspective at all strengthened your impression of the paper? If so, please consider raising your score.

Of course, if you have any remaining concerns or follow-up questions prompted by our revisions or the discussion with the other reviewers, we welcome further discussion.

Thank you for your review and for recognizing several strengths of our analysis approach. We also appreciate you sharing your concerns and questions about the work to which we offer the following comments.

On the relation between loss changes and LLC: We are not making a claim to a definitive connection between our developmental stages and phase transitions in the sense of statistical physics in this paper, although we note that prior work $1$ does explore this in smaller models. Could you please clarify what you mean by the “ICL loss and LLC do not match”? The only implied match in this paper is between critical points of the LLC curve and the stage boundaries.

On novel insights into transformer training dynamics: We respectfully disagree with the claim that our work “primarily reinforces existing findings from Olsson et al. (2022) rather than presenting novel insights into transformer learning dynamics”.

• We do more than reinforce the results of Olsson et al. (2022), we extend them. It’s true that we replicate the emergence of induction heads in two-layer attention-only transformer language models observed by Olsson et al. (2022). However, we also show that before induction heads form, our transformer adopts simpler strategies, following a progression akin to that found by Olsson et al. (2022) for fully-developed models with increasing depth. To our knowledge this is a novel contribution to the literature on transformer learning dynamics.
• Our contributions include more than just our observations about the specific dynamics of these transformers. Please see the discussion in the top-level comment on primary takeaways, where we enumerate several other novel implications of the extensive analysis that we contribute.

We appreciate your comment because we realized we could have made the significance of our contributions and the relationship to prior work much clearer in the paper. We have elaborated on these points in the revised discussion section. Please take a look at the revised discussion and let us know if you have any further concerns.

On the LLC vs. per-token loss: Great question. The primary advantage of the LLC over the per-token loss is that the LLC reveals additional information that is hidden from the per-token loss, offering a complementary perspective on model development. We can see this empirically in our results and it is also consistent with singular learning theory.

• Our results in in-context linear regression demonstrate this empirically: the in-distribution input-output mapping largely stabilizes after LR2. This is apparent in the per-token loss as well as per-token prediction norm (Figure 6a (was 4a), top). However, LLC continues to reveal significant developmental changes in stages LR3–LR5, including phenomena like layer normalization collapse that are invisible to purely behavioral metrics.
• In general, we expect (per-token) loss and the LLC to be highly complementary metrics. The LLC is grounded in singular learning theory through the free energy formula (Section 4 (was 3.1)), which provides a principled mathematical basis for using it to detect developmental stages and transitions in model complexity. Singular learning theory suggests that the LLC is, after the loss, the second most important observable for neural networks. Thus the two metrics would best be studied together.

On estimating LLC away from local minima: Another great question. We have addressed this in the top-level comment. Does the response there address your question?

Thanks again! We once again appreciate your feedback and questions. We are not 100% sure that we have addressed the relationship between loss changes and LLC, but look forward to further discussion. We feel that through the above response and our top-level comment and revisions we have addressed the claim that we have not offered novel results, and we are wondering if this is enough for you to reconsider your recommendation to reject the paper. Of course, if you have any other follow-up questions we are happy to elaborate.

$1$ Chen et al. (2023); as cited in manuscript.

Dear Reviewer E5MY,

Thank you once again for your review. While there are still a few days left of the discussion period, we wanted to quickly follow-up regarding your review. We are still hoping to discuss further with you about your listed concerns and questions:

1. We are still seeking clarification of your first stated concern, regarding quantifying the relationship between ICL, loss, and LLC. We have clarified that we make no claims about establishing "a definitive connection to phase transitions".
2. Regarding our contributions relative to prior work, we feel we have addressed this concern. We have clarified in our individual response, in our top-level comment, and in the new introduction in our revised paper that we draw several novel takeaways from our extensive analysis.
3. In response to your question about the advantages of LLC over per-token loss, we have clarified how we see the two metrics as complementary, which is a perspective supported by our empirical results and by the framework of singular learning theory.
4. In response to your question on the use of LLC estimation for points away from local minima, we have clarified that our methodology is empirically supported by prior work in smaller models, which we see as sufficient for us to draw the conclusions we draw (and while a full theoretical justification for this approach is currently lacking, we have prominently flagged this as a limitation in the revised paper).

Please let us know whether our responses have been sufficient to allay your concerns about the paper. In the case of point (1), we would greatly value clarification so that we can better understand and then address this concern. If we have indeed addressed your concerns, we once again invite you to consider increasing your rating. Otherwise, if you have remaining or new questions or concerns, we invite you to share them so that we can make use of the remainder of the discussion period.

Thank you again for the time and attention you have paid to our work.

On longer training: We believe that there are no additional stages to be found by lengthening training, at least with the current hyperparameters. Notice the log scale on the x-axis for e.g. Fig 1, so that in the language model setting the final stage is from 17k-50k steps, and in the linear regression setting from 320k-500k. So the model spends in both cases a significant fraction of the training time in the final stage, and we see no evidence that longer training would reveal interesting new development.

On the mathematical basis for d(LLC)/dt = 0: Great question. The link between SGD training and the way the Bayesian posterior changes in the singular learning process remains unclear, so it is not possible to give a precise justification for looking for points where d(LLC)/dt = 0 at the present time. However the motivation is as follows: in smaller models (such as $1$ ) the developmental stages correspond to the learning trajectory passing from the neighborhood of one critical point of the population loss to another, and each of these critical points has an associated LLC value. If the trajectory is experiencing “critical slowing down” near one of these critical points, the LLC estimate should stabilize near this value and this is exactly what was observed in $1$ . In larger models the picture is more complex and not fully understood, and a simple picture of the learning trajectory passing near critical points is unlikely to be correct. Nonetheless our overall expectation is still that at some points in training the LLC estimate should stabilize when important structures are “nearing completion” and this motivates the d(LLC)/dt = 0 criterion.

We appreciate your question because this reflects the fact that we did not spell out this connection in the submitted version of the paper. In our latest revision, our updated methodology section (now Section 4) spells out this motivation.

On learning rate schedulers: We agree that a dynamic learning rate could influence the developmental stages and think this is an important direction for future work. Our speculation would be that this would mostly affect the later stages in training, on the theory that these are about learning “finer” patterns in the data and the discovery of these patterns is what a decaying learning rate is enabling. So a reasonable expectation would be that with a decaying learning rate the early stages would look the same, but (returning to your earlier question) there might be additional structure to be uncovered in the final stage which is currently invisible.

Thanks again: We hope we have been able to address most of your questions and concerns about the paper. If so, we invite you to reconsider your recommendation to reject the paper. If you still have concerns or if you have any follow-up questions, please do let us know. Looking forward to further discussion.

$1$ Chen et al. (2023); as cited in manuscript.

Thank you for your detailed review and for sharing your concerns and questions about the paper. Thanks especially for acknowledging the potential of this research to complement work on mechanistic interpretability. We are looking forward to further discussion.

On theoretical foundations: We appreciate you raising this concern. We think we have a strong response.

• First, we claim that our methodology is deeply principled. The core theoretical foundation of our approach is based on singular learning theory, which is an established body of work published in mathematics, statistics, and machine learning communities for over two decades. It’s true that the particular tool that we use for quantifying degeneracy, Lau et al.’s SGLD-based LLC estimator, is a recent idea that, to our knowledge, has not yet been published in a peer-reviewed deep learning venue.
• Second, in light of the clarification in our top-level comment about the primary takeaway of our work, we contend that in order to establish our main claim (that our results point to a link between geometry and development), it is not necessary to have a perfect LLC estimation methodology. We contend that we have contributed evidence supporting this link to a sufficient degree that warrants publication, so that future work (including work on improved LLC estimation approaches and more extensive replication of our findings in a broader class of settings) can continue to investigate this link and its implications for better understanding neural network development.

In recognition of the fact that the LLC is not a widely-used metric within the deep learning community, we have rewritten the methodology section in our latest revision. This includes offering a more detailed and accessible introduction to the LLC, to Lau et al.’s SGLD-based LLC estimation procedure, and to the singular learning theory ideas motivating the other parts of our stage-identification methodology (e.g. using critical points of the estimated LLC curve to mark stage boundaries). We hope this makes our paper more self-contained.

We hope that this explains why we feel confident using Lau et al.’s SGLD-based LLC estimation procedure as one part of our methodology even though it has not been peer-reviewed. Have we allayed your concern on the matter?

On figure 4b (was 2b): Your main question as we understand it is, why are the previous token heads identified as such when there are other heads that increase in their previous token attention score in LM2? Our identification of heads 1:2 and 1:5 as previous token heads is due to the combination of the previous token score and the composition score of those heads with the induction heads, which can be seen in Figure B.5. The other attention heads that increase in previous token attention score in LM2 do not have a high composition score with the induction heads. So, they do not form part of the induction circuit. (As for why other heads have their previous token score increase in LM2, these heads are responsible for predicting the $n$-grams that are learned in LM2, which naturally requires an attention head to attend to recent token positions.)

We appreciate you raising this concern and giving us the opportunity to improve the section by clarifying these points explicitly. We have attempted to do so in the revision, please let us know what you think.

On data distribution for computing loss: We appreciate your feedback. Regarding your questions about the validation data used for LLC estimation:

• We have documented some of these details in the appendices, namely in section D.1.4 and table 2 (was 3) for language modeling, and section D.2.4 and table 3 (was 4) for in-context linear regression, which describe the data distributions and numbers of samples used in SGLD-based LLC estimation. We realized we didn’t link these appendices clearly from the main text—we have corrected this in the revision.
• As described in these appendices for SGLD we use an independently sampled data set (independent from the training data set; a different subset of The Pile for language modeling, and a fresh sample of synthetic data from the same distribution for in-context linear regression). In the process of SGLD we compute the loss based on minibatches drawn from these data sets.
• Of course, there will be small fluctuations in the loss landscape if we generate different data samples or different batches, but since we consider large samples from the same underlying distribution, we don’t think this kind of resampling would cause the SGLD loss landscape to change enough to create a meaningful difference in our estimates.

Is this information satisfactory, or is there anything else we can clarify? In our revision we have more clearly indicated the location of this information in the appendices at the appropriate point in the main text. Does that resolve your concern about the paper not having enough clarity on important experimental details?

Dear Reviewer sxSg,

Thank you again for your engagement with our paper and for your openness to further discussion. While there are still a few days left of the discussion period, we wanted to quickly follow-up on this.

We have attempted to address most of your stated concerns about the paper:

1. Regarding the concern that the theory is not based on widely-accepted prior work, we have pointed out that the LLC estimation techniques we use are new but, like the rest of our methodology, are soundly based on a firm foundation of published work in singular learning theory. We have also attempted to clarify the presentation of our methodology and how it is derived from this prior work in sections 3 and 4 of our revised paper, which we invite you review.
2. Regarding the concern about the classification of previous token heads in figure 2b (now 4b), we have clarified that we identified these heads based on their role in the induction circuit and not only on their previous token score.
3. Regarding the concern about missing experimental details for measuring the loss during LLC estimation, we have pointed out these details in the appendix and we have included more prominent links to these details in the revised main text at the appropriate place.

We invite you to share any follow-up questions, or let us know if the above responses have not sufficiently addressed your concerns, or if you have any other remaining concerns. Otherwise, we invite you to reconsider your rating. We are very much looking forward to a fruitful discussion over the remainder of the discussion period.

Thank you for your review and for sharing your concerns and questions about our work. We would like to respond to each of your concerns and questions in detail.

On limitations of model selection: Please see our top-level comment (on the scale of experiments). Importantly, we agree that our stage identification methodology may not transfer to larger models, and we do not claim that it will. The models are merely case studies, chosen to have non-trivial scale while still being within the realm of interpretability (which is required for our validation step). We claim their scale is sufficient for our results to demonstrate meaningful evidence for the existence of a fundamental link between geometry and development in deep learning. As you point out in your review, we also demonstrate this link in two different settings. Do you agree?

Regarding your question (Q2), in theory our method should scale to larger models, with the caveats in the top-level comment. In more detail, LLC estimation works in principle for any architecture as long as we can run SGLD. The free energy formula underlying our stage identification methodology holds equally for larger models.

As discussed in the top-level comment, the primary barrier to investigating larger models is not to stage identification, but only to validation: interpreting the developmental stages might be more challenging due to the increased model complexity.

On LLC estimation away from minima: We agree this is a limitation, but we do not think it invalidates our approach. Please see the top-level comment for details, and let us know what you think.

On sensitivity of LLC estimation: We adhered to guidelines established in prior work on LLC estimation for selecting SGLD hyperparameters while minimizing bias in LLC estimates. We detail our procedure in appendix D.3. Does this approach address your concern?

On layer normalization collapse: We appreciate your interest in this phenomenon. You mentioned that you saw it as a strength that it can bring new insights to model stability and robustness. We wanted to clarify that we don’t see layer normalization collapse as necessarily having implications for model stability or robustness. Rather, we see it as a possible mechanism behind the phenomenon that in some stages, the LLC decreases (please see the top-level comment, extra takeaway 3, for further clarification).

As you also note under weaknesses, we could have done more to clarify the implications of this phenomenon, and to provide an in-depth analysis. In our latest revision, we have included additional discussion on how layer normalization collapse (along with other observed changes in structure in this model) connects to an increase in degeneracy in appendix C.4. In C.4.6 we discuss how these degeneracy decreases can be consistent with the free energy formula, but the precise interpretation remains an open theoretical problem.

Regarding your question (Q3) about the rate and extent of LN collapse:

• Extent: The LN collapse is near-complete for the unembedding (see insets in Figure C.7): more than 62/64=97% of the unembedding LN weights are less than 0.1 in magnitude by the end of training. The next-largest collapse occurs in the block 1 pre-attention LN weights, where about half of the weights collapse. In the remaining LN layers, the weights do visually appear to collapse (main plots in Figure C.7), but the weights do not reach as small a norm.
• Rate: The collapse occurs over the course of stages LR3 and LR4 at the same time for all of these different layers. Individual weights can collapse in as few as ~10k steps and in other cases closer to ~100k steps (main plots in Figure C.7).

In other seeds we observe the same extent and rate of collapse. We have not explored the collapse in other training regimes, though this is an interesting direction for future work, especially as it relates to the link between geometry and development.

On stages in image models: Thank you for this intriguing question (Q1). We have not yet studied links between geometry and development in image models. However, we see this as a promising avenue for future work, as the more disparate settings in which we see evidence for links between geometry and development, the more support this could provide for the fundamental nature of this link.

Thanks again. Has the above reply and the clarification of the aims of our work in the top-level comment resolved most of your questions and concerns about the paper? If so, please consider raising your score. Also, please let us know if you have any follow-up questions, we are looking forward to further discussion.

Dear Reviewer PMP3,

While there are still a few days left of the discussion period, we wanted to quickly follow-up regarding your review and our response. We feel we have addressed most of your stated concerns about the paper:

1. Regarding the generalization to larger models, we have clarified that we do not claim generalization to larger models, and our top-level comment and revised paper clarifies what we do see as the takeaways from our work, centering around the idea that our experiments in two settings provide evidence of a fundamental link between geometry and development.
2. Regarding the limitations of local learning coefficient estimation, we have clarified how we went about tuning the estimates to ensure our results are robust.
3. Regarding the lack of details on layer normalization collapse, we have clarified that we see this as a potential mechanism for the phenomenon of model simplification and we have added additional details in our response and in our revised paper.

Could you please confirm whether the above has addressed your concerns? If so, we invite you once more to reconsider your rating and share any follow-up questions prompted by our response or the discussion with other reviewers. Thank you again for your consideration, we look forward to further discussion.

On the scale of experiments: Reviewers PMP3 and 3gci raised concerns that while our results hold for two-layer transformers, our methodology might not generalize to LLMs. We want to clarify our choice of models and the ways in which we do and do not expect our results to generalize.

As explained, our primary claim is that our results point to a link between loss landscape geometry and deep learning. With this in mind, we chose models to occupy a ‘sweet spot’ between interpretable toy models and practical LLMs.

• Prior work $1$ studied the link between geometry and development in a toy autoencoder with around 20 parameters. Because of the small size, they could analyze this model’s loss landscape and the development of its internal structure in precise detail, and the connection between geometry and development was clearly demonstrated.
• In this paper, we studied models 2,500x to 150,000x larger trained on substantially more complex data (in-context linear regression, natural language). Our models are of comparable scale to those studied in other recent published works on the science of deep learning $2,3$ . This work provides useful prior knowledge about behavioral/structural developments that we can use to validate that our geometric stage divisions are meaningful.
• Analyzing a substantially larger model (e.g. on the order of BERT or a GPT) would offer a more direct demonstration of the practical implications of the links between geometry and development. However, we expect that in larger models the manifestation of this link will be more complicated than the clear developmental stages identified in the present work and may require the introduction of more refined techniques (see e.g. line 497 (was 480) of the paper).

Our paper demonstrates that there is a link between loss landscape geometry (as measured by the LLC) and neural network development in transformers at the scale of 3M parameters. At this scale, the link can be summarized by the clear developmental stages that we study. As the scale of models increases, we observe that the developmental stages become less clear, with our hypothesis being that larger models learn multiple things in parallel; for example the stages in the present work are more gradual than the sharp transitions in the smaller models studied in $1$ . At some scale we expect that the link between loss landscape geometry and development is better described by a more refined notion of developmental stage, and the simple methodology in the present paper may no longer be applicable.

Nonetheless we view the present paper as providing an essential foundation for the development of more sophisticated techniques that are applicable to larger models, since it validates the basic premise carefully in smaller models where the ground truth is easier to access.

LLC estimation away from local minima: Reviewers PMP3 and E5MY also noted that a rigorous theoretical justification for estimating the LLC away from local minima is currently missing. The meaningfulness of LLC estimates in these cases is empirically supported by $1$ . We agree that studying the LLC away from local minima would be valuable and represents an important future direction, though we consider it beyond our scope, and have explicitly acknowledged it as a limitation of the current work.

Thank you again. We look forward to the remainder of the discussion.

$1$ Chen et al. (2023); $2$ Olsson et al. (2022); $3$ Raventós et al. (2023); as cited in manuscript.