Linguistic Collapse: Neural Collapse in (Large) Language Models

I don’t understand what you mean by ambiguous samples, ``which are not soft- or multi-label’’ (line 128). When a context appears several times in the training set, each time followed by a different next token, this training sample has a soft label, where the label (next token) can take on different values with some non-zero probability for each.

The reviewer is correct that such ambiguous contexts do resemble soft-label data examples, and we will edit the relevant paragraph in the Related Works to reflect this.

It is not clear to me whether the suggested NC metrics are suitable for the language setup. For instance, since in a language dataset, a given fixed context might be followed by different next tokens, a model that is accurate in predicting the labels on the training set cannot achieve near-zero NC1. Or regarding UNC3, see question 4 below.

We agree that such ambiguous (soft-label) samples such as “Once upon a time ___.” do exist. However, we note that such contexts are more often shorter subsequences typically found in the beginnings of sequences. In longer sequences, the probability of having two identical contexts followed by a different word is diminishingly small. Additionally, longer contexts significantly outnumber the shorter sequences in the training dataset, especially when data and models work with ever longer contexts; this is progressively true in models with increasingly long context windows, such as Llama 3.1 (https://ai.meta.com/research/publications/the-llama-3-herd-of-models/). In the abstract, shorter (ambiguous) contexts could even be excluded or relegated as outliers in the data. Therefore, we conclude that it is likely not an issue overall.

On the other hand, we believe it would still be interesting to extend our work to incorporate soft labels and multi-label approaches. This could be explored using recent advancements in multi-label neural collapse (https://arxiv.org/abs/2310.15903) and mixup neural collapse (https://arxiv.org/abs/2402.06171). We hope that our work serves as a strong first step in modeling NC at the next-token level.

Is your setup only specific to causal language modeling? Particularly, a) Is any part of the formulation or results influenced by the causal/autoregressive training, or is it simply related to the different nature of language datasets in general, such as the large number of classes and the possibility of having multiple labels for a given training sample?, b) Do you expect similar behavior if the experiments were conducted using other language training methods, such as masked language modeling?

We appreciate the questions. We don’t believe our setup is strictly specific to causal language modeling.

1. The adverse conditions we list in our Introduction are mostly due to the imbalanced nature of tokens in natural language data. Intuitively, they would apply beyond autoregressive modeling.
2. We expect to see similar results in masked language modeling. However, it may depend on which tokens are frequently masked in the data (if a specific scheme is used). We hypothesize that one could conduct experiments with similar results on bidirectional encoders such as BERT or T5 (or their derivatives) but we imagine they’d present additional challenges such as lower token sample efficiency of the MLM paradigm, or the simultaneous prediction or fill-in of multiple dependent/correlated token predictions.

(Related to the previous question) What is classification error (line 130) in the context of next-token prediction and what do you mean by "irreducible noise’’? How do you define the error for contexts that appear several times with different next tokens/labels?

The loss is token cross-entropy and the error we use is the average mis-classification rate over all tokens. We do not treat ambiguous samples differently and instead defer to the diminishing proportion argument previously made regarding soft-labels. Irreducible noise refers to some minimum loss and error that can't be reduced by more scaling simply because of soft-label contexts.

Intuitively, why do you expect that a model minimizing UNC3 has better generalization? Minimizing NC3 means the classifiers and mean vectors are aligned which is consistent with the classification objective. However, minimizing UNC3 (CoV of the term in Eq (8)), only implies the degree of misalignment between the classifiers w and mean embeddings μ is uniform across classes/samples. How does this connect to generalization?

Our initial manuscript didn't include our reasoning behind UNC3 and we neglected to provide it in our Appendix, so we thank the reviewer for raising this oversight.

The goal of (self-)duality is to minimize the angles between mean vectors and their corresponding classifiers. This can be measured through the expectation (over class or token pairs) of the squared angle between vectors: $\mathbb E[\theta^2]$.

$\mathbb E[\theta^2] = \mathbb E[\theta]^2 + \text{Var}[\theta]$

