Paper: Feature Collapse (ICLR-2024)

Reviewer comment (Experiments):
While the key contributions of the paper are theoretical, the motivation for the theoretical results would have been clearer if a stronger connection could be made to more realistic settings. For example, the connection between feature collapse and generalization outside the context of the toy task does not seem to be well established. Additionally, the degree to which feature collapse sufficiently occurs or does not occur in the context of more realistic models and tasks is not clear. While perhaps not strictly necessary for such a theoretical paper, addressing such questions would have improved the motivation for this work.

Author response:
**Weakness1**:
>  The motivation for the theoretical results would have been clearer if a stronger connection could be made to more realistic settings... The degree to which feature collapse sufficiently occurs or does not occur in the context of more realistic models and tasks is not clear. 

We have studied the feature collapse phenomenon in a controlled setting where proofs are analytically tractable. We agree that, to better motivate this analytical study, we could have provided empirical evidence that this phenomenon occurs in more complex settings.

In the updated version of the paper we:

1) Empirically investigate feature collapse in the setting of a Context Free Grammar (CFG), and where the network is a deep neural network with ReLU nonlinearities and LayerNorm module at each layer. Context Free Grammars [1] are synthetic data models that generate sentences with probabilistic trees and are widely used to understand natural language models (e.g. [2,3,4]). A CFG can be viewed as a `deep' version of our data model, in that there is a hierarchy of latent variables (words, concepts, meta-concepts, meta-meta-concepts, etc), and the latent variables from one level generate latent variables in the level below. The empirical results shows that, if the architecture of the neural network matches that of the CFG, then the representations at each layer of the networks collapse according to the corresponding latent variables.

2) Train a 2-layer transformer on the classification task depicted on Figure 1 of our paper. Feature collapse does take place in this setting as well.

3) Train a 2-layer transformer (with GPT-2 architecture) to predict the next token on sentences generated by our data model (i.e. a language modeling task). Feature collapse also occurs in this situation. 

These three experiments are briefly described in the conclusion of the updated version of our paper (page 9, in blue), with precise details and visuals in the first section of the appendix (page 12-15, in blue). They demonstrate that feature collapse still occurs with a more complex data model (CFG), more complex networks (such as transformers and deep networks with ReLU), and a different task (language modeling).


[1] Chomsky. Three models for the description of language

[2] Kim, Dyer, Rush. Compound probabilistic context-free grammars for grammar induction

[3] Allen-Zhu, Li. Physics of Language Models: Part 1, Context-Free Grammar

[4] Liu, Xie, Li, Ma. Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models

Reviewer comment (Experiments):
The role of the L2 regularization term in equation 2 could be clarified. The authors state "The regularization terms play no essential role apart from making proofs easier; the empirical picture remains the same without weight decay". Could the authors include the experimental results that support this claim? This also seems to diminish (although not completely) the key theoretical contribution. It is not surprising that in the presence of L2 regularization, feature magnitudes would depend on their frequency in the training data. It would be more interesting if we had theory that elucidated why the learned embeddings for rarer tokens have lower magnitude in the absence of explicit regularization.

Author response:
**Weakness2**:
>  The role of the L2 regularization term in equation 2 could be clarified

Indeed, we could have been clearer on this point, and we have updated the paper accordingly (see paragraph in blue on page 4). Without weight decay, the objective (2) does not have a global minimum, which makes the theoretical study much more challenging. In such a situation, gradient descent leads to solutions where the magnitude of the weights grows logarithmically to infinity and becomes more and more aligned with the gradient, c.f. [1,2]. In Appendix A1 of the new version of the paper, we reproduce all our experiments but without weight decay. To train the networks, we performed 5 million iterations of stochastic gradient descent with a batch size of $100$ and a learning rate of $0.1$. As can be seen in Figure 5 and 6 in the appendix, the results are qualitatively very similar to those obtained with weight decay.

[1] Ji, Telgarsky. Directional convergence and alignment in deep learning

[2] Soudry at al. The Implicit Bias of Gradient Descent on Separable Data

Reviewer comment (Theory):
Is there any setting where feature collapse may *not* be desirable? For example, consider a task with a significant degree of label noise and with limited data samples to learn from. Intuitively, features that are observed more frequently during training may be more reliable predictors of the label. Therefore, to the extent that the magnitude of word embeddings relates to their salience in making predictions, it seems that feature collapse may not be strictly desirable in such a setting.

Author response:
**Question**: 
> Is there any setting where feature collapse may not be desirable?  For example...

We fully agree that, in some settings, feature collapse might not be desirable. To put it simply, in our setting the frequency of a word is irrelevant to the classifications task. Using a network architecture that 'removes' this irrelevant information is therefore beneficial. In a setting where the word frequency is meaningful, such as the one described by the referee, a network that ignores it would perform poorly.

Generally speaking, we believe that the range of applications for deep learning is so varied that no one mechanism could possibly be right for all of them. Rather, the right network architecture is heavily task dependent. Our contribution, we hope, is to shed some light on how different architecture choices leads to different types of representations.

Reviewer comment (Experiments):
While the reviewer appreciates the feature collapse concept and believes this phenomenon is not limited to the simple network settings used in the manuscript, the absence of empirical validation within more complex networks or datasets constrains the soundness and significance of the findings.

Author response:
**Weakness1**: 
> The absence of empirical validation within more complex networks or datasets constrains the soundness and significance of the findings.

We agree. It is important to show (empirically) that the feature collapse phenomenon is not limited to the simple controlled setting where proofs are analytically tractable. In the updated version of the paper, we conduct such an empirical investigation with more complex data models and more complex networks. Specifically, we ---

1) Empirically investigate feature collapse in the setting of a Context Free Grammar (CFG), and where the network is a deep neural network with ReLU nonlinearities and LayerNorm module at each layer. Context Free Grammars [1] are synthetic data models that generate sentences with probabilistic trees and are widely used to understand natural language models (e.g. [2,3,4]). A CFG can be viewed as a `deep' version of our data model, in that there is a hierarchy of latent variables (words, concepts, meta-concepts, meta-meta-concepts, etc), and the latent variables from one level generate latent variables in the level below. The empirical results shows that, if the architecture of the neural network matches that of the CFG, then the representations at each layer of the networks collapse according to the corresponding latent variables.

