ICML 2024 | Feature pollution: Neural networks learn irrelevant features and fail to generalize

• Paper title: Feature Contamination: Neural Networks Learn Uncorrelated Features and Fail to Generalize
• Paper link: https://arxiv.org/pdf/2406.03345
• Code link: https:/ /github.com/trzhang0116/feature-contamination

With the great success of large models represented by GPT in recent years, the machine learning paradigm of deep neural network + SGD + scaling once again proved its dominance in the field of AI status. Why are paradigms based on deep neural networks successful? The more common view is that neural networks have the ability to automatically learn abstract and generalizable features from massive high-dimensional input data. Unfortunately, limited by the shortcomings of current analysis methods and mathematical tools, our current understanding of "how (deep) neural networks implement such a feature learning process" is still not very deep. Because of this, most of the current relevant research in the academic community still remains at the level of "explaining" the features that the model has learned, and it is difficult to obtain more data-efficient and more generalizable models through "intervention" in its learning process. Model. When we discuss the feature learning process of neural networks, one of the most basic questions is: What features will the neural network learn from the input data? From a goal perspective, the feature learning of neural networks is a "by-product" driven by tasks, and its purpose is to minimize training errors. Therefore, we would intuitively think that the neural network should extract "task-relevant" features from the data, while the remaining "task-irrelevant" features are equivalent to data noise. Then, because neural networks have the characteristic of "not learning unless necessary" (more precisely, simplicity bias), neural networks should tend not to learn them. This is also a common view in the current literature.

However, in our recent work accepted by ICML 2024, we found that such intuitive cognition is actually

wrong

! Specifically, we found that when nonlinear neural networks learn task-related features, they also tend to learn task-irrelevantfeatures (we call it "feature pollution"), and this tendency can lead to neural It is difficult for the network to generalize to scenarios with distribution shift. Theoretically, we proved that feature contamination occurs even in a simple two-layer ReLU network and is closely related to the category asymmetry of neuron activation in neural networks; experimentally, we also gave a series of evidence that features Contamination also exists in deep networks such as ResNet and Vision transformer, and will adversely affect their generalization. It is worth mentioning that the failure mode we discovered is completely orthogonal to the mainstream analysis based on spurious correlations in the current out-of-distribution (OOD) generalization literature. Therefore, from a larger perspective, our findings demonstrate the importance of the inductive bias of the neural network itself for OOD generalization. It also shows that many of our studies on neural network feature learning and generalization have been Intuition may also need to be rethought.

Next, let’s introduce the specific content of the article:

Research background

The generalization ability in scenarios where data distribution changes (that is, OOD generalization ability) is a measure of whether a machine learning system can perform in reality One of the key indicators of deployment in the environment. However, current neural networks often suffer significant performance losses in OOD generalization scenarios. As for the reason why OOD generalization fails, the more mainstream statement in the literature is that spurious correlations exist in the representation, that is, the model will learn features that are related to the task goal but have no causal relationship. Therefore, when the correlation between these features and task objectives changes due to distribution shifts, models that rely on these features for prediction cannot guarantee the original performance.

The above theoretical explanation is quite intuitive and natural, and has also become the main line guiding OOD algorithm research in recent years, that is, by designing better optimization objective functions and regular terms, the model can learn better representations without false correlations. In order to obtain stronger generalization performance. In recent years, there has been a lot of work along this main line trying to improve the OOD generalization of the model through algorithm design. However, recent work shows that many algorithms with built-in theoretical guarantees have very limited performance improvement on OOD generalization tasks based on real data. Why does this happen? We believe that the current difficulties in OOD generalization research may stem from

two limitations

of existing analyses:

• Most of the existing research only considers the failure mode caused by spurious correlation;
• Most of the current research is limited to linear models and does not consider the nonlinearity of neural networks and the inductive bias of SGD, so the existing analysis results are not necessarily Suitable for the neural network we actually use.

In other words, current explanations and theoretical models of OOD generalization may not accurately reflect real-world distribution shift scenarios. Therefore, we believe that considering the inductive bias of neural networks and SGD is very necessary to understand the generalization of OOD based on deep neural networks.

Experiment

First, we try to estimate the "performance upper bound" that can be achieved by the current OOD generalization algorithm designed based on representation learning goals through experimental design. Under the guidance of spurious correlation theory, existing work mainly attempts to constrain the model to learn representations that can be generalized by OOD by designing auxiliary representation learning objective functions. In order to study whether optimizing such an objective can actually extract the desired representation, we designed an idealized scenario:

