自監(jiān)督對比學習的泛化性理論

自監(jiān)督學習范式已經(jīng)在各種 CV 和 NLP 等任務(wù)上超越了直接用監(jiān)督學習的性能。盡管近來自監(jiān)督學習的算法層出不窮，但是對它泛化性好的理論理解仍然十分缺乏，只能盲目嘗試，因此限制了自監(jiān)督學習進一步的發(fā)展。本工作嘗試對自監(jiān)督對比學習的泛化性進行一些理論上的理解。

懶人版：下面視頻的前 18 分鐘

【2022.03 北大報告】華為諾亞黃維然：自監(jiān)督學習理論

https://www.zhihu.com/zvideo/1483773806724997120

自監(jiān)督學習 Self-Supervised Learning 是目前人工智能領(lǐng)域大熱的一個方向，圖靈獎得主 Yann Lecun 和 Yoshua Bengio 都認為自監(jiān)督學習是機器智能達到人類水平的關(guān)鍵[1]。這種學習范式已經(jīng)在各種 CV 和 NLP 等任務(wù)上超越了直接用監(jiān)督學習的性能。盡管近來自監(jiān)督學習的算法層出不窮，但是對它泛化性好的理論理解仍然十分缺乏，只能盲目嘗試，因此限制了自監(jiān)督學習進一步的發(fā)展。我們最近的工作嘗試對自監(jiān)督對比學習的泛化性進行一些理論上的理解。論文地址如下：

Towards the Generalization of Contrastive Self-Supervised Learning

https://www.weiranhuang.com/publications/contrastive2021.html

我們對自監(jiān)督對比學習的泛化性進行分析，提出了一個統(tǒng)一的理論框架來解釋對比自監(jiān)督學習泛化性好的原因。我們對唯一的自監(jiān)督信息——數(shù)據(jù)增強，進行數(shù)學上的定量刻畫，并基于此推導出自監(jiān)督對比學習在下游任務(wù)的泛化誤差上界。根據(jù)該誤差上界，得出泛化性好的三個條件：正樣本對齊 Alignment of positive samples、類中心錯開 Divergence of class centers、增廣數(shù)據(jù)的聚焦性 Concentration of augmented data。其中，前兩個條件是可以通過對比學習算法進行優(yōu)化的——因此我們證明了兩個常見算法 SimCLR 和 Barlow Twins 都在隱式地滿足前兩個條件；而第三個條件只依賴輸入的訓練數(shù)據(jù)——因此我們在真實數(shù)據(jù)集上測試，發(fā)現(xiàn)用我們理論預(yù)測出不同訓練數(shù)據(jù)的泛化性變化趨勢能跟真實性能變化趨勢一致。

自監(jiān)督學習是通過大量無標簽數(shù)據(jù)來預(yù)訓練一個模型，學到數(shù)據(jù)表示 representation，然后根據(jù)不同的下游任務(wù)，把學到的表示接在不同的分類器 / 回歸器上，并用下游任務(wù)的有標簽數(shù)據(jù)對分類器 / 回歸器進行訓練，從而適用于下游任務(wù)。自監(jiān)督對比學習首先把 n 個訓練數(shù)據(jù)進行兩次數(shù)據(jù)增廣得到 n 對“正樣本”，然后把“正樣本對”在特征空間中拉近從而學到特征提取器。由于只拉近正樣本會導致學到常數(shù)平凡解（所有樣本被映射到特征空間的同一個點），為了避免塌縮，不同的損失函數(shù)（比如 SimCLR/MoCo 等算法的 InfoNCE 引入負樣本并將其遠離正樣本；Barlow Twins 等算法的互相關(guān)損失把特征的各個分量解耦）和訓練方法（比如 BYOL/SimSiam 引入 stop-gradient）被提出。

盡管自監(jiān)督對比學習已經(jīng)被實驗廣泛驗證了它在下游任務(wù)的泛化性，但是它的理論分析還十分有限。有一些工作從互信息最大化的角度理解 InfoNCE，指出 InfoNCE 的相反數(shù)是正樣本互信息的下界，因此最小化 InfoNCE 其實是在最大化互信息的下界。然而，最新的結(jié)果[2]指出，如果去優(yōu)化更緊的互信息下界，并不能得到更好的表示。這說明通過最大化互信息的角度也許并不一定能夠完全解釋 InfoNCE 的成功。