2) Train a 2-layer transformer on the classification task depicted on Figure 1 of our paper. Feature collapse does take place in this setting as well.

3) Train a 2-layer transformer (with GPT-2 architecture) to predict the next token on sentences generated by our data model (i.e. a language modeling task). Feature collapse also occurs in this situation. 

These three experiments are briefly described in the conclusion of the updated version of our paper (page 9, in blue), with precise details and visuals in the first section of the appendix (page 12-15, in blue). They demonstrate that feature collapse still occurs with a more complex data model (CFG), more complex networks (such as transformers and deep networks with ReLU), and a different task (language modeling).  


[1] Chomsky. Three models for the description of language

[2] Kim, Dyer, Rush. Compound probabilistic context-free grammars for grammar induction

[3] Allen-Zhu, Li. Physics of Language Models: Part 1, Context-Free Grammar

[4] Liu, Xie, Li, Ma. Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models

Reviewer comment (Experiments):
The authors differentiate their work from the Neural Collapse (NC) concept by emphasizing that, unlike NC, Feature Collapse inherently suggests enhanced generalization and transfer learning capabilities. But once again, this claim lacks strong empirical evidence, the 2-layer network used in the paper is not comprehensive enough to demonstrate this point.

Author response:
**Weakness2**:
> The authors differentiate their work from the Neural Collapse (NC) concept by emphasizing that, unlike NC, Feature Collapse inherently suggests enhanced generalization and transfer learning capabilities. But once again, this claim lacks strong empirical evidence.

We agree that our claim lack strong empirical evidence. We had wanted to emphasize the fact that feature collapse and neural collapse are truly distinct phenomena. In doing so we did not intend to over-claim regarding feature collapse and generalization/transfer learning. The offending sentences

`Thus neural collapse has no consequences for generalization or transfer learning. By contrast, feature collapse  describes the emergence of ‘good’ local features, and this phenomenon goes hand in hand with generalization and transfer learning.'  

have been removed from the paper.

Let us clarify the point we inartfully tried to make. In the context of our data model, neural collapse means that the network maps all sentences of category 1 to a single point in its last layer, all sentences of category 2 to a single point in its last layer, etc., and that these points are maximally opposed. In other words, neural collapse states that the network has perfectly learned the labels explicitly given during training. Also, neural collapse occurs for any network that is expressive enough -- it doesn't depend on the architecture. It doesn't provide any information on whether or not the network has learned useful lower-level local features, and so in this limited sense, does not speak directly to transfer learning.

For feature collapse we study how different architectural choices lead to different types of local representations in the first layer (and intermediate layers in the case of the CFG) of the network. The same intermediate representation is simultaneously leveraged by many categories during classification. For example, the representation for the word 'cheese' is used by the upper layer to classify sentences from many different categories. As the representations are not directly tied to the categories of the classification task, they could conceivably be used in a transfer learning scenario.

Reviewer comment (Experiments):
The work primarily studies that features of words with similar semantics would collapse, how about words that sharing similar syntactic in a sentence? Would something similar happen as well?

Author response:
**Question2**:
>  The work primarily studies that features of words with similar semantics would collapse, how about words that sharing similar syntactic in a sentence? 

Our current data model is quite simplistic in the sense that tokens are either 'perfectly similar', or 'perfectly dissimilar' -- there is no middle ground. In future works, we would like to investigate more sophisticated models in which tokens might have different attributes (say color/size/shape, or semantic/syntactic). Some tokens might have some attributes in common, but might differ with respect to some others. The question is then: how is this reflected in their representations? What architectures promote useful representations for such a data model? This is a challenging question, and we currently have no answer. Our hope is that one would see some type of partial collapse along some dimensions coding for the attribute.

Reviewer comment (Theory):
For the theory part, it would be helpful to highlight core technical challenges the authors met and addressed in the proofs, or any new proof techniques the authors developed.

Author response:
Thank you for the question. One of the key difficulties was to realize that the latent structure of the problem exhibits a type of 'symmetry' when we consider a large enough number of latent variables. This symmetry can in turn be leveraged to obtain fully tractable analytical solutions. Specifically, the symmetry on the latent variables allows us to write the loss in a type of 'spherical coordinates', where a `sphere' is defined as the latent variables which are at a Hamming distance $r$ from a reference latent variable. This spherical decomposition for the loss then allowed us to explicitly characterize global minima. To the best of our knowledge, these techniques have not been used before in this context.