• First, during the training process, we allow the model to explicitly fit a teacher model that can be generalized by OOD. The extracted representation is representation distillation. In experiments, this teacher model can be a large-scale pre-trained model (such as CLIP). In order to control variables, in actual operation we control the model structure of the student model and the teacher model to be exactly the same.
• In the second step, we train linear classifiers (linear probing) on ​​the training set based on the representations provided by the teacher model and student model respectively. ,
• Finally, we tested the linear classifiers based on the teacher model and the student model on the identically distributed test set and the OOD test set, respectively, to measure the OOD generalization of the representations extracted by these two models.

The experimental results are shown in the picture above. From the figure we have two main findings:

• Compared with the standard model (blue) that does not directly fit the teacher model representation during the training process, the linear classifier based on the student model (orange) does have better OOD Generalizability;
• However, the OOD generalization performance of the linear classifier based on the student model (orange) is still significantly behind the linear classifier based on the teacher model (purple).

So we naturally ask: Since we have directly fitted the representation of the teacher model, where does the generalization gap between the student model and the teacher model come from? We found that it is difficult to directly explain this experimental phenomenon with the existing theoretical explanations:

• First of all, this gap cannot be directly explained by the spurious correlation theory: since the representations of the student model and teacher model (on the training set) are basically The same, then the linear classifier based on these two representations should be similarly affected by spurious correlation features during the training process, and should not have such a large gap;
• Another possible explanation is the teacher model (such as CLIP) may have "seen" many OOD samples during its own pre-training process, so it can extract some features that are not found on the training set for the OOD samples. However, recent research shows that even if all samples similar to OOD test samples are removed from CLIP's pre-training data, CLIP still has strong OOD generalization [1]. This shows that it is not sufficient to explain the gap between the teacher model and the student model simply from this perspective.

In short, we believe that the existing analysis is insufficient to explain the gap in OOD generalization ability that we actually observed in our experiments. At the same time, since "directly fitting representations that can be generalized by OOD" cannot guarantee a model that can be generalized by OOD, we have to consider the "process" of representation learning in addition to the "goal" of representation learning. ", which is the inductive bias caused by the feature learning dynamics of neural networks. Although it is very difficult to directly analyze the feature learning process of deep neural networks in theory, we found that even a two-layer ReLU network will show an interesting feature learning tendency, that is, "feature pollution", and this tendency It is also directly related to the OOD generalization of neural networks.

Theory

In this section, we prove the existence of the "feature pollution" phenomenon on a binary classification problem based on a two-layer ReLU network, and analyze the source of this phenomenon. Specifically, we assume that the input to the network is composed of a linear combination of two features: "core features" and "background features". Among them, the distribution of core features depends on the category label (can be visualized as the object to be classified in the image classification problem), while the distribution of the background features has nothing to do with the label (can be visualized as the picture background and other elements in the image classification problem). In order to eliminate the interference of other factors, we also make the following assumptions about these two types of features:

• 背景特征和标签不相关（这样我们就排除了由虚假相关性导致的failure mode）。
• 通过核心特征可以对标签实现100%准确率的预测（这样我们就排除了由于训练集的特征不够导致的failure mode）。
• 核心特征和背景特征分布在正交的子空间中（这样我们就排除由于不同特征难以解耦导致的failure mode）。

我们发现，即使在以上的条件下，神经网络仍然会在学习核心特征的同时学习和任务完全不相关的背景特征。由于这两种特征在网络权重空间的耦合，在背景特征上发生的分布偏移也会导致神经网络的误差增大，从而降低网络的OOD泛化性。我们也因此把这种神经网络的特征学习偏好称之为“特征污染”。以下，我们详细介绍特征污染现象的出现原因。整体分析思路的示意图如下:

我们分析中的关键点在于：特征污染实际上和神经网络中的神经元往往对不同类别具有不对称激活（asymmetric activation）有关。具体而言，我们可以证明在经过足够的SGD迭代后，网络中至少有相当一部分的神经元都会被倾向于而与一个类别的样本保持正相关（我们称之为该神经元的正样本，并用y_pos表示其类别），而与另外一个类别的样本保持负相关（我们称之为该神经元的负样本，并用y_neg表示其类别）。这就会导致这些神经元的激活具有类别不对称性，如定理4.1所示：