Arora 等人[3]提出一個理論框架來直接分析 InfoNCE 在下游任務(wù)的泛化性，但是他們基于一個假設(shè)：正樣本是從同一個潛在類別 latent class 中選取的，比如兩張不同的狗的圖片作為一對正樣本。而在實際中，一對正樣本是同一個圖片的兩個不同的增廣，不同的圖片（即使屬于同一個潛在類別）是作為負樣本的，因此這個假設(shè)和實際情況出入較大。Wang 等人[4]提出 alignment 和 uniformity 作為預(yù)測下游性能的經(jīng)驗性指標，并沒有嚴格的泛化性理論保證。

同時，上述兩個工作都避開了去分析數(shù)據(jù)增廣在自監(jiān)督對比學習中的角色，然而它卻是自監(jiān)督對比學習有效的關(guān)鍵——因為僅有的人類先驗知識是通過數(shù)據(jù)增廣構(gòu)造正樣本來注入學習過程的（即人類可以識別出增廣后的樣本來自同一個圖片，說明人類認知這兩個增廣后的圖片所用的特征在這個兩個圖片中都具備，而自監(jiān)督對比學習正是通過把上述人類經(jīng)驗教給機器從而學到好的特征）。

另外，Haochen 等人[5]從矩陣分解的角度來分析自監(jiān)督對比學習，但是他們的理論分析僅僅適用于他們的 spectral contrastive loss，無法分析廣泛引用的 SimCLR 或者 Barlow Twins 的損失函數(shù)。

除了上述現(xiàn)有自監(jiān)督對比學習泛化性理論的局限性之外，還有一些有趣的實驗現(xiàn)象無法完全解釋。

• 更加豐富的數(shù)據(jù)增廣能夠讓學到的特征空間具有更好的聚類性質(zhì)。上圖從左到右分別是 (a) 訓練前的特征空間、(b) 只用一種類型的數(shù)據(jù)增廣學到的特征空間、(c) 使用多種類型的數(shù)據(jù)增廣學到的特征空間?？梢钥吹?(a) 幾乎看不出聚類，(b) 可以看出紅色和藍色有一定的聚類，但是還有很大的重疊區(qū)域，而 (c) 聚類性質(zhì)則十分明顯。而由于更好的聚類性質(zhì)往往代表更好的下游分類精度（因為下游任務(wù)只訓一個線性分類器），因此在下游任務(wù)上的表現(xiàn) (c) 是優(yōu)于 (b) 的。SimCLR 論文中也觀測到類似的現(xiàn)象——更加豐富的數(shù)據(jù)增廣具有更有的下游任務(wù)表現(xiàn)。直觀上這個現(xiàn)象似乎是理所當然的，這能否從理論上給出定量刻畫？此外，SimCLR 論文也做了數(shù)據(jù)增廣兩兩組合的實驗，發(fā)現(xiàn) crop&color 效果最好。這能否有理論解釋？
• 樣本層面的拉近為什么會得到類別層面的聚類？比如 SimCLR 在優(yōu)化的時候，僅僅讓同一樣本的兩個數(shù)據(jù)增廣靠近，而把同一潛在類別中其它樣本視作該樣本的負樣本從而被推開，這等價于是讓 2n 個增廣后的樣本做 n 分類，n 是訓練數(shù)據(jù)量。但是在 SimCLR 學到的特征空間中，樣本是按類別顯著聚類在一起的（比如在 cifar-10 上的結(jié)果就是 10 類而不是 n 類）。這一點是如何保證的？
• 前面提到如果特征空間按類別聚類越好，下游任務(wù)的性能通常也越好。但是 Barlow Twins 并沒有直接優(yōu)化特征空間的幾何性質(zhì)讓它聚類，而是優(yōu)化特征的相關(guān)性性質(zhì)（將特征的各個分量解耦），這為什么能保證特征空間在幾何上有聚類的效果？

本文我們主要聚焦在理論分析自監(jiān)督對比學習的泛化性 (Motivation 1)，有了我們的理論結(jié)果，就可以輕松地解釋上述提到的實驗現(xiàn)象 (Motivation 2)。首先，最重要的一點就是理解并用數(shù)學刻畫出數(shù)據(jù)增廣的作用。

直觀來看，同一個類別的樣本很可能具備相似的增廣圖片。比如上圖中的兩個狗的圖片（綠框和藍框），它們在像素層面是完全不同的（不同的姿勢，不同的背景等等）。但是如果我們對它們做數(shù)據(jù)增廣比如隨機裁剪 random crop，那么它們有可能得到在像素層面十分接近的兩個狗頭（灰框）。兩個相似的圖片，它們的特征也應(yīng)當接近 (Lipschitz 連續(xù))，對應(yīng)特征空間中的兩個灰點。自監(jiān)督對比學習的目標是把正樣本在特征空間中拉近，因此以灰點為錨點，綠點和藍點都往灰點靠攏，最終這兩個狗的圖片（綠框和藍框）在特征空間聚到一起。

根據(jù)上述想法，我們定義一個新的兩個圖片之間距離叫做 augmented distance，它是兩個圖片在各自的增廣中各選取一個增廣，使得它們的像素距離最小。從上圖的例子可以看出，augmented distance 從某種程度上刻畫了兩個圖片的語義距離。緊接著, 我們可以定義 -增廣, 即超過  部分的樣本, 它們之間兩兩 augmented distance 都小于  。當  越大或者  越小的時候, 代表更加 sharp 的 concentration, 見下圖。

有了上述-增廣來定量刻畫數(shù)據(jù)增廣，我們通過理論推導出自監(jiān)督對比學習泛化性好的三個條件：

1. Alignment of positive samples：這是不同自監(jiān)督對比學習算法的共同目標，即正樣本在特征空間中要被拉得越近越好；
2. Divergence of class center：不同類中心在特征空間里離得足夠遠，這保證了學到的表示不會塌縮到平凡解。這跟 Wang 等人[4]提出的 uniformity 并不相同，后面 3.1 節(jié)會具體討論。
3. Concentration of augmented data：數(shù)據(jù)增廣的聚焦性要好 sharp concentration，即越大的 \sigma\sigma 和越小的 \delta\delta 。

注意到，其中 1 和 2 兩個條件是可以通過設(shè)計好的對比學習算法來滿足的，但是條件 3 不依賴于算法。為了驗證上述 3 個條件，我們分析了 SimCLR 和 Barlow Twins 這兩個算法的 loss，并從理論上證明他們均可以滿足 1 和 2 兩個條件；為了驗證條件 3，我們通過實驗選取不同的數(shù)據(jù)增廣，來驗證越 sharp 的 concentration 確實能帶來更好的下游任務(wù)效果。

我們用粗體的  來表示一個（隨機）樣本, 由它進行數(shù)據(jù)增廣后的圖片集合記作 , 那么 InfoNCE loss 可以寫成：

而 Barlow Twins 中的 cross-correlation loss 可以寫成：

其中  且  對每一維  是歸一化的。

標準的自監(jiān)督學習算法的評估方法是在學到的表示后面接一個線性分類器，并在下游數(shù)據(jù)上對線性分類器進行訓練，由此得到的下游任務(wù)分類器，并以它的性能來作為評估學到的表示的好壞。為了簡化數(shù)學分析，我們考慮不需要訓練的最近鄰分類器作為下游任務(wù)的分類器，即

其中  是第  類  的類中心。事實上，最近鄰分類器可以看作是線性分類器的特例（見原文），因此實際中用線性分類器得到的下游任務(wù)性能會高于最近鄰分類器。我們用誤差率來測量最近鄰分類器的性能，即

有了上述定義，下面我們研究為何自監(jiān)督對比學習可以得到一個小的誤差率。

我們首先給出上面 1.3 Main Idea 中 augmented distance 的數(shù)學定義：

以及 -增廣的嚴謹定義：

當我們有一個更豐富的數(shù)據(jù)增廣  時, 根據(jù) augmented distance 定義, 我們有 , 因此當  不變時，  會變小。也就是說，更豐富的數(shù)據(jù)增廣會帶來更加 sharp 的 concentration。后面我們會看到更加 sharp 的 concentration 意味著更好的下游性能, 由此我們可以得出, 更豐富的數(shù)據(jù)增廣帶來更好的下游性能。

為了刻畫正樣本拉近的程度, 我們定義集合 里面的元素的兩個增廣都足夠接近。進一步定義  越小表示 alignment 越好。

注：我們可以證明  可以用 alignment 來 upper bounded 如下, 具體細節(jié)見原文。

有了數(shù)據(jù)增廣的定義，下面我們不加證明地直接給出本文的第一個定理：

為了更直觀的理解上述定理, 我們首先來看一個簡單的例子：任何兩個來自同一個潛在類別的樣本 都存在一個相同的增廣圖片 , 而且通過自監(jiān)督對比學習正樣本都 perfect aligned , 見下圖：

在這種情況下, 來自同一個潛在類別的樣本都被映射到了同一個點上, 因此我們只需要任意小的角度就能把不同的類別都準確區(qū)分開, 也就是  。事實上, 經(jīng)過簡單計算我們有 , 根據(jù)定理 1 得到, 這和我們直觀分析的結(jié)論完全一致。

從定理 1 可以看出自監(jiān)督性能好  即  小）的三個條件:

1. Alignment of positive samples：這也是不同自監(jiān)督對比學習算法優(yōu)化的共同目標。由定理 2 可 知 alignment 越小  越小, 再由定理 1 可知  也越小;
2. Divergence of class centers：不同的類中心的距離要足夠大, 也就是足夠小的  。從定理 1 可以看到, divergence 和 alignment  、concentration  都有關(guān)。越好的 alignment 和越 sharp 的 concentration（從而更小的  ) 會放松 divergence 的條件。這也可以從上圖中的簡單例子看出：perfect alignment 和 prefect concentration 下, 只需任意小的 divergence 即可。
3. Concentration of augmented data：當  固定的時候, 越 sharp 的 concentration 意味著越大的 , 從而直接降低了  的上界。

定理 1 提供了一個理論框架在分析不同的自監(jiān)督學習算法的泛化性。上述條件 1 和 2 可以通過算法來優(yōu)化（后面第 4 節(jié)我們將通過兩個具體的例子 SimCLR 和 Barlow Twins 來揭示這一點），而條件 3 不依賴于算法，是由預(yù)先給定的增廣數(shù)據(jù)來決定（后面第 5 節(jié)我們將通過實驗來驗證這一點）。

兩個工作中都涉及 alignment，它們具有相同的含義，因為這是不同自監(jiān)督對比學習算法優(yōu)化的共同目標；但是我們提到的 divergence 和他們提出的 uniformity 在概念上完全不同。

• Uniformity 需要所有的數(shù)據(jù)點均勻的分布在特征球面上；我們的 divergence 刻畫的是類中心之間的 cosine distance。我們不要求所有的數(shù)據(jù)點盡可能散開分布，相反，我們希望同類別的數(shù)據(jù)點盡量接近。同時，好的 alignment 和 concentration 可以放松對 divergence 的要求。因此，我們的 divergence 是一種更加準確和精細的刻畫。
• Alignment 和 uniformity 是一種用來預(yù)測下游性能的經(jīng)驗指標，而我們的 alignment 和 divergence 是由定理 1 導出的，具有泛化性的理論保證。
• 另外，當 concentration 較差時（比如用恒等變換作為數(shù)據(jù)增廣的方法），即使有 perfect alignment 和 perfect uniformity，下游性能仍然很差，這再次說明 alignment 和 uniformity 指標的局限性。

自監(jiān)督對比損失函數(shù)一般可以分解成兩部分，一部分是為了拉近正樣本，另一部分是正則項為了保證不會學到塌縮解，即

其中, 第一項也就是 alignment, 即  。不同的損失函數(shù)的差別主要在于第二項, 一個有效的正則項應(yīng)當讓 divergence 滿足定理 1 的條件。原文中的本節(jié)針對不同 loss 的數(shù)學分析較多, 在這里我們直接給出最后的結(jié)論:

首先, 我們證明了兩種常見算法 SimCLR 和 Barlow Twins 的損失函數(shù)確實可以被分解成上面相關(guān)的兩部分  和 ; 其次, 我們可以證明, 對于  有 , 對于 Barlow Twins 有 ; 最后, 我們給出下游性能的理論保證, 對于 SimCLR 有 , 對于 Barlow Twins 有

對于 SimCLR 的 InfoNCE loss，把它拆成 pos 和 reg 兩項后得到

其中第二項正則項是 LogExp 形式，這對于避免塌縮很關(guān)鍵。如果用線性函數(shù)來替換第二項（有些文獻成為 simple contrastive loss），

由于  和  的獨立性, 第二項可以化簡成 , 因此最小化第二項只能要求學到的表示均值為零。但是均值為零加上  無法避免表示發(fā)生維度塌縮 dimensional collapse（維度塌縮的概念最早由 @清華MARS Lab [6] 提出）, 比如所有的數(shù)據(jù)點被映射到特征球面上相對的兩個點上, 如下圖：

這也和 Wang 等人[7]發(fā)現(xiàn)的現(xiàn)象吻合：當 InfoNCE 的溫度增加時，uniformity 變差（正則項變差）。因為當溫度趨于正無窮時，InfoNCE 退化到上述 simple contrastive loss。

Barlow Twins 的 loss 被設(shè)計用來做特征分量解耦，那它為什么能滿足 alignment 和 divergence 呢？

首先，把它拆成兩項