NC3 measures the average angle $\mathbb E[(\theta)]$, whereas our newly introduced UNC3 measures the variances in the angles $\text{Var}[\theta]$. Achieving duality ultimately requires minimizing both terms.

Once again, we appreciate this comment and will add a short discussion on duality decomposition in $\S$ 3.6 and $\S$ 4.6.

Do you observe a decrease in the NC metrics and validation loss as you train the models longer (e.g., 1 vs 10 epochs)?

Yes. Appendices D, E, F, G, H, I, J, L, and M all show trends towards NC across training (left-to-right) in sufficiently large models. We show more explicit scatter plots (red = more parameters) in the 1-page PDF supplied in the general rebuttal. That PDF also includes validation loss with respect to training. We will also add these figures to the Appendix of our main manuscript.

Overall, we appreciate the constructive comments of the reviewer. Should our response address the concerns raised, would you consider raising the score? Thank you.

(Major) I appreciate the extensive experimental work on establishing a connection between the proposed NC metrics and validation accuracy; however, I believe the coefficient of determinations provided in Table 1 suggests a low correlation between NC properties and validation loss. While I see the higher correlation for NC1 and NC4 in Fig 1, I’m having a hard time convincing myself that other factors such as model size and architecture don’t play a role in improving the loss.

To clarify, we don't claim that size and training don't contribute to reducing the loss. Our principle observation is that scaling improves both performance (as expected) and NC, and that the two are associated (shown in Figure 1).

To address the concern that scaling alone is the confounding factor, we conducted a seed sweep ($\S$ 4.5) to test an independent correlation. Table 1 shows that most correlations between NC and generalization are statistically significant. Equinormness is not correlated but only accounts for some of NC, a property that is superseded by GNC2 anyway. The only real exception is therefore UNC3, which seems to be entirely confounded by model scale. (Although this does highlight the effect of NC3 on performance independent of scale.)

(Major) While I understand the attempt at making the autoregressive next token prediction of LLMs comparable to a simple classification setup, I do not think I agree with the “not soft label, not multi label” argument. As the authors suggest, the language format is rather ambiguous. Considering a simple next word prediction, “I went to the ___ “ can be followed by a variety of valid words. Same applies to token baked predictions. I believe not having a mechanism to account for this and dismissing the multi-lableness or soft-lableness of language removed much of the uniqueness and difficulty of dealing with LLM training.

Thanks for the detailed comment. Firstly, we agree with the reviewer that there are indeed ambiguous samples, particularly with short sentences, such as the example “I went to the ___.” These cases can resemble soft label data examples, and we will edit the relevant paragraph in the Related Works to reflect this.

We note however that ambiguous contexts are typically shorter and found at the beginning of sequences. For longer sequences, the probability of two identical contexts followed by a different word is low. Additionally, longer contexts outnumber shorter ones in the training data, especially as LLMs (such as Llama 3.1) handle longer contexts. Shorter contexts could even be excluded or treated as outliers. Thus, we conclude that these anomalies likely represent a minority of the data and aren't a significant issue overall. Nonetheless, we take the first step in NC at the next-token level, and future work might incorporate multi-label (arXiv:2310.15903) and mixup (arXiv:2402.06171) NC.

(Major) Considering how languages themselves have an imbalanced nature, I wonder whether the suggested ETF of hypersphere features and weights geometric structure are actually optimal for LLMs. As an example, in the case where for a CIFAR10 classifier, we train the model under artificial imbalance, I would understand why symmetric features could potentially help improve test accuracy given the balanced nature of the validation set. However, I have a hard time being convinced the same is true in language.

Our primary objective is to investigate the structures that form in standard LLM training. We recognize that the convergence geometry in current LLMs may not be optimal, likely due to using CE loss (adapted for balanced data) on an imbalanced dataset. We don't not claim the simplex ETF or hypersphere are optimal; rather, we focus on characterizing the patterns that emerge at scale in relation to baseline geometries. However, recent works have proposed methods for addressing data imbalance (arXiv:2301.00437v5, arXiv:2301.01100, and MLR: Liu et al. (2023)). We hope our insights may lead to adaptations of such methods to optimize these structures that address the imbalanced nature of language

