Why do small language models underperform? Studying Language Model Saturation via the Softmax Bottleneck

We are truly grateful for the very positive feedback from reviewer gtfa.

Reasons to Reject

1. We agree that the causality is not perfectly established in the paper, as it is very hard to derive clear causality from such highly-parameterized models for which reproducing pretraining is rather costly and conducting ablations is rather costly. We never explicitly mention such a causality, and describe correlations wherever required in the paper. Based on your comment, we will include a causality/correlation statement in the Discussion section of our final version.
2. As some other reviewers have pointed out, it may be tough for less knowledgeable readers to understand the paper as it is in its current form. As we described in the other answers, we have worked on a revised version of the paper in which the writing has been improved for publication to a broader audience.

Questions

1. The purpose of Figure 1b is to show that if we extrapolate the expected performance of smaller Pythia models from the final performance of larger ones, we observe that models ranging from 14m to 160m parameters perform worse than expected. However, one could argue that extrapolating the larger models performance from the smaller models would lead to the opposite conclusion (i.e. large models underperform). Hence, we add another extrapolation (using the best checkpoints) to show that smaller models were closer to the performance predicted by the scaling laws mid-training, before the saturation phenomenon. We will add a sentence to make this clearer in the corresponding paragraph.
2. In our revised version, Figure 7 is after section 5.2.

We warmly thank the reviewer for its thorough feedback.

1 - The writing of the paper is quite frankly a mess

Based on the feedback of other reviewers, our understanding of the writing issue rightly identified by reviewer vM6v is that the paper is hard to follow for readers that are less knowledgeable in this specific subfield. We thus propose improvements for the camera-ready version that should make the article more inclusive:

• The representation degeneration problem : We immediately provide two references that explore this specific problem in-depth, and we proceed to explain what it consists in. We will insert an explanation between the first two sentences of the papers, to make the connection between representation degeneration and anisotropy clearer: ”It [the rep. deg. problem] consists in the emergence of degenerated structures in the intermediate latent spaces of language models throughout training. In particular, many observations…”.
• Low angular variability : We agree that readability could be improved by expanding the phrase, which we will do in the revised version: “...have shed light on their low angular variability, by showing that cosine-similarity between pairs of intermediate embeddings tend to be unexpectedly high.”
• In general, realizing that the page limit was greater than 8 (which we initially thought) let us reorganize and expand some paragraphs further. Although we cannot share the new version because of the conference guidelines, we obtained positive feedback from readers on the writing quality of this revision. The theoretical section was also reorganized and detailed in a new version of the article.

2. I would be surprised if the relevant threshold wasn't dependent on other factors

The models we experimented with in Figure 6 are: the 410m and 1B models from the Pythia suite, which were trained on The Pile with a vocabulary of 50k tokens, and the 7B Llama-2 model, which was trained on another unspecified dataset using a vocabulary of 32k tokens. We chose The Pile for the low-rank head training as it is expected to match the pretraining dataset of Llama-2 closely (and exactly for the Pythia models).

In all three cases, we see that the performance on The Pile validation data seems to noticeably deteriorate around a rank constraint of roughly 1000. Hence, we believe that the 1000 number is a good estimation of the range where the head rank starts being a significant in-domain performance bottleneck.

We are grateful to reviewer sGwj for its time and effort, and for raising interesting points about generalization.

Answering the why question in a general way

Although our paper has a strong empirical stance, we provide several elements in an attempt to prove the general nature of our observations:

• An extensive analysis of the intrinsic rank of contextual probability manifolds (Figure 7), that allows an estimation of the dimensionality of target probability distributions in several natural language scenarii.
• Rank-constrained experiments, performed on several models of different sizes and initially trained on different data (Figure 6).
• A theoretical exploration (section 5.2) that proves that the connection between the softmax bottleneck and the achievable performance affects all language models trained using cross-entropy on a linear logit language modeling head, especially in the low-dimensional case.

What would be the immediate impact of the study and analysis?

There can be several immediate impacts, among others:

• Including last-layer anisotropy monitoring when training language models to detect saturation;
• Further exploration of mitigation strategies for the softmax bottleneck issue, specifically designed for language models, with potential performance benefits;
• Incentivizing the use of larger hidden dimension choices.

We sincerely thank the reviewer for its highly valuable comments, and especially for its care about the theoretical aspects of the article.

1. The explanation does not lead to new methodology design

Our paper is framed as an extensive study of the connection between the softmax bottleneck and the representation degeneration issue. The drawn conclusions show that there is, to a certain extent, an intrinsic limitation to the performance a language model using a small hidden dimension can achieve. Designing a novel methodology to go around this limitation is definitely a path future works should explore, but as shown in the literature, it represents a difficult problem that probably could lead to a full paper on the subject if such a method was designed. We believe that describing this limitation in depth can pave the way for such methods, and represents an interesting contribution as such.

2. The theory seems not to add new insights about the phenomenon

The goal of the theoretical section is to show that the connection between both phenomena is theoretically grounded. Hence, while we agree that we get no novel insight from this section, we argue that it strengthens our argument and shows its general nature.

3. Mathematical issues

In a revised version of the paper, we adjusted and corrected some parts of the theoretical part. Notably, we removed lemma 5.1 as it is not necessary for the conclusion, we unified the notations and explained them in details from the start of the article.

Lemma 5.2 can indeed be proven by invoking the smoothness of the loss function, likely by using the $C^1$ nature of cross-entropy to argue that it is Lipschitz continuous, as regular continuity seems insufficient e.g. if the gradient in $W^*$ is undefined. We will add a remark in the proof to underline this. However, our proof relies on limited development, as we wanted to display a more precise relation between the loss gap and the error on the W approximation, and its quasi-linear nature. This formulation will hopefully pave the way for additional exploration of this relation in the future.