那么

經(jīng)過簡單的計算，不難得到

也就是說，互相關(guān)矩陣的對角元其實決定了 alignment。

當 alignment 較好時, 可以預(yù)期 , 于是第二項中, 而對于每個類別中的不同樣本基本都聚集在所屬類中心的周圍, 因此可以用  來近似代替 , 于是上述過程可以總結(jié)成

可以看到  距離我們想要的 divergence  只有一步之遙了。事實上, 利用矩陣 trace 的輪換性, 我們就能得到

這樣就能保證 divergence 得到滿足。

我們主要驗證：更 sharp 的 concentration 預(yù)示著更好的下游性能。因此我們設(shè)計三組實驗：

1）在第 3 節(jié)的定義 1 下面我們論證了更豐富的數(shù)據(jù)增廣能推導出更 sharp 的 concentration，因此第一組實驗我們觀測越來越多的數(shù)據(jù)增廣種類對下游性能的影響。我們考慮 5 種增廣類型：(a) random cropping (b) random Gaussian blur (c) color dropping (d) color distortion (e) random horizontal flipping。

可以看到，隨著數(shù)據(jù)增廣種類越來越豐富，下游性能越來越好，且不同算法下趨勢一致。有趣的點是加入 (c) 增廣之后，下游性能有明顯的提升。我們回過頭去看 (c) 是什么操作，是把彩色圖片變成灰度圖片。由于這個操作讓圖片在像素層面有劇烈變動，因此由 augmented distance 的定義可以預(yù)期該距離會變小很多，從而導致 concentration 變得 sharp。一個直觀的例子，在 Figure 2 中，如果把右邊那個狗的圖片替換成哈士奇（黑狗），那么僅僅做裁剪操作并不能讓兩個狗獲得像素層面相似的增廣圖片，所以 augmented distance 較大，concentration 較差。而當我們引入 color dropping 操作之后，我們可以獲得兩個黑灰色的狗頭，這樣我們又可以得到像素層面相似的增廣圖片，從而能把它們的原始圖片拉近。這時對應(yīng) augmented distance 變小，concentration 變得 sharp。

2）當我們固定數(shù)據(jù)增廣的類型，而改變增廣的強度，那么更強的數(shù)據(jù)增廣會導致更小的 augmented distance，從而獲得更 sharp 的 concentration 和 更好的下游性能。因此，第二組實驗，我們固定 (a) 和 (d) 兩種增廣的種類，改變 (d) 的強度，觀測對下游性能的影響。

可以看到，隨著數(shù)據(jù)增廣強度的增加，下游性能越來越好，且不同算法下趨勢一致。

3）上面兩組實驗雖然一定程度上能夠驗證我們的結(jié)論，但是有人可能覺得實驗結(jié)果是理所當然的。第三組實驗，我們固定數(shù)據(jù)增廣種類的個數(shù)為二，嘗試所有可能的兩兩組合，并直接計算出不同組合的 ，從而觀測它與下游任務(wù)的相關(guān)性。

可以看到, 下游性能  和 concentration (  當固定  時) 具有高度相關(guān)性：越 小的  代表越 sharp 的 concentration, 從而有更小的下游誤差率。如果我們固定一個增廣為 (a), 觀測它與其他增廣的組合, 可以看到不論是下游性能 (綠線) 還是 concentration (紫線、橙 線) 都具有  的排序。另外, 在全部的組合中, concentration 最 sharp 的 組合  成功預(yù)測出它也是下游性能最好的組合, 這也和 SimCLR 論文中的發(fā)現(xiàn) crop&color 是 最有效的組合相吻合。

參考

1. ^https://cacm.acm.org/news/244720-yann-lecun-yoshua-bengio-self-supervised-learning-is-key-to-human-level-intelligence/fulltext
2. ^Michael Tschannen, Josip Djolonga, Paul K Rubenstein, Sylvain Gelly, and Mario Lucic. On mutual information maximization for representation learning. arXiv preprint arXiv:1907.13625, 2019.
3. ^Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:1902.09229, 2019.
4. ^abcTongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning, pages 9929–9939. PMLR, 2020.
5. ^Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. Advances in Neural Information Processing Systems, 34, 2021.
6. ^Tianyu Hua, Wenxiao Wang, Zihui Xue, Sucheng Ren, Yue Wang, and Hang Zhao. On feature decorrelation in self-supervised learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9598–9608, 2021.
7. ^Feng Wang and Huaping Liu. Understanding the behaviour of contrastive loss. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2495–2504, 2021.