(Minor) I find it hard at times to follow the paper’s ideas. Particularly with regards to section 4, I believe a bit more structure or separation of ideas could help readers better digest the conclusions.

We sympathize with this sentiment. We'll move 4.5 through 4.7 (including Table 1) into a separate section.

Is it true that NC properties are attributed to better generalization for classification? Could the authors please provide some references?

The original NC showed some relationship to generalization. In particular, we highlight the results in transfer learning (arXiv: 2112.15121) and a generalization bound that can estimate test performance (arXiv: 2202.09028).

We aren't aware of other works that explicitly studied the relationship between NC and generalization. However, our results appear to reinforce the generalization narrative as we observe correlations long before the terminal phase of training, some of which are independent of scale. We hope to inspire future works that might definitively answer this open question.

Have the authors considered using smaller or simpler language real datasets to combat the problems with large numbers of classes or have models be trained for a longer time to see better convergence results for the NC metrics as illustrated in the appendix?

Thank you for your insightful suggestion. We considered using smaller or simpler language data to potentially improve convergence. However, we focus on large-scale language modeling scenarios to ensure our findings are robust and reflective of contemporary practices. This would render our results applicable and generalizable to real-world applications. The synthetic nature of TinyStories is a caveat to our study (as the reviewer implies), but we provide a lengthy explanation for its use over real data in the general rebuttal above.

I’m somewhat unclear on the motivation for using TinyStories, a purely synthetic dataset. While I understand the benefit of using a small dataset to ease the computational burden of the experiments, I would have liked to see some further experiments on real data confirming the observed trends. For example, you could sample a small dataset of human-written stories.

This crucial question arose during our research, as using a synthetic dataset may limit the generality of our results to human-written data (see our lengthy explanation in the general rebuttal). We acknowledge the importance of validating our findings with real data and are aware of follow-up efforts to scale our experiments with appropriately sized models and datasets.

I would have liked to see more practical guidance on the value of the proposed metrics. For example, it’s not clear if they provide any value in terms of assessing model fit beyond what a held-out validation would (more cheaply) provide.

NC has been studied in and in relation to many areas of machine learning. We defer the vast related works and applications to the Related Works and Discussion sections of our manuscript, while our response here will therefore be tailored towards practical consequences and applications of studying NC in CLMs.

1. Better optimization strategies: LLMs are typically pre-trained with CE loss in teacher-forcing. It stands to reason that works listed in the “Learning to collapse” Discussion paragraph can be applied or inspire better objective functions towards potentially better geometric configurations for language models. Hierarchical structures like those presented by Liang & Davis (2023). The implication is that future LLMs could be trained in fewer steps or to better performance.
2. Interpretability of contextual token-wise behaviors: our measurements provide insight into behaviors of individual token embeddings. Token-wise NC3 can reveal how confident an LLM is about certain tokens in their top-layer classifiers. For example, the number and density of clusters for a particular token can shed light into its various meanings and uses. This would be particularly useful as LLMs adapt to ever-evolving use of language and further expansion into non-English domains.
3. Interpretability of pair-wise interactions: our measurements primarily center around pair-wise relationships between tokens in a vocabulary. The pair-wise arrangements are critically important to modeling interference between tokens because their context vectors live on a lower-dimensional hypersphere. For example, you can tell how related or interchangeable two words are based on their noise and interference (NC1,2), or how antithetical or unrelated they are based on orthogonality (NC2).
4. Interpretability for fairness: natural language is inherently imbalanced. Individual or pair-wise analyses of representations can reveal inequalities within the vocabulary or across topics. These insights can guide researchers and engineers to address inequities in a targeted and measurable way.
5. Interpretability for ethics: if there are concerns about biases or safety issues in general language modeling, NC analysis can aid researchers in interpreting how LLMs could pose risks to users, and how to safely mitigate risks without degrading model quality.
6. Continual learning: [19, 20, 21] cited studied NC in (task-)incremental learning problems and provided insights or improvements. Since an LLM (especially a foundation model) is to continually learn language patterns and develop capabilities, NC can provide carefully structured configurations that allow for graceful lifelong learning.

