Tending Towards Stability: Convergence Challenges in Small Language Models (2024)

Abstract: Increasing the number of parameters in language models is a common strategy to enhance their performance. However, smaller language models remain valuable due to their lower operational costs. Despite their advantages, smaller models frequently underperform compared to their larger counterparts, even when provided with equivalent data and computational resources. Specifically, their performance tends to degrade in the late pretraining phase. This is anecdotally attributed to their reduced representational capacity. Yet, the exact causes of this performance degradation remain unclear. We use the Pythia model suite to analyse the training dynamics that underlie this phenomenon. Across different model sizes, we investigate the convergence of the Attention and MLP activations to their final state and examine how the effective rank of their parameters influences this process. We find that nearly all layers in larger models stabilise early in training - within the first 20% - whereas layers in smaller models exhibit slower and less stable convergence, especially when their parameters have lower effective rank. By linking the convergence of layers’ activations to their parameters’ effective rank, our analyses can guide future work to address inefficiencies in the learning dynamics of small models.

## introduction scaling the number of parameters in language models (lms) has provided impressive performance gains on a variety of tasks @xcite and has become the de facto standard to make progress in model design (e.g., @xcite . small lms, however, remain essential as they are more practical: lower training and inference costs result in a smaller environmental impact @xcite . small lms empower individuals to train on proprietary data by requiring fewer resources, enhancing data privacy @xcite and democratising access to language modelling technology @xcite . however, for the same data and computational budget, small lms (unsurprisingly) underperform larger ones @xcite and (importantly) their performance tends to degrade in the late pretraining phase, a phenomenon termed saturation by @xcite . 1 saturation is typically attributed to the "limited representational capacity" of small lms; besides this anectodal justification, our understanding of its causes is still limited. recently, @xcite linked saturation to the reduced variability of the output embeddings of lms caused by the mismatch between the hidden model dimension and the vocabulary size @xcite . specifically, the last layer of lms maps the hidden representation of random tokens to output embeddings with high cosine similarity. 2 in this paper, we use the pythia model suite @xcite to provide orthogonal analyses that consider models' training dynamics. first, we study how the activations of the attention and mlp layers converge to their final state across lms of different sizes. then, we relate the difference in convergence behaviour across sizes to the effective rank of their parameters: layers whose activations converge later in training span a smaller fraction of their dimensions. specifically, we first use the centered kernel alignment (cka; @xcite metric to measure the similarity of layers' activations across checkpoints. we observe that larger lms converge faster and more smoothly to their final state. as shown in fig. 1 , within the first 20% of training nearly all layers in the larger lm (2.8b) resemble their final state, while most layers in the smaller lm (160m) remain different for most of training. we then find a strong correlation between the convergence pattern of a layer's activations and the rank of its parameters and gradients. we introduce the concept of proportional effective rank ( §3) to consistently compare these effective ranks across model sizes. our analyses highlight training inefficiencies in small-scale lms, paving the way for targeted improvements in future work. ## related work prior work has studied various learning dynamics of the pythia suite, including memorisation @xcite , training data influence @xcite , and statistics of learned embeddings @xcite . related to our work, @xcite examine the differences in the rank of the unembedding matrix (mapping from hidden representations to tokens) across model sizes, known as the softmax bottleneck @xcite . unlike their findings, we focus on the convergence dynamics of all layers. similarity metrics like cka and singular vector canonical correlation analysis (svcca) are widely used to analyse language model properties. @xcite find that architectural decisions, such as model width and depth, affect hidden representation similarity. @xcite show that models within the same architectural family share similar hidden structures, a similarity that per-2 this issue is termed anisotropy @xcite . ## methodology we first describe the residual stream view of transformer-based models and define layers' activations. then, we introduce the cka and proportional effective rank metrics. the residual stream view. the residual stream view of the transformer architecture @xcite is an analytical framework to study how information flows through its layers @xcite . this conceptualisation focuses on the residual connections as they provide a direct reference to the inputs. specifically, the set of residual connections across layers is termed the residual stream. each layer can be seen as providing modifications to the residual stream via addition operations. layers have two main components, attention and mlp, that sequentially update the residual stream. formally, a sequence of t tokens @xmath0, ..., t t ⟩ is first converted into a matrix x 0 ∈ r t ×d by the embedding layer: each column is a token representation of size d. then, each layer @xmath1, ..., l} updates these representations as follows: finally, the t -th column of x l is used to predict the (t + 1)-th token. more details in app. a. activations and parameters. the updates to the residual stream-underlined in eq. ( 1 )-are the layer's activations and have the same dimensions as the residual stream, i.e., r t ×d . both attention and mlp first project, or "read", the residual stream into lower-dimensional intermediate representations; then project these representations back, or "write", into the residual stream. here, we study the behaviour of the parameters that write to the residual stream. we use a att and a mlp to denote the activations and θ att and θ mlp to denote the parameters of, respectively, attention and mlp. activations' similarity. given a set of activations, either a att or a mlp , of a layer l at a particular checkpoint c, a l,c , we measure how similar they are to those at the last checkpoint c, a l,c , using the linear variant of the centred kernel alignment metric (cka; @xcite : ## experimental setup we use the pythia model suite @xcite , composed of 8 transformers of different sizes trained for 143k steps on the deduplicated 3 version of the pile dataset @xcite . intermediate checkpoints are available every 1k steps and at log-spaced intervals early in training. to comply with our computational budget, we consider models up to 2.8b parameters-i.e., 70m, 160m, 410m, 1.4b, and 2.8b-evaluated at the following steps: 0, all logspaced steps {1, 2, 4, ..., 512}, 1k, 3k, and then every 10k steps up to 143k. we evaluate each checkpoint on the last batch of the training set and collect its activations. more details in app. b. ## results our analyses reveal quantitative differences in the learning dynamics of layers across model sizes. result 1. activations of larger models converge faster and more monotonically to their final state than those of smaller models. as observed in fig. 2 (first column), larger models show, on average, earlier convergence of attention and mlp activations. for example, by 20% of training, the cka score in 2.8b is 0.8 for mlp and 0.7 for attention, where in 70m and 160m it is around 0.5. this fast convergence pattern holds across layers, as shown by the distributions in fig. 1 . result 2. activations of earlier layers converge faster, regardless of the model size. across model sizes, earlier layers' activations converge faster to their final state than those of later layers. as shown in fig. 3 (app. c), the faster average convergence in larger models is due to more of their later layers converging earlier, whereas smaller models' layers only reach their final state towards the end of training. based on recent work that identifies parameter rank differences across model sizes @xcite , in the next paragraphs, we study whether the different convergence behaviours are related to the effective rank of layers' parameters and gradients. result 3. parameters of layers in larger models proportionally span more dimensions. parameters in layers of larger models span a slightly larger fraction of their available dimensions compared to smaller models, as shown in fig. 2 (second column). moreover, the per of larger models stabilises early, while it keeps decreasing throughout training for smaller ones. this finding is further .20 .40 .60 .80 1 .90 .95 .50 .60 .70 .80 0 13k 43k 93k 143k .40 .60 .80 1 0 13k 43k 93k 143k .97 .98 .99 1 0 13k 43k 93k 143k .50 .60 .70 .80 .90 checkpoint step attention cka attention per (weight) attention per (grad) mlp cka mlp per (weight) mlp per (grad) 70m 160m 410m 1.4b 2.8b figure 2: cka similarity (current vs. last checkpoint) of layers' activations (first column), per of layers' parameters (second column) and gradients (third column) for attention (top row) and mlp (bottom row) in pythia 70m, 160m, 410m, 1.4b, and 2.8b averaged (mean) across layers per each checkpoint. parameters' @xmath2.95 by the end of training (applies to both the attention and mlp layers). we observe a strong correlation for the attention layers across model sizes. for the mlp layers, the correlation with the gradients' per is strong for models up to 1.4b, while the correlation with the parameters' per is strong only for the 70m model. we hypothesise that this discrepancy can be explained by the fact that mlp layers have a large per throughout training across all model sizes, apart from those of the 70m model. while these results are correlational, they provide a foundation for future work to test whether methods that specifically increase the per of layers' parameters and gradients induce faster convergence of the layers' activations in small models. ## conclusion our study highlights disparities in the learning dynamics of small and large lms. using the pythia model suite, we demonstrate that layers' activations in larger models converge faster and more monotonically to their final state. we correlate this phenomenon with the larger per in the parameters and gradients of larger models. our analyses expand our understanding of training inefficiencies in small models and provide insights for future work to address them, e.g., by developing methods that increase the per of layers' parameters.