Stagewise Development in Transformers and the Geometry of the Loss Landscape

Jesse Hoogland, George Wang, Matthew Farrugia-Roberts, Liam Carroll, Susan Wei, Daniel Murfet

Keywords: Science of deep learning, loss landscape geometry, training dynamics, singular learning theory

TL;DR: Transformers learn in discrete developmental stages that can be discovered by studying the local geometry of the loss landscape.

Abstract: Deep learning involves navigating a high-dimensional parameter space guided by the loss landscape. In the process, complex computational structures form and re-form inside the neural network, leading to shifts in input--output behavior. It is a priority for the science of deep learning to uncover principles governing the development of neural network structure and behavior. Drawing from the framework of singular learning theory, we propose that model development is governed by the local geometry of the loss landscape. We investigate this link by monitoring the geometry of the loss landscape throughout training for transformers trained as language models or for a synthetic in-context regression task. We divide training into ``developmental stages'' marking discrete shifts in loss landscape geometry. We then confirm that these stages coincide with significant changes in the internal computational structure and the input--output behavior of our models. Our findings provide new insights into transformer development and underscore the potential of a geometric perspective for understanding modern deep learning.

Primary Area: other topics in machine learning (i.e., none of the above)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6078