Of course it remains to be seen in exactly what form or role NC will play in these areas of further studies. It is also probable there are benefits of NC analysis beyond our knowledge. Our work here is to take the first step in extending the body of NC work into the more irregular and demanding settings of causal language modeling.

We are grateful to the reviewer for probing this response. We will further expand our Significance and Discussion (sub)sections to explicate these practical guidelines.

In several places, the wording is unusual / confusing. For example, L47-L48 says that "[...] LLM learn to model aleatoric uncertainty, which can be viewed as stochastic token prediction," which is a peculiar/wordy way to describe a standard MLE procedure.

We agree with the reviewer and will revise L47-48 to “[…] learn to stochastically predict tokens to generate text.” We also identified the following opportunities for rewording or clarification:

1. L107-108 will be adjusted: “[…] classes number in the tens of thousands.”
2. L128 will be clarified such that ambiguous contexts are “soft-label” samples.
3. L230-231 will be rewritten: “These noise reductions are associated with generalization (Fig. 1, left, “N C1”); this relationship grows stronger with model size.”

Should the reviewer feel there are more instances that we did not cover, we will happily review further during the discussion period. And of course, we will perform more passes on this paper to improve the writing for future revisions.

While the experiments are fairly comprehensive, the technical novelty is relatively low. The paper mostly consists of methods from prior works being applied to a new classification setting.

This is a fair point, that we mostly applied previous expressions and methods to the causal language modeling setting. Our contribution focuses on applying the canonical NC framework to more adverse conditions (larger and imbalanced vocabulary, ambiguous contexts, undertraining), the effort of which is reflected in our work. We hope that this first attempt at analyzing NC in this area inspires further work that better adapts to the ever-changing landscape.

We are grateful for the helpful comments. Should our response warrant it, we would greatly appreciate the reviewer raising the score.

Apart from evaluation loss, model performance metrics from commonly used LLM benchmarks are not provided, making high level comprehension of the technical observations more complicated.

LLM researchers are indeed interested in performance metrics beyond the cross-entropy (CE) evaluation loss that we use. However, as our work concerns the measurement of NC under model pre-training only, other benchmarks (particularly those on downstream tasks) would not necessarily be appropriate as they measure specific capabilities rather than generic stochastic token prediction; therefore, they would be out-of-scope for our work.

Furthermore, according to several recent works on LLM evaluations (https://arxiv.org/abs/2403.15796, https://arxiv.org/abs/2404.09937, https://arxiv.org/abs/2407.06645), downstream capabilities of LLMs are roughly correlated with their abilities to compress their pre-training data. Based on these findings, we find that CE loss is the most sensible metric with which to measure generalization.

We thank the reviewer for raising this question as we see fit to include a short discussion on this and potential future directions towards the end of our manuscript (citing the aforementioned works).

Considering CLMs to be classifiers is technically sound, nonetheless unlike with classification tokens can mean different things (ex:homonyms) these might also explain divergence from the classic NC model. Did you observe any such groupings or effects?

Appreciate the question! We inspected token-wise NC values in our largest and most-trained model (12-layer d=1024, 10 epochs). Based on your suggestion, we chose 15 homonyms and found that most have much shorter mean vector norms (meaning they’re closer to the global center) than the average token. This makes sense as homonyms present conflicts and interference.

We also observed that the individual variability and interference of some English first names all other tokens were far below the average. This is also intuitive as names are distinct from one another and aren't typically used in the same contexts as other words (aside from articles).

There are probably thousands of groupings or linguistic relationships one could observe in the NC measurements, so we’ll leave such interpretability to future applications.

We thank the reviewer for the thought-provoking questions. We hope our responses adequately address them and warrant a slightly higher score.