Linguistic Collapse: Neural Collapse in (Large) Language Models

Authors: robert  wu
vardan  papyan
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:10:47

Abstract: Neural collapse ($\mathcal{NC}$) is a phenomenon observed in classification tasks where top-layer representations collapse into their class means, which become equinorm, equiangular and aligned with the classifiers.
These behaviors -- associated with generalization and robustness -- would manifest under specific conditions: models are trained towards zero loss, with noise-free labels belonging to balanced classes, which do not outnumber the model's hidden dimension.
Recent studies have explored $\mathcal{NC}$ in the absence of one or more of these conditions to extend and capitalize on the associated benefits of ideal geometries.
Language modeling presents a curious frontier, as \textit{training by token prediction} constitutes a classification task where none of the conditions exist: the vocabulary is imbalanced and exceeds the embedding dimension; different tokens might correspond to similar contextual embeddings; and large language models (LLMs) in particular are typically only trained for a few epochs.
This paper empirically investigates the impact of scaling the architectures and training of causal language models (CLMs) on their progression towards $\mathcal{NC}$.
We find that $\mathcal{NC}$ properties that develop with scale (and regularization) are linked to generalization.
Moreover, there is evidence of some relationship between $\mathcal{NC}$ and generalization independent of scale.
Our work thereby underscores the generality of $\mathcal{NC}$ as it extends to the novel and more challenging setting of language modeling.
Downstream, we seek to inspire further research on the phenomenon to deepen our understanding of LLMs -- and neural networks at large -- and improve existing architectures based on $\mathcal{NC}$-related properties.
Our code is hosted on GitHub: [`https://github.com/rhubarbwu/linguistic-collapse`](https://github.com/rhubarbwu/linguistic-collapse).

Main paper: 
 introduction 1.neural collapse: 
 A learning phenomenon known as neural collapse (N C) emerges during the terminal phase of training (TPT) deep neural networks with cross-entropy (CE) loss for classification tasks. [1] It was originally characterized as the co-occurrence of the following properties in a model's top-layer (also known as last-layer) representations (also known as features or embeddings) and linear classifier weights:(N C 1 ) Within-class variability collapse: Top-layer representations collapse to their class means.(N C 2 ) Convergence to a simplex ETF: Class means tend towards equinorm and equiangular vectors when centered about the global average. The resulting geometry -known as a simplex equiangular tight frame (ETF) -maximizes pairwise angles and distances.(N C 3 ) Convergence to self-duality: Linear classifier weight vectors converge to their corresponding top-layer class mean vectors, and thus also form a simplex ETF.(N C 4 ) Nearest decision rule: Linear classifiers approximate a nearest-class center (NCC) classifier: top-layer representations predict the class with the closest mean (implied by N C1-3).These behaviors, often associated with improved generalization and robustness [3] [4] [5] (among other benefits, such as those discussed in §1.4), traditionally manifest under the following conditions:Rq1) Few classes: The number of classes is upper-bounded by the embedding dimension plus one: C ≤ d + 1; this is required to construct a perfect simplex ETF.Rq2) Balanced classes: The number of samples is equal across classes:N c = N c ′ , ∀c ̸ = c ′ .Rq3) Noise-free labels: Identical (or very similar) embeddings should belong to the same class.The model is trained past zero error, towards zero loss.Absent these conditions, one does not typically anticipate N C. Since then, follow-up studies have extended beyond and proposed techniques of quantifying or achieving N C (discussed in Section 2).N C is a phenomenon observed specifically in classification tasks. While not traditionally thought of as classifiers, language models -including large language models (LLMs) -learn to model aleatoric uncertainty, which can be viewed as stochastic token prediction [6] . For instance, masked language models (MLMs) such as BERT [7] predict one or several masked tokens within an input sequence based on the surrounding context. Likewise, autoregressive or causal language models (CLMs) such as GPT [8] predict the next token in a sequence given the context of previous tokens. Most of these models are essentially (pre-)trained by token classification on their vocabularies. This parallel -also drawn by [9] -raises a few natural questions:1. Does the pre-training stage of a language model exhibit N C? To address these, we first examine the specific settings of training CLMs as they are opposed (¬) to the traditional prerequisites (R1-4, §1.1) for N C to manifest.¬Rq1) Many classes: The unique tokens found in language modeling vocabularies are vast, usually numbering in the tens of thousands and far exceeding the hidden dimension [10] .¬Rq2) Imbalanced classes: The distribution of tokens in natural language is typically very imbalanced [11, 12] , as is the case in TinyStories [2] , the dataset we use (Appendix Figure 4 ). It has an average of 16K samples per class but a standard deviation of over 32K.Ambiguous contexts: There may exist very similar or even identical contexts followed by different tokens in the natural language data [13] . For instance, over half of the sequences in TinyStories [2] lead with "Once upon a time", but only three-quarters of these follow with a comma ("Once upon a time,"). In other words, there is almost certainly some ambiguity to contend with in our context embeddings. ¬Rq4) Undertraining: Most practical language models (including LLMs) are not trained for more than a few epochs [14, 15] . Further optimization typically renders diminishing returns in improving evaluation performance [16] long before any possible TPT.We train a suite of Transformer-based [17] CLMs1 across a grid of model widths, depths, and training epochs on the TinyStories dataset [2] to assess the degrees to which N C properties develop and how they relate to generalization performance. Our findings (summarized in Figure 1 ) reveal:• Emergence of N C with scale: As we scale model size and training, the properties of N C emerge; within-class variability (N C 1 ) and interference are reduced while hyperpsherical uniformity (GN C 2 ), uniform duality (UN C 3 ), and classifier agreement (N C 4 ) improve. • Progression towards hyperspherical uniformity: Class means, while unable to form a simplex ETF (N C 2 ), nonetheless tend towards uniform dispersion on a hypersphere, a geometry theorized by [18] and formalized by [19] as hyperspherical uniformity (GN C2). • Tendency towards uniform duality: Direct alignment (self-duality) between class means and classifiers (N C 3 ) does not appear to develop with scale. Instead, the variability of (mis)alignment across classes decreases with width and training, suggesting its minimization -which we term uniform duality (UN C 3 ) -may be more cohesive with N C. • Correlation between N C and generalization: The developments of N C properties are correlated with improved validation performance. We show these correlations to persist even when fixing the (potential confounders of) model architecture and training by simply varying the random seeds for initialization and data shuffling. This suggests that N C is not simply a side-effect of training, but possibly a factor of generalization in its own right.Recently, methods building on N C have found use in deep learning at large. We highlight areas such as federated learning [20] , graph neural networks [21] , incremental/continual learning [22] [23] [24] , metalearning [24, 25] , out-of-distribution detection [26] [27] [28] , privacy [29, 30] , learning theory [31] [32] [33] [34] [35] [36] transfer learning [3, 5, [37] [38] [39] [40] , and even LLM prompting [41] . With our results, we aim to extend insights from such contributions and other related works to the autoregressive language modeling domain and ultimately assist researchers in improving and interpreting their (L)LMs. 
related works: 
 N C was initially observed in image classification tasks such as on CIFAR-10/100 [42] and ImageNet [43] . Since then, the phenomenon has been further studied both theoretically and empirically [5, 18, 35, , with several works venturing into settings without some of the traditional prerequisites (¬Rq1-4, §1.2) and proposing adaptations of the analysis framework or optimization procedures:A large number of classes (¬Rq1) N C established the simplex ETF as an optimal configuration. However, a perfect simplex ETF (N C 2 ) requires that the number of classes C not exceed d + 1 where d is the embedding dimension. This requirement that d be sufficiently large is not practical when the classes number beyond the thousands. 2 For instance, GPT-2 [68] and Llama 3.1 [69] have vocabularies of around 50K and 128K tokens, respectively.In such a scenario, one might still expect class means to be uniformly distributed on a d-dimensional hypersphere [18] . [19] formalize this as hyperspherical uniformity within a generalized neural collapse (GN C) framework, which [9] then empirically demonstrate. These latter two works mention language modeling as applicable for GN C; [9] even framed token prediction as a classification task, just as we do. We remark however that both earlier works simulate a large number of classes by drastically shrinking the embedding dimension. In contrast, we study next-token prediction, using the full class space (vocabulary) with an imbalanced token distribution and ambiguous samples.Class imbalance (¬Rq2) N C traditionally assumed that classes were sample-balanced. Since then, follow-up works have investigated the effect of class imbalance on the prevalence of N C. [47] studied the layer-peeled model (LPM) and discovered that minority collapse occurs when classes are imbalanced across two equally-sized groups of classes; a threshold for minority collapse was later characterized by [35] . [52] showed that N C still occurs in such an LPM when the classifier is initialized as a fixed ETF. [70] introduced simplex-encoded-label interpolation (SELI) but noted that more severe imbalance worsens even this geometry. Recently, feature regularization has been employed to induce N C and improve generalization in class-imbalanced settings [58, 61, 62] .In some problems, one might encounter mixed or multi-label samples, be they natural or due to noise or augmentation [71, 72] . N C was also recently studied for such data by [73] , who observed that multi-label class means arrive at the average of their labels' respective means. They also devise an augmented CE loss function to accommodate such samples.Likewise, most of our ambiguous samples could be considered multi-or soft-label: identical (or very similar) context samples with different hard token labels (¬Rq3). Under popular CLM pre-training (teacher-forcing with CE loss), this effectively precludes the prospect of achieving zero classification error and potentially introduces irreducible noise.A recent study showed that memorization of noisy labels likely leads to degradation of N C [60] .Early stages of training (¬Rq4) [53] studied N C in small ResNet [74] models that had not yet converged (similar to most LLMs). They show that within-class variability drops and "plateaus" (N C 1 ) earlier than other N C metrics, a result that we also observe ( §4.1, Figures 6, 7 ).Natural language processing (NLP) An earlier study reported that the ratio of within-class to between-class covariances of word embeddings increases with model depth [75, 76] , seemingly at odds with N C 1 . It can however be distinguished from literature studying layer-wise N C in that it does not center the mean representation vectors (i.e. subtract the global mean).[19] fine-tuned BERT [7] using their hyperspherical uniformity gap (HUG) objective on binary classification tasks from the GLUE benchmark [77] . [78] conducted a tangentially-related investigation of convolutional neural networks for few-class semi-supervised clustering in which they identify N C. Our work is distinct from these in several ways: a) our class space far exceeds our embedding dimension (C ≫ d + 1) because we classify on the full token vocabulary; b) we analyze embeddings at a token-level granularity rather than at the sequence level; and c) our next token prediction is causal (context-dependent) as opposed to the per-sample independence of their few-category classification.A more related work studied feature collapse in individual word representations [79] , but the authors note that their analysis is limited to shallow NLP modeling on more rigid and tabular text data.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We discuss the limitations of the experiments themselves (e.g. small models, limited sample size) as well as the possibility that our formulations could be improved in the future (there could be more suitable metrics). See Sections 6, 7.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methodology: 
 Below we describe the training setup for our CLMs ( §3.1), procedures 3 for collecting top-layer embeddings ( §3.2, 3.7), and measurements of N C and generalization ( §3.4, 3.5, 3.6, 3.7, 3.8).TinyStories [2] is a synthetic 4 dataset generated by GPT-3.5 and GPT-4 using around 1500 English words a child might use. Next-token prediction is performed by sampling from the token vocabulary V = 1, 29233 , which for our purposes can therefore be framed as classification among C = 29,233 classes. 5 Following the GPT-style preprocessing regime [8] , raw sequences are packed into S chunks of size T , providing N = S(T -1) token samples for training. 6 Details are listed in Appendix A.We use 30 CLM architectures based on GPT Neo [80] , configured similarly to [2] . They vary in width (embedding dimension) d ∈ {64, 128, 256, 512, 768, 1024} and depth (number of self-attention layers) L ∈ {1, 2, 4, 8, 12}. Our models were trained by teacher-forcing 7 using CE loss. For each architecture, we trained multiple models for 1, 3, and 10 epochs ablated over weight decay factors β = 0.0005 [51] and β = 0.1 [81] . Further details are listed in Appendices B, C.Base CLMs perform next-token prediction: given a sequence of tokens x 1:t ∈ V t , a top-layer context embedding h(x 1:t ) ∈ R d is used to predict the next token x ′ t+1 ∈ V where 1 ≤ t ≤ T . A classifier for class c with weights w c and bias 8 b c would make maximum a prior (MAP) predictions asEQUATIONClass embedding means For each token class c, we are interested in the mean embedding µ c ∈ R d across sequences s and their contexts x t+1 is ground-truth (t < T ) and equal to c. To centre these means, we compute their unweighted 9 average μ ∈ R d . Precisely, we takeEQUATIONwhere N c is the number of samples of class c and I is the (binary) indicator function.Class embedding variances In a second pass, we accumulate the sample variance norms: 10EQUATIONFor some collapse measures (such as (G)N C 2 and N C 3 ), we are primarily interested in the variation rather than the average of pairwise relationships. To that end, we also include in our analysis the coefficient of variation (CoV) of several measurements, which is the ratio of their standard deviations to their means: σ(•)/µ(•). We can interpret this as a normalized measure of variability.The ability to disambiguate between classes depends on the ratio of within-class to between-class variabilities. Building upon foundational works [85, 86] , N C originally measured variability through an inverse signal-to-noise ratio (SNR), whose minimization constitutes within-class variability collapse (N C 1 ). We instead employ a class-distance normalized variance (CDNV) similar to [3] :EQUATIONOur metric differs in that we divide by an additional power ∥µ c -µ c ′ ∥ k 2 of the mean distance norm. This further downweights the CDNV within well-separated class pairs in favour of emphasizing more 5 Although the GPT-Neo [80] tokenizer has over 50K tokens, only a subset vocabulary appears in TinyStories. 6 We cannot use the first ground-truth nor last predicted token in any chunk. 7 Parallelized training using sequences' ground-truth labels for context as opposed to predicted tokens. 8 Similar to many causal LLMs [80] [81] [82] [83] [84] , our classifiers do not include additive biases, so bc = 0. 9 Different from most literature where classes were balanced and weighting is already equal. 10 This sample variance is computed across all unnormalized dimensional entries. mutually noisy pairs. We found this augmented CDNV with k = 2 to be especially useful in our setting of many imbalanced and variably confusable token classes. These pairwise measures constitute the off-diagonal 11 entries of a symmetric matrix in R C×C , whose average we use as an inverse SNR. Within-class variability collapse is then re-characterized by the minimization of this quantity: σc,c ′ → 0, ∀c ̸ = c ′ . This alternative convergence is empirically faithful to N C 1 but more robust and numerically stable [3] .The separability of our representations also depends on the geometric structures found in our embeddings. [1] characterize N C 2 as convergence to a simplex equiangular tight frame (ETF) [87, 88] .Equinormness Such a near-orthonormal configuration firstly implies that class means are equinorm:log ∥µ c -μ∥ 2 -log ∥µ c ′ -μ∥ 2 → 0, ∀c ̸ = c ′ .(5) We measure CoV in the logarithms of class mean norms to assess "equinormness".Equiangularity N C 2 also entails that class means are equiangular about their center μ: pairwise distances and angles between their class means should be maximized and similar. Following [1] , we measure interference (sometimes known as similarity or coherence [89, 90] ). Its minimization,EQUATIONtogether with equinormness (Equation 5) constitute convergence to a simplex ETF. Although this geometry is not ultimately attainable since there are too many classes (C > d + 1), it can still be meaningful to measure a model's tendency towards one. As with CDNV noise (Equation 4), pairwise interferences form off-diagonal 12 entries in a symmetric matrix in R C×C . The minimization of CoV in interferences therefore expresses the degree of "equiangularity".Hyperspherical uniformity A relaxation from the ETF is hyperspherical uniformity (GN C 2 ), with mean vectors µ c uniformly distributed on the d-dimensional hypersphere [18, 19] . We gauge the degree of this uniformity with pairwise interactions under a logarithmic inverse distance kernel: 13EQUATION)The linear classifiers {w c } C c=1 lie in a dual vector space to that of the class means {µ c } C c=1 . While convergence to self-duality (N C 3 ) was initially measured as distances [1] between class means and classifiers (Equation 11), we follow [5] to inspect class-wise cosine similarities: 14 wEQUATIONFor intuition analogous to that for equinormness and equiangularity ( §3.5), we also measure uniform duality (UN C 3 ) as the minimization of the CoV of these similarities (Appendices N, O).Finally, N C 4 is described as the simplification of the linear classifier's MAP predictions (Equation 1) for test 15 points to those of the implicit nearest-class center (NCC) classifier:EQUATION)11 The main diagonal of CDNVs is undefined (due to the zero denominator) and ignored. 12 The main diagonal of interferences is equal to 1 (maximal coherence or self-similarity). 13 Following [19] , we employ the logarithmic inverse distance kernel K log (a, b) = log ∥a -b∥ -1 2 for its ability to emphasize gaps between small distances while scaling down the effect of larger distances. 14 Dot-product would be confounded by norms and therefore inappropriate for similarity up to rescaling. 15 We used the validation set since it was not used for training or hyperparameter optimization.We calculate 16 agreement as the proportion of samples on which the classifiers agree: 17Tval-1 t=1EQUATION)To isolate the effect of N C on generalization independent of model scaling and training (if it exists), we selected a two-layer 768-wide architecture of which to train twenty more instances with weight decay β = 0.0005, each with a different data shuffling seed. We then followed the remainder of the pipeline described above to collect and analyze embeddings with respect to N C. Finally, we performed a permutation test with 10 4 trials to determine the statistical significance of any correlation between N C and generalization that remains when we hold constant all factors but shuffling seeds.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: The equations and equivalent decompositions used are all clear and numbered in the paper. Their sources are all provided as references. We do not introduce new proofs or non-trivial expressions. See Section 3.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experimental results: 
 In this section, we present the results from our empirical study on scaling and generalization:(N C 1 ) Within-class variability is reduced across model scale (more so by width than depth) and training (up to 6 epochs), and is tightly correlated with validation performance. (N C 2 ) Equinormness/equiangularity shows some improvement with scale, training, and performance. Hyperspherical uniformity (GN C 2 ) also improves but more clearly and consistently. (N C 3 ) Our models fail to achieve self-duality: class means do not align with classifiers. However, uniform duality (UN C3) is correlated with model width, training, and performance. (N C 4 ) Larger or more trained models exhibit closer agreement between their linear and implicit NCC classifiers. Agreement is also associated with validation performance.Scaling our models dramatically reduces normalized variance, which is further aided by more training epochs and stronger weight decay (Appendix Figs. 6, 7 ). These noise reductions tightly associate with generalization (Fig. 1 , left, "N C 1 "). The relationship is most apparent at larger model scale. Equinormness Class mean embedding norms grow with model width and training (Appendix Fig. 8 ), and subtly with depth (Appendix Fig. 9 ). Despite this growth, the variation of norms consistently decreases (improving equinormness) with scale (Appendix Figs. 10, 11) . Both trends correlate with improved generalization (Appendix Fig. 20 ).Equiangularity Scaling model dimensions reduces average interference (Appendix Figs. 12, 13) down to an empirical plateau of approximately 10 -3 , which is more apparent in less-trained models. However, the variation of interference rises and persists when scaling (Appendix Figs. 14, 15) , suggesting that interference becomes more concentrated between some pairs of classes. These results could be due to various factors, including but not limited to unfriendly conditions of language modeling ( §1.2) or the impossibility of a simplex ETF ( §3.5).We only see a rough correlation between average interference and generalization and almost none with its variation, equiangularity (Appendix Fig. 21 ). In other words, the limited trends we observed toward a simplex ETF do not appear to be associated with generalization when C > d + 1.Hyperspherical uniformity Logarithmic distances drop more gradually and consistently with scale (Appendix Figs. 16, 17) , implying this quantity is more robust or may not face the same geometric barriers seen in conventional interference (Appendix Figs. 14, 15) . Variation of these logarithmic distances is also consistently reduced with scale (Appendix Fig. 18, 19 ).And finally, generalization has much stronger correlations with logarithmic distances than it has with regular interference (Fig. 2 ), validating the applicability of GN C [19] when C > d + 1.Model width (d) is correlated faintly with average similarity between class means and their respective classifiers (Appendix Fig. 23 ), but strongly with variation in similarity (Appendix Fig. 25 ). The relationships to generalization follow the same pattern (Fig. 3 ), suggesting that uniform duality (UN C 3 ) might serve as a better N C property than self-duality (N C 3 ) overall; we discuss this in §5.1. The linear and NCC classifiers agree on far more examples than random chance. Scaling encourages agreement (Appendix Figs. 29, 30) . Increasing width for certain depths happens to plateau or even regress the agreement rate, but this limitation is overcome with further training. And finally, agreement is a strong indicator of generalization (Fig. 1 , right, N C 4 ).Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: Our methods for training the models, collecting embeddings, and analyzing metrics are clearly described in the paper, with supporting details in Appendices A, B, C. Links to code are provided in the Abstract and Section 3. 
analysis: 
 We find that N C is generally promoted by model size and training and correlated with generalization (validation performance). We also discern some of this correlation independent of scale ( §5.1). Table 1 presents the correlation scores of N C metrics with generalization alongside their associated p-values from the permutation tests described in §3.8. The majority of the correlations are statistically significant (p < 0.05) independent of scale, affirming that N C is correlated with generalization.The dichotomy of self-duality (N C 3 ) and uniform duality (UN C 3 ) is rather conspicuous. Our main experiments find that N C 3 does not consistently develop with scale while UN C 3 does (Fig. 3 ). However, within fixed scale, the opposite is true, implying that UN C 3 may be confounded by model capacity while N C 3 is a subtle and fine-grained indicator of generalization.Our models trained with either weight decay factor exhibited very similar patterns in the emergence of N C (or lack thereof), but the more aggressive factor β = 0.1 resulted in stronger development of N C properties than with β = 0.0005 (Appendices E, F, G, H, I, J, K, M, N, P). These findings empirically affirm β = 0.1 weight decay as sensible for LLM pre-training [81] , and concur with [56] on the pivotal role that appropriate regularization plays in the emergence of N C. 
limitations: 
 Neural collapse While to the best of our knowledge, no previous work has studied realistic stochastic token prediction, it is possible that the quantities that we measure are not perfectly suited for N C in language modeling. As we described in §1.1, the N C framework does not translate neatly to the language modeling space due to many adverse conditions, so full convergence to N C in the TPT was highly improbable. This paper leaves much room for future work to better adapt N C for next-token prediction, which we discuss further in Section 7.Language modeling Our work focused on autoregressive pre-training in its most basic form. We did not conduct experiments into encoder, multi-modal, or instruction-tuned models. Post-training techniques such as supervised fine-tuning, reinforcement learning with human feedback [91] or direct preference optimization [92] are also out-of-scope. This paper uses validation CE loss as the sole indicator of performance, leaving out any downstream task evaluations.The models that we used in our permutation test ( §5.1, Table 1 ) are only of a single small architecture trained for one epoch with relatively weak weight regularization (β = 0.0005). Therefore, our experimental results on scale-independent links between N C and generalization may not necessarily translate to larger models. Further investigation on (foundation) LLMs orders of magnitude larger than our CLMs trained with modern NLP methods would provide more robust insight into any direct correlations. 
discussion: 
 Layer/depth-wise neural collapse Past works have established that properties resembling N C evolve as a function of model depth [4, 36, 53, 59, 60, [93] [94] [95] [96] [97] [98] [99] [100] [101] [102] [103] [104] [105] . Layer-wise N C -sometimes dubbed deep neural collapse (DN C) [53, 59] -and related phenomena at intermediate layers remain an interesting subtopic. We leave their study and induction in CLMs (like [106] ) as future work.Learning to collapse Given the evidence for the development of N C and associated generalization under various loss functions [19, 49, 60, 101, [107] [108] [109] in other domains, NLP researchers may still benefit from analyzing, adapting or training towards N C. As alluded to earlier, the simplex ETF and even the CE loss may not be truly optimal for this problem setting, so we anticipate future works to both construct more amenable geometries with better-suited objectives and then capitalize on their benefits downstream. As discussed in §2, there is an abundance of literature in N C, some of which could potentially adapt N C to be useful for NLP; we hope to inspire such investigations.LLM Evaluations Researchers in NLP and multimodal settings are ultimately interested in measuring model performance on practical tasks. Notable benchmarks include GLUE [77] , MMLU [110] , and BIG-bench [111] . However, several contemporaries [112] [113] [114] have demonstrated that models' downstream capabilities are roughly correlated with their abilities to effectively compress their pretraining data. Based on their findings, our application of the N C framework to the pre-training stage of CLMs against validation CE loss should be an appropriate first step in this intersection. Looking forward, we anticipate exciting analysis for language modeling tasks or benchmarks, especially on creativity and retrieval for natural language understanding and generation.Conversely, some form of N C could even serve as an alternative evaluation. Although it would be prohibitively expensive to measure N C on the vast and sometimes obscure pre-training data of most frontier production LLMs, doing so on a small set of unseen data (i.e. test data) would be realistic.Interpretability At a high level, the number and density of clusters for a token can reflect its learned meanings and uses. This would be particularly useful as LLMs adapt to ever-evolving language and further expansion into non-English domains. Our formulae (Section 3) and results (Section 4) expose the pairwise token class interactions in noise, interference, classifier duality, and classifier agreement in the top-level features. Similarly to works in other domains [21, 55, 72, 115] , these N C metrics can serve as a form of low-level interpretability to aid understanding certain behaviors of (L)LMs. Between tokens, one can often discern how related or interchangeable words are based on their pair-wise interactions, or how antithetical or unrelated they are based on orthogonality. For example, we present a rudimentary analysis of homonyms and English first names in Appendix Q.Fairness Foundation models are ubiquitous for their comprehensive capabilities and adaptability.As previous work discussed class imbalance [35, 58, 61, 62] , our work may extend these strategies to measure and perhaps promote fairness in foundation LLMs, some of which are designed to be multilingual or multicultural. For example, [116] contemporarily explores the use of UN C3 to mitigate biases in BERT-based [7] models.While N C itself would not lead to unfairness, its potential interpretability may, in theory, enable an (adversarial) agent to measure and optimize for (un)fairness as they manipulate an LLM.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: We trained multiple copies of the same models and observed patterns consistent with the ones presented in our paper. We also state the number of data points and permutation trials we performed for our significance test. See Figure 1 and Sections 4 and 5.Guidelines:• The answer NA means that the paper does not include experiments. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
conclusion: 
 In this paper, we apply the neural collapse (N C) framework to the next-token prediction problem, where models are undertrained and next-tokens are variably drawn from numerous and imbalanced token classes. We leverage canonical and more recent metrics to demonstrate that N C emerges as we scale the size and training of hundreds of causal language models. Our results show a correlation between N C and generalization, a relationship that persists even when the model scale is fixed.In the short term, our work presents rudimentary techniques to analyze and interpret token-level properties of (L)LMs. We anticipate future work to suitably adapt N C (and related frameworks) to the still-fresh frontier of autoregressive language modeling. Researchers could then effectively capitalize on previous learnings from N C to better understand and improve the pre/post-training processes of increasingly complex and large language (multimodal) models. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: The paper includes a discussion on how N C might be extended to measure fairness in LLMs. Naturally this may inform positive or adversarial manipulations of LLMs.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not use pretrained models or scraped data. Our dataset is synthetic, public and transparent with no risks. The training procedure we use is the standard causal language modeling regime.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: We do not work with human subjects. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
: 
 Neural collapse (N C) is a phenomenon observed in classification tasks where top-layer representations collapse into their class means, which become equinorm, equiangular and aligned with the classifiers. These behaviors -associated with generalization and robustness -would manifest under specific conditions: models are trained towards zero loss, with noise-free labels belonging to balanced classes, which do not outnumber the model's hidden dimension. Recent studies have explored N C in the absence of one or more of these conditions to extend and capitalize on the associated benefits of ideal geometries. Language modeling presents a curious frontier, as training by token prediction constitutes a classification task where none of the conditions exist: the vocabulary is imbalanced and exceeds the embedding dimension; different tokens might correspond to similar contextual embeddings; and large language models (LLMs) in particular are typically only trained for a few epochs. This paper empirically investigates the impact of scaling the architectures and training of causal language models (CLMs) on their progression towards N C. We find that N C properties that develop with scale (and regularization) are linked to generalization. Moreover, there is evidence of some relationship between N C and generalization independent of scale. Our work thereby underscores the generality of N C as it extends to the novel and more challenging setting of language modeling. Downstream, we seek to inspire further research on the phenomenon to deepen our understanding of LLMs -and neural networks at large -and improve existing architectures based on N C-related properties. Our code is hosted on GitHub: https://github.com/rhubarbwu/linguistic-collapse. C ) [1] properties in 230 causal language models trained on TinyStories [2] , alongside improvement in generalization (i.e. validation performance). Left to right: N C 1 ) within-class (representation) variability collapse; GN C 2 ) hyperspherical uniformity of class means; UN C 3 ) uniform duality between class means and corresponding classifiers; and N C 4 ) agreement between token (maximum a prior) classifiers and implicit nearestclass center classifiers. Colored by model size and annotated with coefficient of determination (R 2 ). 
 
 Appendix: TinyStories is a dataset of short children's stories generated by GPT-3.5 and GPT-4 [2] , released with the CDLA-Sharing-1.0 licence. We trained and evaluated on the first version, as described:• The 2,141,709 stories are split into 2,119,719 train and 21,990 validation stories.• Their experimental setup [2] called for the GPT-2 [68] tokenizer, of which only a subset vocabulary V = 1, 29233 appears in TinyStories. • Following the GPT-style teacher-forcing regime for training/evaluation [8] , raw sequences (stories) from the train set are packed (by two preprocessing workers) into 229,367 (S) chunks of 2048 (T ) tokens each. This setup provides 469,514,249 (N ) ground-truth 18 token samples for training. # Samples (log scale)Figure 4 : The 500 most frequent classes from TinyStories [2] exhibit significant sample imbalance. Despite the synthetic nature of TinyStories, such a distribution is typical of natural language [11, 12] .The study of N C in causal language modeling at the token level would be very expensive, so the motivation to use a small dataset is clear. However, most commonly used text datasets such as WikiText [117] , BookCorpus [118] , CommonCrawl [119] , or most subsets from the Pile [120] are much too complex and broad to be effectively compressed by CLMs of the scale that we work with.WikiText-2 and WikiText-103 present significant drawbacks for our experiments. Both datasets contain a considerable amount of low-quality data that does not concentrate on essential linguistic structures such as grammar, vocabulary, facts, and reasoning. WikiText-2 has a similar empirical vocabulary to TinyStories under the GPT-Neo [80] tokenizer (27K vs. 29K) but only has around 34K rows of training data compared to 2.1M in TinyStories. Our small-scale NC experiment on WikiText-2 revealed that the models were very brittle and prone to overfitting. On the other hand, WikiText-103 is comparably sized to TinyStories but utilizes around 44K unique tokens. Our CLMs trained on WikiText-103 struggled to produce coherent sentences, likely due to the excessive breadth and information, as noted by the authors of TinyStories. Beyond these two, we were unable to find any real datasets that both followed established scaling laws [14, 15] for CLMs at our scale and are simple enough to suit the analysis of N C.According to its authors, TinyStories [2] is explicitly designed to preserve the essential elements of natural language, such as grammar, vocabulary, facts, and reasoning, while being smaller and more refined in terms of its breadth and diversity. Unlike large corpora that can overwhelm small language models (SLMs) due to their excessive breadth and diversity, TinyStories offers a concentrated dataset that hones in on core linguistic structures and reasoning capabilities. This is evident in its small word vocabulary, consisting of approximately 1500 words that a child would use, and in its 29K empirical vocabulary under the GPT-Neo tokenizer.Despite its concentrated nature, TinyStories enables models trained on it to produce grammatically correct, factual, and reasonable stories. Additionally, these models can be finetuned on specific instructions found in the TinyStories-Instruct dataset. The authors of TinyStories also demonstrate that their models can creatively produce stories that are dissimilar enough to their training data, indicating a balanced capability for generalization and creativity.One particular advantage of TinyStories is the small vocabulary relative to total training tokens, rendering a reasonable number of classes with higher average token counts. This is relevant because the possibility of N C and a CLM's ability to compress language data into distinct geometries depend partially on the ratios between embedding dimension, vocabulary size, and average token frequency. Conveniently, frequency analysis of the overall dataset produced a distribution (Figure 4 ) similar to real human language, so TinyStories should provide a good balance an initial study of this N C.Additionally, TinyStories has a more regular structure as GPT-3.5/4 was instructed to produce children's stories with certain themes and forms with a conservative vocabulary. We believe this would reduce the amount of clustering noise from the breadth of information and structures in real general data, and allow our smaller CLMs to exhibit some clear trends toward N C. Furthermore, TinyStories was created using GPT 3.5/4, which are advanced language models with significantly larger architectures trained on orders of magnitude more tokens, helping minimize the effect of the synthetic nature of the generated dataset. We also considered a possible effect of model collapse as a result of training on synthetic data [121] and follow-up work [122] suggest that a single iteration of data generation (as generated TinyStories) has a very negligible model collapse. Training was performed using an adaptation of an open-source causal language modeling script from Huggingface: https://github.com/huggingface/transformers/blob/main/ examples/pytorch/language-modeling/run_clm.py• Each model was trained on a single NVIDIA A100 (40GB) GPU for up to 8 hours per epoch.• Learning rates were set by a linear schedule based on the number of steps with no warm-up.• Training was performed in bfloat16 [123] mixed precision.• The results presented in this work are from two sets of models trained with weight decay β = 0.0005 [51] and β = 0.1 [81] . A previous set of models was trained without weight decay and the results are very similar to β = 0.0005. Codes for (post-)training analysis are hosted on GitHub:• Main code https://github.com/rhubarbwu/linguistic-collapse • Auxillary package: https://github.com/rhubarbwu/neural-collapseOne pass over the train set for embeddings collection can take up to 6 hours on a single NVIDIA A100 (40GB) GPU. Analysis of a single metric for a given model takes no longer than 5 minutes. M Self-Duality with Scale -(Against N C 3 ) Self-duality (N C 3 ) was originally the convergence of classifiers to means up to rescaling [1] :EQUATIONInstead, we use class-wise cosine similarity (Equation 8) and its variation (UN C3). This section presents token-wise interpretability results found in top-layer embeddings from our most learned model (https://huggingface.co/rhubarbwu/TinyStories-12x1024_10L). Table 4 : Under TinyStories-12x1024_10L, these fifteen homonyms have much shorter mean embedding norms (i.e. closer to the global centre) than the average token. This is expected since homonyms typically present conflicts and interference.Table 5 : Under TinyStories-12x1024_10L, the variability and interference of some English first names were far below those of the average token. This might be because names are distinct from one another and are not typically used in the same contexts as other words (aside from articles). The only names to have CDNV close to that of the average token are "Anna" and "Tim". Note that the positive interference of the last row (average token) is not a typo.GPT-Neo Token ID Token Text CDNV (Eq. 4) Interference (Eq. Justification: The claims of this paper listed in the abstract and introduction are simply the emergence of the N C phenomena and its relationship with generalization in CLMs, which are the scope and results of the paper.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material.Justification: We have read the Code of Ethics and determined that our paper does not violate the guidelines.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Our models range from 3.4M (small) to205M (large), so we inclusively use "CLM" instead of "LLM".2 Following[67], one might describe such a setting (C > d + 1) as a model "in superposition".Leveraging our generic N C package: https://github.com/rhubarbwu/neural-collapse.TinyStories was developed and evaluated as a faithful emulation of large language modeling, so we chose it for experimentation to train hundreds of CLMs and analyze embeddings at low cost. See Appendix A.In practice, we use an equivalent decomposition (Eq. 12).Note that agreement is not equivalent to accuracy.N = S(T -1) as we cannot use the first ground-truth nor the last predicted token in any chunk. 


Decision: The paper explores the phenomenon of neural collapse (NC) in the context of causal language models (CLMs). The authors aim to understand how NC properties, which have been predominantly studied in well-behaved classification settings that are trained on noisy, unbalanced data with a large and often imbalanced vocabulary. The paper applies the concept of NC to a new domain, CLMs, which differs significantly from the classification settings where NC has been traditionally studied, which is novel and valuable. The authors provide a comprehensive empirical evaluation, demonstrating a correlation between the proposed NC metrics and generalization in CLMs. One main concern raised by the reviewers is the use of a synthetic dataset (TinyStories), which may limit the generalizability of the findings to real-world human-written language data.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: questions: See "Weaknesses."
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Limitations are to my knowledge well adressed
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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 ?
limitations: limitations: None
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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)?
limitations: limitations: Yes

Comments: comment: I would like to thank the authors for the elaborate response. With regards to the the multi-label context I understand the arguments regarding the impact of context lenght and I appreciate the explanation. I would like this point to be further clarifier in future revisions. I would further like to encourage evaluation metrics other than test loss to validate the relationship between NC properties and improving LLM quality. Having said this, I would like to increase my score to a 5.
Thanks for the response! I maintain my accept recommendation.
Thanks for your response. 

Thanks for clarifying NC3. The intuition appears to rely on both NC3 and UNC3 being correlated with generalization, but Figs. 3, 19, and 21 show negligible/no correlation for NC3. This makes arguments like the one in line 255 somewhat misleading. Defining a metric that correlates with generalization doesn't make it a better replacement for an NC metric that (based on your arguments) is also likely to be tied to generalization. I recommend clarifying this issue in the updated manuscript.

Overall, I still feel like the technical novelties for adequately addressing the intricacies of language setups are limited in this paper. However, I also acknowledge that the paper, along with its extensive experiments, can motivate further investigations into extending the NC literature to language tasks. Thus, I will slightly increase my score
Thank you for your response, those additional results kind of also fit in with my other question on common LLM benchmarks, they provide examples of use cases of the mathematical framework. I maintain my opinion that this is an excellent work, and maintain the high score.