ICML2024|特徵汙染：神經網路會學習不相關特徵而泛化失敗

©PaperWeekly 原創 · 作者 | 章天任、趙處傑

單位 | 清華大學 VIPLAB

研究方向 | 表徵學習、多模態模型

隨著以 GPT 為代表的大模型在近年來取得的巨大成功，深度神經網路 + SGD + scaling 的機器學習正規化再次證明了其在 AI 領域的主導地位。為什麼基於深度神經網路的正規化能夠取得成功？

比較普遍的觀點是：神經網路具有從海量的高維輸入資料中自動學習抽象而可泛化的特徵的能力。例如，近期 LLM 領域的一些可解釋性工作表明，LLM 的一些隱層表徵確實可以和人類可理解的高層語義概念相對應 [1]。

遺憾的是，受限於當前分析手段和數學工具的不足，目前我們對於“（深度）神經網路如何實現這樣的特徵學習過程”這一問題的理解仍然很不深刻。也正因如此，目前學界的相關研究大多還停留在“解釋”模型已經學到的特徵的層面，而很難透過對於其學習過程的“干預”來得到更加資料高效、泛化性更強的模型。

當我們討論神經網路的特徵學習過程時，一個最基本的問題是：神經網路會從輸入資料中學到什麼特徵？

從目標上看，神經網路的特徵學習是任務驅動下的“副產品”，其目的就是為了最小化訓練誤差。因此我們會直觀地想，神經網路應該從資料中提取“任務相關”的特徵，而剩餘的“任務無關”的特徵則相當於資料噪聲。那麼，由於神經網路自帶“非必要不學習”（更準確地說是 simplicity bias [2]）的特點，神經網路就應當傾向於不對它們進行學習。這也是目前文獻中比較普遍的觀點 [3]。

然而，在我們最近的一篇被 ICML 2024 接收的工作中，我們發現這樣的直觀認知實際上是錯誤的！具體而言，我們發現非線性神經網路在學習任務相關的特徵時還會同時有學習任務不相關特徵的傾向（我們稱之為"特徵汙染"

），並且這種傾向會導致神經網路難以泛化至具有分佈偏移（distribution shift）的場景。

理論上，我們證明了特徵汙染即使在簡單的兩層 ReLU 網路中都會出現，並且和神經網路中神經元啟用的類別不對稱性息息相關；實驗上，我們也給出了一系列證據表明特徵汙染在 ResNet、Vision transformer 等深層網路中也同樣存在，並且會對其泛化性產生不利影響。

值得一提的是，我們發現的這種 failure mode 和當前分佈外（out-of-distribution, OOD）泛化文獻中主流的基於虛假相關性（spurious correlations）的分析是完全正交的。因此，從更大的角度看，我們的發現側面表明了神經網路自身的歸納偏置（inductive bias）對於 OOD 泛化的重要性，同時也表明：我們關於神經網路特徵學習和泛化的很多已有直覺可能也都需要被重新思考~

接下來，我們介紹一下文章的具體內容。

論文標題：

Feature Contamination: Neural Networks Learn Uncorrelated Features and Fail to Generalize

論文連結：

https://arxiv.org/pdf/2406.03345

程式碼連結：

https://github.com/trzhang0116/feature-contamination

研究背景

在資料分佈發生變化的場景中的泛化能力（也即 OOD 泛化能力）是衡量機器學習系統能否在現實環境中部署的關鍵指標之一。然而，當前的神經網路在 OOD 泛化場景中經常會遭遇顯著的效能損失。

關於 OOD 泛化失敗的原因，文獻中比較主流的說法是表徵中存在的虛假相關性（spurious correlations），也即模型會學習到與任務目標相關但無因果關係的特徵 [4]。於是，當這些特徵和任務目標之間的相關性由於分佈偏移而發生變化時，依賴於這些特徵進行預測的模型就無法保證原有效能。

以上這個理論解釋相當直觀且自然，也成為了指導近年 OOD 演算法研究的主線，即透過設計更好的最佳化目標函式和正則項來使得模型學習到更好的、不帶有虛假相關性的表徵，以期得到更強的泛化效能。近年來，已經有大量工作沿著這條主線試圖透過演算法設計來提升模型的 OOD 泛化性。然而，近來的工作表明很多自帶理論保障的演算法在基於真實資料的 OOD 泛化任務上的效能提升卻非常有限 [5]。

為什麼會出現這樣的情況？我們認為，目前 OOD 泛化研究的困難可能來源於現有分析的兩個侷限性：

現有研究大部分僅考慮虛假相關性導致的 failure mode；
目前研究大部分侷限於線性模型，而沒有考慮神經網路的非線性和 SGD 的 inductive bias，因而已有的分析結果也未必適用於我們實際使用的神經網路。

換言之，目前對 OOD 泛化的解釋和理論模型可能無法準確地反映真實世界的分佈偏移場景。因此我們認為，考慮神經網路和 SGD 的 inductive bias 對於理解基於深度神經網路的 OOD 泛化是十分必要的。

實驗：神經網路難以學習可OOD泛化的表徵

