Stagewise Development in Neural Networks
> TLDR: This post accompanies The Developmental Landscape of In-Context Learning by Jesse Hoogland, George Wang, Matthew Farrugia-Roberts, Liam Carroll, Susan Wei and Daniel Murfet (2024), which shows that in-context learning emerges in discrete, interpretable developmental stages, and that these stages can be discovered in a model- and data-agnostic way by probing the local geometry of the loss landscape. Four months ago, we shared a discussion here of a paper which studied stagewise development in the toy model of superposition of Elhage et al. using ideas from Singular Learning Theory (SLT). The purpose of this document is to accompany a follow-up paper by Jesse Hoogland, George Wang, Matthew Farrugia-Roberts, Liam Carroll, Susan Wei and Daniel Murfet, which has taken a closer look at stagewise development in transformers at significantly larger scale, including language models, using an evolved version of these techniques. How does in-context learning emerge? In this paper, we looked at two different settings where in-context learning is known to emerge: * Small attention-only language transformers, modeled after Olsson et al. (3m parameters). * Transformers trained to perform linear regression in context, modeled after Raventos et al. (50k parameters). Changing geometry reveals a hidden stagewise development. We use two different geometric probes to automatically discover different developmental stages: * The local learning coefficient (LLC) of SLT, which measures the "basin broadness" (volume scaling ratio) of the loss landscape across the training trajectory. * Essential dynamics (ED), which consists of applying principal component analysis to (a discrete proxy of) the model's functional output across the training trajectory and analyzing the geometry of the resulting low-dimensional trajectory. In both settings, these probes reveal that training is separated into distinct developmental stages, many of which are "hidden" from the loss (Figure