Linguistic Collapse: Neural Collapse in (Large) Language Models

Summary

This paper empirically investigates the emergence of Neural Collapse (NC) properties during the training of causal language models. NC is a phenomenon observed in the top layer of deep-nets trained on one-hot classification problems, where the last-layer class-mean embeddings become equinorm, have maximal angular separation and align with their respective last-layer classifier. However, NC only emerges when: 1) the models are trained beyond the zero training error regime, 2) the number of classes $C$ is larger than the last-layer embedding dimension $d$, 3) the classes are balanced (they have the same number of samples in the training set). Considering causal language training as the classification of contexts across $C$ words in the vocabulary set, this work explores the extent to which NC emerges in language models, given that none of the three conditions above hold for language models: 1) language models are typically trained for a few epochs, 2) the number of classes $C$ is large, 3) the words frequency in the training set is heavily imbalanced. Since the original NC geometry is not achievable, particularly without condition 2, the authors define *generalized* NC metrics to measure the geometrical properties of the last-layer embeddings and classifiers. They train NeoGPT with varying numbers of layers and width $d$, and analyze the correlation between the (generalized) NC metrics and validation loss across different training regimes.

Strengths

The problem setup is interesting as it attempts to extend previous observations in deep learning to the trending language models, despite their distinct configurations. The scope of the study, its connection to previous works, and the limitations of the setup are well-presented in the paper.

Weaknesses

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.

Questions

1. 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? 2. 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. 3. (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? 4. 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 $\mathbf{\mu}$ is uniform across classes/samples. How does this connect to generalization? 5. Do you observe a decrease in the NC metrics and validation loss as you train the models longer (e.g., 1 vs 10 epochs)?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes

Summary

This work focuses on studying the Neural Collapse phenomenon in the context of language model training. The author first introduces the original NC properties, explaining how such metrics may not apply to the case of LLM training given 1) the ambiguity of language next token prediction, 2) large number of possible tokens , 3) token imbalance and 4) under-parametrization or undertraining of LLM models. Having established this, the authors introduce new metrics based on the NC ones. Further, the authors analyze the relationship between these metrics and validation loss on a range of models trained on TinyStories dataset. In this regard, 1) collapse of token representation ($ NC_1 $), 2) Hyper-spherical uniformity which represents maximal separation ($ G-NC_2 $), 3) alignment between token feature mean and token classifier ($ U-NC_3 $) and 4) NCC accuracy of features ($ NC4 $) have positive (and stronger) correlation with validation loss.

Strengths

The connection established between Neural Collapse and LLM training is an interesting point. The introduction of the new NC metrics are not only viable for LLMs but also for other classification and training regimes that suffer from similar problems as the language domain. Further the experiments on TInyStories are performed on a large number of models to establish a connection between the metrics and validation loss and the results could be used to develop evaluation metrics for LLM training.

Weaknesses

1) **(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 suggest 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. 2) **(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-lablenss of language removed much of the uniqueness and difficulty of dealing with LLM training. 3) **(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. 4) **(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.

Questions

1) Is it true that NC properties are attributed to better generalization for classification ? Could the authors please provide some references. 2) Have the authors considered using smaller or simpler language datasets to combat the problems with large numbers of classes or have models be trained for a longer time in order to see better convergence results for the NC metrics as illustrated in the appendix ?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

None

Summary

The final layers of classifiers show a property called neural collapse (NC), which are seen as beneficial to model performance. The authors study it’s appearance in causal language models (CLMs), and point out that CLMs do not respect those conditions (CLM are trained on noisy, unbalanced data, made on more tokens than the models have dimensions, and training is stopped before loss reaches 0). They then provide a theoretical framework to adapt the NC properties to CLMs (arguing for two relaxations from past works: measuring hyperspherical uniformity, and uniform duality), and consider show that despite expectations, evidence of NC exists, and is stronger for better performing models.

Strengths

Provided measured and well formulated novel results regarding the impact of scale, training, number of parameters on the different components of generalisation. While a highly technical read which required a lot of concentration, all the required information for understanding the concepts are introduced. Conclusions are measured and entirely based on evidence from theory and empirical evidence, yet discussion manages to highlight the implication of such research. Exciting mathematical tools I hope to use in the future.

Weaknesses

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.

Questions

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?

Rating

8

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

Limitations are to my knowledge well adressed

Summary

This paper investigates neural collapse (NC) -- properties of the penultimate feature representation of DNN -- in causal language models (CLM). Previous NC studies have primarily focused on classification problems with balanced classes with few labels compared to the feature dimensionality. This paper finds that NC properties in CLM are correlated with generalisation. This finding is based on extensive experimental evaluations on the TinyStories dataset.

Strengths

* The NC analysis is interesting and it's noteworthy that properties of NC appear to correlate with generalisation in CLMs, which differ in important ways from other classification settings where NC has been studied. * The empirical results seem to paint a clear trend in terms of NC and generalisation. * The presentation is generally quite good, with contributions and significance clearly spelled out.

Weaknesses

* 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. * 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. * 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. * 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.

Questions

See "Weaknesses."

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes.