一篇解決表示學(xué)習(xí)坍塌問題的工作報告

1 表示學(xué)習(xí)中的坍塌問題

在表示學(xué)習(xí)中，一種很常見的做法是利用孿生網(wǎng)絡(luò)的結(jié)構(gòu)，讓同一個樣本的不同數(shù)據(jù)增強(qiáng)后的表示相似。比如對于一張圖像，使用翻轉(zhuǎn)、裁剪等方法生成另一個增強(qiáng)圖像，兩個圖像分別經(jīng)過兩個共享參數(shù)的編碼器，得到表示，模型的優(yōu)化目標(biāo)是讓這兩個表示的距離近。

這種方法一個比較大的挑戰(zhàn)在于，模型在訓(xùn)練過程中容易出現(xiàn)坍塌問題。模型可以把所有樣本的表示都學(xué)成完全相同的常數(shù)向量（比如每個樣本模型的輸出都是全0向量），這樣就能滿足上述兩個表示距離近的要求了。

2 解決坍塌問題的方法

業(yè)內(nèi)有很多解決坍塌問題的方法，主要可以分為4種類型：

對比學(xué)習(xí)方法：在訓(xùn)練一對正樣本對時，同時采樣大量的負(fù)樣本，讓正樣本之間離得近，負(fù)樣本之間離得遠(yuǎn)，避免模型偷懶把所有樣本的表示都學(xué)成一樣的。

聚類方法：在訓(xùn)練過程中增加一個聚類過程，將樣本分配給不同的類簇，然后在類簇級別進(jìn)行對比學(xué)習(xí)。

基于蒸餾的方法：通過模型結(jié)構(gòu)的角度避免坍塌問題，學(xué)習(xí)一個student network來預(yù)測teacher network的表示，teacher network是student network參數(shù)的滑動平均，teacher network不通過反向傳播更新參數(shù)。

信息最大化方法：讓生成的embedding中每一維的向量相互正交，使其信息量最大化，這樣可以避免各個維度的值信息過于冗余，防止坍塌問題。

目前常用的負(fù)樣本采樣方法，一個比較大的問題是計算開銷大，取得好的效果往往需要大量負(fù)樣本，因此有了MoCo等對比學(xué)習(xí)框架。本文提出的方法基于信息最大化的思路，能夠只使用正樣本對實現(xiàn)表示學(xué)習(xí)的同時，防止坍現(xiàn)象的發(fā)生。關(guān)于對比學(xué)習(xí)的常用經(jīng)典方法，可以參考這篇文章：對比學(xué)習(xí)中的4種經(jīng)典訓(xùn)練模式。

3 信息最大化方法歷史工作

基于信息最大化的方法典型的工作有兩篇ICLM 2021的文章，分別是Whitening for Self-Supervised Representation Learning（ICML 2021，W-MSE）和Barlow twins: Self-supervised learning via redundancy reduction（ICML 2021）。下面介紹一下這兩篇文章的整體思路，F(xiàn)acebook的這篇論文也是基于這個思路設(shè)計的。

W-MSE的整體網(wǎng)絡(luò)結(jié)構(gòu)如上圖，模型中輸入一對互為正樣本的樣本對（例如一個圖像的不同增強(qiáng)形式），使用共享參數(shù)Encoder分別編碼后，增加一個whitening模塊，對每個batch內(nèi)的所有embedding進(jìn)行白化，讓embedding的各個維度變量線性無關(guān)，后面再接norm處理。下圖形式化表明了W-MSE的用途，通過白化+norm讓樣本形成一個球形分布，正樣本之間距離近，每個樣本需要調(diào)整自己在圓周上的位置拉進(jìn)正樣本之間的距離，形成最終的簇。

Barlow twins方法也是同樣的思路，只是做法不同。在得到一對正樣本的兩個embedding后，計算兩個向量各個變量之間的相關(guān)性矩陣，優(yōu)化這個矩陣接近對角矩陣，以此實現(xiàn)embedding中各個變量之間線性無關(guān)。

4 VICREG

Facebook團(tuán)隊提出的VICREG是上述基于信息最大化表示學(xué)習(xí)方法的一個擴(kuò)展，損失函數(shù)主要包括variance、invariance、covariance三個部分。模型的主體結(jié)構(gòu)如下，仍然是一個孿生網(wǎng)絡(luò)結(jié)構(gòu)，輸入一個樣本的不同view互為正樣本。

variance部分主要是讓batch內(nèi)每個樣本embedding向量的每一維變量之間的方差不小于一個值，這樣可以有效防止每個樣本的向量都對應(yīng)同一個點的情況，防止坍塌的發(fā)生。

invariance部分即讓正樣本對的表示embedding距離盡可能小，是正常的表示學(xué)習(xí)loss。

covariance借鑒了Barlow twins中的思路，讓batch內(nèi)的embedding非對角線元素盡可能為0，即讓表示向量各個維度變量線性無關(guān)。

最終的表示學(xué)習(xí)loss是由上述3個loss的加和：

5 實驗結(jié)果

作者進(jìn)行了一些實驗驗證VICREG表示學(xué)習(xí)方法的效果。在ImageNet數(shù)據(jù)集上，使用不同方法預(yù)訓(xùn)練的表示接一個線性層進(jìn)行分類，或者基于預(yù)訓(xùn)練的表示使用部分?jǐn)?shù)據(jù)進(jìn)行finetune，對比圖像分類效果，實驗結(jié)果如下：

作者也對比了在不同方法上是否采用variance或covariance等loss的效果變化：

6 總結(jié)

本文介紹了Facebook AI團(tuán)隊在2022年ICLR的一篇解決表示學(xué)習(xí)坍塌問題的工作，順帶介紹了基于信息最大化的防坍塌工作，有助于幫助我們更深層次理解表示學(xué)習(xí)的原理，以及如何解決實際應(yīng)用表示學(xué)習(xí)、對比學(xué)習(xí)時遇到的坍塌問題。

審核編輯：劉清