不平衡數(shù)據(jù)下對比學習的理論分析：從訓練動態(tài)到剪枝解決方案

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

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

一、研究背景

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

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

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

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

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

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

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

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

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

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

二、剪枝算法

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

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

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

三、關鍵發(fā)現(xiàn)的核心洞見

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

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

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

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

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

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

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

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

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

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

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

四、理論結果

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

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

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

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

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

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

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

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

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

五、實驗結果

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