Another surprise was the fact that, when the words are not uniformly distributed, we could still derive analytical formula for the critical points of the loss (both with and without layer norm). We did not expect such solutions to have tractable analytical formula.

We have added such a discussion in the new version of the paper (see blue paragraph on p9).

Reviewer comment (Evaluation):
For the experiment part, I am not sure how it could help design or explain applications. I totally agree that semantically similar instances would lead to similar representations. But is it benign or not for the network to assign exactly the same representation to equivalent instances? If it is benign, how to ensure it happens in practice? If not, then how to avoid it? It would be helpful for the authors to give some convincing scenarios that people care about this phenomenon.

Author response:
In general the desirability of full collapse depends on the task at hand. If collapse destroys information relevant to a particular task then it is clearly not benign. A simple but concrete example of this might be the overall frequency of a word in our data model. In our setting this frequency is irrelevant to our classification task; so, full collapse (i.e., words with distinct frequencies but the same meaning receive the same representation) is beneficial and empirically improves generalization. Had we changed the task to one where learning the frequency is needed to solve it, full collapse is not benign since the frequency information is destroyed. 

One of the outcomes of our study, which we did not expect going in but makes sense in retrospect, is that proper normalization helps similar tokens receive similar embeddings. 

Personally, we care about the phenomenon to the extent that it sheds light on the mechanisms by which semantically similar instances receive similar representations. Ideally, we would want a theoretical framework that explains the `if, when and how' of semantically similar instances receiving similar representations. This is a challenging question, and in particular, it requires a careful and precise mathematical definition of what 'semantically similar' means in the form of a metric and a proof that smallness in this metric implies small distance between representations.

In the current work, we study an extreme version of the above problem. Specifically, we consider a setting in which tokens are either semantically identical or semantically unrelated (i.e., the semantic distance between tokens is binary). In this simple setting, we can prove that representations are either identical or maximally opposed.

We are very interested in extending our analysis to settings in which the semantic distance between tokens in not binary. In future works, we would like to investigate more sophisticated models in which tokens might have different attributes (say color/size/shape). Some tokens might have some attributes in common, but might differ with respect to some others. The question is then: how is this reflected in their representations? What architectures promote useful representations for such a data model? This is a challenging question, and we currently have no answer. Our hope is that one would see some type of partial collapse along some dimensions coding for the attribute.

Reviewer comment (Experiments):
Maybe a simpler and more intuitive experiment to demonstrate this “feature collapse” behavior is to assign two embeddings (to be learned) to one token, but during training this token is randomly mapped to these two embeddings, and see if these two embeddings become similar after training.

Author response:
Note that this experiment is essentially the same than the one we conduct in our work. Indeed, in our setting, each `concept token' is randomly mapped to one of out of 400 embeddings (the 400 embeddings corresponding to the 400 words contains in the concept).

Reviewer comment (Experiments):
Can we empirically verify the feature collapse in deeper layers?

Author response:
Thanks for the question. We conducted a careful experiment to verify this. The data model presented in the paper extends to one with a deeper hierarchy of latent structures. Recall that words are partitioned into concepts and that the latent variables are sequences of concepts. We can further partitioned the latent variables into 'meta-concepts' and create 'deeper latent variables' that are sequences of 'meta-concepts'. We can iterate this process to obtain a hierarchy of any depth. Such data model is a particular instance of a Context Free Grammar [1], which generate sentences with probabilistic trees and are widely used to understand natural language models (e.g. [2,3,4]). In Figure 9 in the appendix, we provide an illustration of a simple depth 3 context free grammar.

We ran experiments with a context free grammar of depth 4 (meaning that we have words, concepts, meta-concepts and meta-meta-concepts). We used a deep neural network with ReLU nonlinearities and LayerNorm module at each layer. The architecture of the neural network was chosen to match that of the context free grammar; see Figure 10 in the appendix. In Figure 11 we plot the activations after each of the three hidden layers and readily observe the expected feature collapse phenomenon. All segments of the input sentence that correspond to same concept, meta-concept, or meta-meta-concept, receive the same representations in the appropriate layer of the network (layer 1 for concepts, layer 2 for meta-concepts, and layer 3 for meta-meta-concepts). This shows that the feature collapse phenomenon also occurs in deeper layers.

This experiment is briefly described in the conclusion of the updated version of our paper (p9 in blue). The precise details and visuals can be found in appendix A5 (page 13). 

[1] Chomsky. Three models for the description of language

[2] Kim, Dyer, Rush. Compound probabilistic context-free grammars for grammar induction

[3] Allen-Zhu, Li. Physics of Language Models: Part 1, Context-Free Grammar

[4] Liu, Xie, Li, Ma. Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models