首先，我們嘗試透過實驗設計對當前基於表徵學習目標設計的 OOD 泛化演算法所能取得的“效能上界”進行預估。現有工作在虛假相關性理論的引導下，主要嘗試透過設計輔助的表徵學習目標函式來消除特定於某種分佈的虛假相關表徵，進而約束模型學習到可 OOD 泛化的表徵。為了研究最佳化這樣的目標是否能真正提取到期望的表徵，我們設計了一個理想化的場景：

首先，在訓練過程中，我們允許模型顯式擬合一個可 OOD 泛化的 teacher model 所提取出的表徵。實驗中，這個 teacher model 可以是一個大規模預訓練模型（如 CLIP [6]）。本質上看，這是一種表徵蒸餾操作，不過為了控制變數，在實際操作時我們控制待訓練模型（student model）和 teacher model 的模型結構完全一樣。
第二步，我們在訓練集上分別基於 teacher model 和 student model 所提供的表徵訓練線性分類器（linear probing）。
最後，我們在同分布（in-distribution）測試集和 OOD 測試集上分別對基於 teacher model 和 student model 的線性分類器進行測試，從而度量這兩個模型所提取的表徵的 OOD 泛化性。

實驗結果如上圖所示。從圖中我們有兩個主要發現：

與訓練過程中沒有直接擬合 teacher model 表徵的標準模型（藍色）相比，基於 student model 的線性分類器（橙色）確實有更好的 OOD 泛化性；
然而，基於 student model 的線性分類器（橙色）的 OOD 泛化效能仍然明顯落後於基於 teacher model 的線性分類器（紫色）。

於是我們自然會問：既然都已經直接擬合了 teacher model 的表徵，那 student model 和 teacher model 之間的泛化性 gap 從何而來？我們發現，其實目前已有的理論解釋是很難直接解釋這個實驗現象的：

首先，這個 gap 不能直接被虛假相關性理論解釋：既然 student model 和 teacher model 的表徵（在訓練集上）基本一樣，那麼基於這兩種表徵的線性分類器在訓練過程中受到虛假相關性特徵的影響也應該是類似的，而不應該出現這麼大的 gap；
另一個可能的解釋是 teacher model（如 CLIP）在它自己的預訓練過程中可能已經“見過”許多 OOD 樣本了，所以可以針對 OOD 樣本提取一些在訓練集上沒有的特徵。然而最近有研究表明即使從 CLIP 的預訓練資料中把所有和 OOD 測試樣本相似的樣本都去掉，CLIP 仍然有很強的 OOD 泛化性 [7]。這說明單純從這個角度來解釋 teacher model 和 student model 之間的 gap 也是不充分的。

簡言之，我們認為現有的分析不足以解釋我們在實驗中實際觀測到的 OOD 泛化能力 gap。同時，既然“直接擬合可 OOD 泛化的表徵”都不能保證得到可以 OOD 泛化的模型，那麼我們也就不得不在考慮表徵學習的“目標”之外同時考慮表徵學習的“過程”

，也就是神經網路的特徵學習動力學帶來的 inductive bias。

儘管從理論上直接分析深層神經網路的特徵學習過程是十分困難的，但我們發現，即使是兩層 ReLU 網路也會表現出很有趣的特徵學習傾向，也即“特徵汙染”，而這一傾向又和神經網路的 OOD 泛化有著直接的聯絡。

理論：神經網路的特徵汙染現象

在本節，我們在一個基於兩層 ReLU 網路的二分類問題上證明了“特徵汙染”現象的存在性，並且理論分析了這種現象的來源。具體而言，我們假定網路的輸入由兩種特徵線性組合而成：“核心特徵”和“背景特徵”。其中，核心特徵的分佈取決於類別標籤（可以形象理解為影像分類問題中的待分類物體），而背景特徵的分佈和標籤無關（可以形象理解為影像分類問題中的圖片背景等其他要素）。

為了排除其他因素的干擾，我們還對這兩類特徵作如下假設：

背景特徵和標籤不相關（這樣我們就排除了由虛假相關性導致的 failure mode）。
透過核心特徵可以對標籤實現 100% 準確率的預測（這樣我們就排除了由於訓練集的特徵不夠導致的 failure mode）。
核心特徵和背景特徵分佈在正交的子空間中（這樣我們就排除由於不同特徵難以解耦導致的 failure mode）。

我們發現，即使在以上的條件下，神經網路仍然會在學習核心特徵的同時學習和任務完全不相關的背景特徵。由於這兩種特徵在網路權重空間的耦合，在背景特徵上發生的分佈偏移也會導致神經網路的誤差增大，從而降低網路的 OOD 泛化性。我們也因此把這種神經網路的特徵學習偏好稱之為“特徵汙染”。以下，我們詳細介紹特徵汙染現象的出現原因。整體分析思路的示意圖如下：

我們分析中的關鍵點在於：特徵汙染實際上和神經網路中的神經元往往對不同類別具有不對稱的啟用

（activation）有關。

