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

Summary

The paper find that there are strong correlation between the following things 1. Loss saturation (i.e., instability of the loss after training for a long time) 2. Performance drop in the downstream tasks 3. Anisotropy distribution in the embedding space 4. Low effective dimension in the word embedding space (i.e., problem caused by softmax bottleneck)

Rating

9

Confidence

4

Ethics flag

1

Reasons to accept

Disclaimer: I did some research in the related topics previously, so I would like to admit that my evaluation might contain some positive biases toward this research direction. Too few papers like this one that studies these fundamental and practical problems. I believe this kind of study is very important and won't be outdated in the near future. 1. Overall, the presentation is clear and easy to read for me (maybe partially due to my prior knowledge in this field). 2. The correlations are very strong. I reviewed several analysis papers for language models before and it is very rare that you can observe such strong correlations. 3. The connection between loss saturation and anisotropy distribution / softmax bottleneck is a novel discover. 4. The authors designed several good experiments and theoretical analysis to quantify the seriousness of the problems. 5. Build the connections between several previous findings or research directions quantitatively, including softmax bottleneck, loss saturation, and anisotropy representation.

Reasons to reject

I think the rejection reasons are fairly minor to me. You don't have to spend time writing the rebuttal. 1. In the title, the why question seems to suggest this paper build causality relation between observations. It would be better to emphasize in the paper that correlation is not causality. In addition, it would be even better if the authors can find whether the loss saturation causes the anisotropy representation or the anisotropy representation causes loss saturation, and why. 2. Except for the connection between loss saturation and anisotropy distribution, the rest of the empirical and theoretical analyses is somewhat novel but not super novel (although showing them quantitatively is pretty useful and the execution is done very well). 3. I am not sure how easy to understand this paper for the researchers who do not study related problems before.

Questions to authors

Small presentation suggestions: 1. The purpose of Figure 1 (b) is not very clear to me. Maybe you could explain more about what we can learn from this result. 2. When I first read Figure 7 and the definition of W-error(d), I don't understand what you tried to analyze. I understand it after reading Theorem 5.3. Thus, I suggest to move Figure 7 and its corresponding explanation to the end of section 5.2 (after your theoretical results).

Summary

The paper provides some experimental and theoretical evidence that the softmax bottleneck explains degenerate behaviour that has been observed when training LLMs with fewer parameters. In particular, it empirically shows that: 1. as the input dimension $d$ of the final classification layer shrinks, training loss tends to worsen and plateau late in training, and 2. as this happens, the average cosine similiarity between vectors in this layer suddenly spikes, and 3. singular value distributions also become suddenly high entropy (that this corresponds to 2 is not surprising), and 4. The the softmax bottleneck (where $d$ is much smaller than the number of classes $n$) theoretically relates to the observed problem with singular value distributions.

Rating

5

Confidence

3

Ethics flag

1

Reasons to accept

- I don't work on this problem, but my outsider view is that there is enough here that it will probably be of interest to those trying to figure out how to build effective compact models. I think the evidence presented here is circumstantial (there might be other ways to explain all of these things happening at once), but that seems ok from a scientific perspective; other people might be able to follow up by exploiting these observations or by coming up with better explanations for them. - The figures are quite clear, and make good use of a color gradient to illuminate relationships between variables. Most of what I understood from the paper was due to the figures.

Reasons to reject

1. The writing of the paper is quite frankly a mess. It appears to have been written in a big hurry, and some parts are difficult to follow, especially at the beginning. For example, at first I thought I would follow the paper pretty well, having understood the abstract, and having myself worked on the softmax bottleneck and on training dynamics. But I found the first paragraph (and much of the introduction) to be utterly incomprehensible, winging about obscure terms like "The representation degeneration problem" and "low angular variability" as if everyone working on LMs knows what these are. I think I followed the paper pretty well after a few pages in, after the central ideas were explained and it became clear that these terms were just needless technical mumbo-jumbo. But this is a problem, because as it stands, I think most readers won't bother reading past the first page if the paper remains in its current form. I certainly wouldn't have if I hadn't been reviewing it. The proof is written as a bunch of math with no intuitive explanation, including of how it connects to the rest of the exploration; you have to squint a bit to see the relationship, or else take it on faith. 2. The paper makes an overly strong, unqualified claim that "a linear head W substantially affects performance when rank(W) < 1000". While this seems to be empirically true _for the models examined in this paper_ (which, as I understand, are all trained on the same data, only varying the number of parameters), I have a hard time believing that 1000 is some kind of magic number that is applicable in all circumstances. I would be surprised if the relevant threshold wasn't dependent on other factors like the size of the training data and vocabulary. I think this claim should be softened to account for these unknowns. tl;dr: the paper exhibits a severe case of the Curse of Knowledge (https://en.wikipedia.org/wiki/Curse_of_knowledge). Fortunately this can be fixed with better writing. I really think a rewrite is in order before it is published.

Summary

The paper delves into a 'why' question concerning LM saturation through the softmax bottleneck. The key finding reveals that models with fewer than 1000 hidden dimensions, often representing the context, tend to adopt degenerate latent representations during late pretraining characterised by a drop after a plateau, resulting in reduced performance. The evaluation of these findings was conducted empirically, and theoretical analysis indicates that the quantification of the performance limitation is induced by a low-rank linear language modeling head.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

The paper is well written for a knowledgable reader. It is difficult to answer a `why' question and the paper showcases a good attempt to characterise a learning problem and provides empirical justifications and theoretical analysis.

Reasons to reject

It is very difficult to provide an answer/justification to `why' question with language modelling that can be generalisable. But that is the nature of the artificial neural network empirical research paradigm. We still don't know why small models worked well on some problems and did much worse on other problems. Even when it worked well on some problems, it might not have worked well for every instance of that problem.

Questions to authors

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

Summary

This paper studies relevant properties that lead to performance saturation in small language models. A particular attention is guided towards the rank of linear layer. Through experiments and some theoretical justification, the rank of the linear layer casts an impact on explaining the limited performance of small language models.

Rating

5

Confidence

3

Ethics flag

1

Reasons to accept

The study here is self-contained and explanatory with both empirical and theoretical evidences. Presentation is relatively good. As discussed in the related work, the discovery of a tie between the softmax bottleneck to the performance saturation is new to the community.

Reasons to reject

This study seems to explain a phenomenon, while the explanation does not lead to new methodology design. The discussion at the end of the paper is rather weak on the encouragement of more expressive softmax alternatives. The theory seems not to add new insights about the phenomenon and the theory itself is very standard in a simplified model. The proof is not very rigorous though, especially the usage of small o and big O notations is not well explained. Presentation can be improved by providing a clearer connection between Sections 3 - 5.

Questions to authors

Dimension of $W$ in Lemmas 5.1, 5.2 and Theorem 5.3 does not match with previous definitions. Lemma 5.1 seems to be the definition of a global minimizer. Why there is the limit of $\epsilon \to 0$ in Lemma 5.2 (resp. $d \to D$ in Theorem 5.3)? Isn't Lemma 5.2 simply the continuity of the loss function?

Thanks

I'm glad to hear that the writing of the paper has been improved. Also, thank you for the clarification about the training data and model sizes. However, I still don't think you can justify the framing of this as a constant on the basis of this handful of datapoints. Couching it as an _estimation of the range_ is an improvement over the strong claim in the paper (in the abstract no less). But even then, any such threshold is likely to be a function of other factors rather than a constant, and there isn't enough data in the paper to accurately determine what that function is, nor is it central to the main message of the paper. So I think this claim really ought to be softened and pointed to as a question for future work.

Thanks for the response

I have read the rebuttal and have no further question.