不平衡數據下對比學習的理論分析：從訓練動態到剪枝解決方案

本文第一作者廖海旭為新澤西理工學院數據科學系在讀博士生，師從Prof. Shuai Zhang。

• 論文標題：Theoretical Analysis of Contrastive Learning under Imbalanced Data: From Training Dynamics to a Pruning Solution
• 論文鏈接：https://openreview.net/forum?id=DUXG9E8dEO

一、研究背景

對比學習已成為表征學習中的一種強大范式，能夠在不依賴標簽的情況下有效利用無標注數據。

在這一框架下，語義相似的樣本被視為正樣本對，而語義不同的樣本被視為負樣本對。通過在表征空間中拉近正樣本對、拉遠負樣本對之間的距離，對比學習使得模型能夠捕捉到豐富且具有判別性的特征。

該方法在廣泛的應用領域中取得了顯著成功，尤其在多模態學習中影響深遠，推動了早期視覺語言模型發展的重大進展。

盡管對比學習具有諸多優勢，但它在現實數據集中常見的類別不平衡的問題下，仍然面臨挑戰。在這類數據中，多數類主導了樣本對的構造，而少數類則代表性不足。這種不平衡會阻礙模型對少數類判別性特征的捕捉，降低表征質量，并導致模型產生偏差行為。

近期的研究開始逐步從理論角度理解對比學習，主要關注其相較于傳統生成式方法的優勢、數據增強在有效表征學習中的必要性，以及其在降低下游任務樣本復雜度方面的能力。然而，這些研究尚未考慮數據分布不平衡所帶來的影響，缺乏對這些影響的理論刻畫。

總體而言，本文在理論和實踐層面都提供了新的洞見，主要理論結論也通過數值實驗得到了驗證。主要貢獻如下：

第一，我們構建了一個理論框架，用以刻畫在數據分布不平衡條件下、基于 Transformer 編碼器的對比學習訓練動態。我們表明學習過程可以分為三個階段。

圖 1：神經元在訓練過程中投影的動態變化。藍色曲線表示神經元在其主導特征方向上的投影增長情況，橙色曲線表示其在非主導特征方向上的投影，綠色曲線表示其在噪聲空間方向上的投影。在第一階段，神經元主要沿特征方向增長，同時抑制噪聲分量。在第二階段，其在主導特征方向上的投影增長速度快于所有其他特征方向，從而形成明顯的分離。在第三階段，神經元逐漸收斂，其最終表示由所學習的主導特征所主導

第二，我們定量刻畫了少數特征的存在如何影響神經元的學習能力，進而影響整體表征學習。我們的分析表明不平衡會從多個方面削弱表征性能。

第三，基于幅值的剪枝能夠增強對少數特征的學習，恢復因數據不平衡而退化的性能。我們的結果顯示，剪枝可以加強沿少數特征方向的梯度更新，鼓勵更多神經元專門學習少數特征，從而獲得更加穩健且更均衡的表示。

二、剪枝算法

為了解決數據不平衡問題，我們在前向傳播過程中動態移除幅值較小的神經元權重，而在反向傳播過程中仍保留所有參數為可訓練狀態。

具體而言，二值掩碼初始設為全 1，表示訓練開始時不進行任何剪枝。在每一個訓練周期中，我們剪除幅值最小的一部分神經元，并相應更新二值掩碼。在前向傳播階段，使用掩碼后的參數對輸入進行編碼。在反向傳播階段，梯度是基于剪枝后的模型計算，但更新作用于完整的參數集合。

需要注意的是，該過程并不會為了提高效率而永久性地刪除任何神經元，盡管可以觀察到一定程度的計算成本下降。

三、關鍵發現的核心洞見

在介紹正式理論結果之前，我們首先總結分析所得的關鍵洞見。我們的研究表明，神經元在訓練過程中會分階段逐步學習特征表示。具體而言，我們得到以下結論：

(K1) 基于 Transformer-MLP 框架的對比學習訓練動態。我們的理論將學習過程劃分為三個階段。

• 在第一階段中，神經元權重沿特征方向增長，增長速率由特征頻率決定，神經元在非特征方向上的分量則被抑制。
• 在第二階段中，幸運神經元進一步和它學習的主導特征方向對齊；而普通神經元則被這些幸運神經元所界定并保持受控，從而使所學習的特征更加純凈，同時非特征分量繼續受到抑制。
• 在最終階段，訓練收斂，每個神經元都會與某一特定特征集合對齊，每個神經元在一個或多個特征方向上強對齊，在其他特征方向上弱對齊，并在非特征方向上保持較小幅度。

(K2) 特征頻率比率決定神經元專門化程度。在收斂時，每個神經元主要由某一特定特征集合中的特征主導，而來自其他方向的貢獻可以忽略。

• 首先，越稀有的特征學習得越弱。
• 其次，該特征集合的大小由特征頻率比率決定：較小的特征頻率比會擴大該集合，導致特征混合；較大的特征頻率比則會縮小該集合，使神經元學習到的特征更加純凈，這對于對比學習而言是更好的。
• 最后，專門學習單一特征的神經元數量與也由特征頻率比率決定，并且隨著特征頻率比的增大而減少。

由于對比學習在神經元專門學習單一特征時效果最佳，數據不平衡會帶來三個相互關聯的消極影響：

• 少數特征以較小幅度被學習。
• 神經元傾向于學習混合的多個特征而非保持與單個特征對齊。
• 專門學習單一特征的神經元的總數減少。

這些因素共同削弱了表征質量，并要求更大的模型規模才能學習所有特征。

(K3) 剪枝增強少數特征的學習。

• 學習少數特征的神經元會獲得更強更新；而學習非少數特征的神經元僅獲得的更新較弱。
• 在收斂時，學習少數特征的神經元的系數可以達到與多數特征同階的規模，從而緩解由數據不平衡帶來的性能下降。

直觀來看，由于少數特征神經元的幅值較小，它們更容易被剪枝；這種機制在梯度更新中放大了包含少數特征樣本的貢獻。因此，剪枝強化了少數特征，使其與其他成分明顯區分開來，并推動更多神經元專門化學習該特征，從而提升表示學習的魯棒性。

四、理論結果

理論分析結果的完整內容請參考文章，此處我們給出精簡的概括。

理論一：在第一階段的訓練中：

• 所有神經元都會沿著特征方向增長，而非特征方向則被忽略。
• 在每個特征方向上的增長速率取決于該特征的出現頻率。

理論二：在第二階段的訓練中：

• 幸運神經元：與單一特征強烈對齊。
• 普通神經元：其在特征上的分量受幸運神經元的約束。

理論三：當學習收斂時，在沒有剪枝的情況下，對比學習中神經元和特征的對齊情況如下：

• 每個神經元都會與一個或多個特征強對齊，與其他特征弱對齊，并在非特征方向上保持較小幅度。
• 每個神經元傾向于學習多個特征的混合表示，將學習更大規模的混合特征集合。

理論四：在引入剪枝的情況下，對比學習中神經元和特征的對齊情況如下：

• 沿少數特征方向的神經元更新得到增強。少數特征對應的神經元權重增長更快。非少數特征對應的神經元權重增長緩慢。

五、實驗結果

實驗結果表明，剪枝在所有數據集上均持續提升準確率，并且隨著不平衡程度加劇，性能提升更加顯著。此外，剪枝也改善了頭部類別與尾部類別之間的性能不平衡。這些結果表明，剪枝不僅提升了下游任務的整體性能，同時也縮小了頭部類別與尾部類別之間的性能差距。