具體而言，我們可以證明在經過足夠的 SGD 迭代後，網路中至少有相當一部分的神經元都會被傾向於而與一個類別的樣本保持正相關（我們稱之為該神經元的正樣本，並用表示其類別），而與另外一個類別的樣本保持負相關（我們稱之為該神經元的負樣本，並用表示其類別）。這就會導致這些神經元的啟用具有類別不對稱性，如定理 4.1 所示：

這樣的類別不對稱性是怎麼影響神經網路的特徵學習過程的呢？我們首先注意到，對於網路隱層的第個神經元，其權重向量在第次迭代後可以被拆分為：

上式中，和分別表示核心特徵和背景特徵的集合，其中每個都對應一個核心特徵或者背景特徵。從該式中我們可以看出，神經元的權重可以分解為其在不同特徵上的投影（這裡我們假設不同的之間都是正交的單位向量）。進一步地，我們可以證明在的負梯度在每一個背景特徵上的投影滿足：

對於具有類別不對稱啟用的神經元，根據定理 4.1 我們有。換句話說，這樣的神經元的梯度主要取決於該神經元的正樣本而和負樣本幾乎無關。這就導致正樣本帶來的核心特徵和背景特徵會同時得到正的梯度投影，而這一過程和不同特徵與標籤之間的相關性無關。

如定理 4.2 所示，我們證明了在經過足夠的 SGD 迭代後，上面這種梯度投影的積累將導致神經元學習到的特徵既包含核心特徵，也包含耦合的背景特徵：

由於核心特徵和背景特徵在神經元權重中的耦合，背景特徵的負向分佈偏移會降低神經元的啟用，導致額外的 OOD 誤差。如定理 4.3 所示，我們定量描述了特徵汙染對 ID 和 OOD 泛化風險的影響：

同時，為了進一步說明特徵汙染源自神經網路的非線性啟用函式之間的關係，我們證明了在去除掉神經網路的非線性後，特徵汙染將不再發生：

如下圖所示，我們透過數值實驗驗證了我們的理論結果。同時，在兩層 ReLU 網路 + SGD 之外，我們也把我們的結論擴充套件到了更一般的設定上，包括其他種類的啟用函式、具有自適應步長的最佳化器等，結果如圖 3(d) 所示，表明特徵汙染在更一般的設定上也普遍存在。

同時，我們也提供了更多的實驗證據和特徵視覺化表明在我們日常使用的 ResNet、Vision transformer 等深度網路中，特徵汙染現象同樣會出現，並且能夠解釋我們實驗中觀測到的 OOD 泛化 gap。對這一部分內容感興趣的大家可以參考我們原論文的第 5 章。

總結與討論

最後，我們列舉一些未來可能比較重要/可以繼續深入做下去的研究點，也歡迎感興趣的大家和我們進一步交流:

更深層的網路：雖然我們從實驗上證明了深層網路也存在特徵汙染問題，但目前我們的理論分析還只做了兩層的 ReLU 網路。我們猜想特徵汙染可能是一個比較 general 的概念，並且神經元對於類別的啟用不對稱性可能只是其發生的原因之一。透過分析更加深層的網路或者更加複雜的網路結構（例如引入歸一化層等），我們或許可以發掘出引發特徵汙染的更多原因，並給出針對性的解決思路。
預訓練的作用：本文的理論分析只考慮了 train from scratch 的情況，但是我們實際使用的模型往往是預訓練模型。目前已有很多實驗證據表明預訓練是有助於提升模型的 OOD 泛化性的，那麼這種泛化性的提升的本質是否和緩解特徵汙染問題有關？預訓練又是如何做到這一點的？
怎麼解決特徵汙染問題：我們的工作雖然指出了特徵汙染問題，但還沒有給出比較明確的解決方案。不過，我們之後的一些工作已經表明，類似的問題在 fine-tuning 大模型的時候也會出現，並且我們也發現一些基於梯度調整的手段確實能夠緩解這個問題，從而顯著提升 fine-tuning 後的模型的泛化能力。關於這部分工作的具體內容我們未來也會放出，歡迎大家持續關注。

參考文獻

[1] Templeton, A., Conerly, T., Marcus, J., Lindsey, J., Bricken, T., Chen, B., Pearce, A., Citro, C., Ameisen, E., Jones, A., and Cunningham, H. Scaling monosemanticity: Extracting interpretable features from claude 3 sonnet. Transformer Circuits Thread, 2024.

[2] Shah, H., Tamuly, K., and Raghunathan, A. The pitfalls of simplicity bias in neural networks. In Advances in Neural Information Processing Systems, 2020.

[3] Pezeshki, M., Kaba, S.-O., Bengio, Y., Courville, A., Precup, D., and Lajoie, G. Gradient starvation: A learning proclivity in neural networks. In Advances in Neural Information Processing Systems, 2021.

[4] Scholkopf, B., Locatello, F., Bauer, S., Ke, N. R., Kalchbrenner, N., Goyal, A., and Bengio, Y. Toward causal representation learning. Proceedings of the IEEE, 109(5): 612–634, 2021.

[5] Gulrajani, I. and Lopez-Paz, D. In search of lost domain generalization. In International Conference on Learning Representations, 2021.

[6] Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, 2021.

[7] 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.

更多閱讀