这样的类别不对称性是怎么影响神经网络的特征学习过程的呢？我们首先注意到，对于网络隐层的第k个神经元，其权重向量w_k在第t次迭代后可以被拆分为：

上式中，S_core和S_bg分别表示核心特征和背景特征的集合，其中每个m_j都对应一个核心特征或者背景特征。从该式中我们可以看出，神经元的权重可以分解为其在不同特征上的投影（这里我们假设不同的m_j之间都是正交的单位向量）。进一步地，我们可以证明在w_k的负梯度在每一个背景特征m_j，j属于S_bg上的投影满足：

对于具有类别不对称激活的神经元，根据定理4.1我们可得其梯度主要取决于该神经元的正样本y=y_pos而和负样本y=y_neg几乎无关。这就导致正样本中存在的核心特征和背景特征会同时得到正的梯度投影，而这一过程和特征与标签之间的相关性无关。

如定理4.2所示，我们证明了在经过足够的SGD迭代后，上面这种梯度投影的积累将导致神经元学习到的特征既包含核心特征，也包含耦合的背景特征：

由于核心特征和背景特征在神经元权重中的耦合，背景特征的负向分布偏移会降低神经元的激活，导致额外的 OOD 误差。如定理4.3所示，我们定量描述了特征污染对 ID 和 OOD 泛化风险的影响：

同时，为了进一步说明特征污染源自神经网络的非线性激活函数之间的关系，我们证明了在去除掉神经网络的非线性后，特征污染将不再发生：

As shown in the figure below, we verified our theoretical results through numerical experiments. At the same time, in addition to the two-layer ReLU network + SGD, we also extended our conclusions to more general settings, including other types of activation functions, optimizers with adaptive step sizes, etc. The results are shown in Figure 3(d) ), indicating that feature contamination is also prevalent in more general settings.

At the same time, we have also provided more experimental evidence and feature visualization to show that in the deep networks such as ResNet and Vision transformer that we use daily, the phenomenon of feature pollution also occurs, and can explain the observations in our experiments The OOD generalization gap reached. Anyone who is interested in this part can refer to Chapter 5 of our original paper.

Summary and discussion

Finally, we list some research points that may be more important in the future/can be continued in depth. We also welcome everyone who is interested to communicate with us further:

• Deeper network:Although We have experimentally proven that deep networks also have feature pollution problems, but so far our theoretical analysis has only done a two-layer ReLU network. We suspect that feature contamination may be a more general concept, and the activation asymmetry of neurons for categories may be only one of the reasons for its occurrence. By analyzing deeper networks or more complex network structures (such as introducing a normalization layer, etc.), we may be able to discover more causes of feature pollution and provide targeted solutions.
• The role of pre-training: The theoretical analysis in this article only considers the case of train from scratch, but the models we actually use are often pre-trained models. There is a lot of experimental evidence that pre-training can help improve the OOD generalization of the model. So, is the essence of this improvement in generalization related to mitigating the feature pollution problem? How does pre-training do this?
• How to solve the feature pollution problem: Although our work pointed out the feature pollution problem, it has not yet given a clear solution. However, some of our subsequent work has shown that similar problems will also occur when fine-tuning large models, and we have also found that some gradient adjustment-based methods can indeed alleviate this problem, thereby significantly improving the fine-tuning model. generalization ability. We will also release the specific content of this part of the work in the future, and everyone is welcome to continue to pay attention.

About the author | The author of this article, Zhang Tianren, is a doctoral candidate in the Department of Automation, Tsinghua University. He holds a bachelor's degree from the Department of Automation, Tsinghua University. His supervisor is Professor Chen Feng. During his PhD, the author mainly conducted theoretical and algorithmic research around representation learning and generalization issues in machine learning. He has published many articles in top machine learning conferences and journals, such as ICML, NeurIPS, ICLR, IEEE TPAMI, etc.

Author affiliation | Tsinghua University VIPLAB

Contact email | zhangtr22@mails.tsinghua.edu.cn

References

[1] Mayilvahanan, P., Wiedemer, T., Rusak, E ., Bethge, M., and Brendel, W. Does CLIP's generalization performance mainly stem from high train-test similarity? In International Conference on Learning Representations, 2024.

The above is the detailed content of ICML 2024 | Feature pollution: Neural networks learn irrelevant features and fail to generalize. For more information, please follow other related articles on the PHP Chinese website!