在大尺度的對(duì)比學(xué)習(xí)訓(xùn)練過(guò)程中，Batch Norm層所造成的過(guò)擬合現(xiàn)象被眾多公開(kāi)論文所報(bào)道過(guò)，比如MoCo[3], SimCLR[4]和其他一些工作[5]。之前筆者在MoCo的筆記中也簡(jiǎn)單談到過(guò)這個(gè)問(wèn)題[1]，然而當(dāng)時(shí)尚未深究，現(xiàn)在在工作中實(shí)際遇到了這個(gè)問(wèn)題，就權(quán)當(dāng)筆記將其紀(jì)錄。

在大尺度的對(duì)比學(xué)習(xí)中，一種常見(jiàn)的實(shí)踐是：設(shè)置一個(gè)較大的batch size，比如4096，顯然一張GPU卡很難塞下，特別是在多模態(tài)模型中，因此通過(guò)數(shù)據(jù)并行將大batch size均分到不同卡上，比如16張卡。在雙塔模型中，我們需要對(duì)兩個(gè)塔輸出的特征進(jìn)行計(jì)算得到打分矩陣，如Fig 1所示。然而分布在不同卡上的雙塔特征的批次大小為256，如果對(duì)每個(gè)卡上的雙塔特征進(jìn)行打分，那么得到的打分矩陣 ??？梢钥吹酱藭r(shí)多GPU并沒(méi)有真正意義上增大batch size，只是通過(guò)數(shù)據(jù)并行的方式更快地遍歷了訓(xùn)練數(shù)據(jù)集，提高了訓(xùn)練速度而已，我們期望中的打分矩陣應(yīng)該是 ?，在這種大尺度的打分矩陣上能更高效地獲得足夠好的負(fù)樣本進(jìn)行對(duì)比學(xué)習(xí)。

Fig 1. Batch Negative的方式從一個(gè)batch中構(gòu)造負(fù)樣本。

一種實(shí)現(xiàn)這種方式的實(shí)踐是通過(guò)各個(gè)深度學(xué)習(xí)框架提供的all_gather機(jī)制，比如pytorch中的torch.distributed.all_gather() [6]或者paddle中的paddle.distributed.all_gather() [7]。這些函數(shù)可以從所有的GPU中匯聚某個(gè)矩陣，比如特征矩陣，將所有的GPU中的特征矩陣進(jìn)行匯聚，可以得到 的匯聚特征矩陣。我們可以對(duì)這個(gè)進(jìn)行打分，然后進(jìn)行對(duì)比學(xué)習(xí)的訓(xùn)練。

然而，這只是大尺度對(duì)比學(xué)習(xí)的第一步，我們發(fā)現(xiàn)在匯聚特征之前，特征的計(jì)算都是在各自的GPU中進(jìn)行計(jì)算的，假如模型中具有Batch Norm層，那么其統(tǒng)計(jì)參數(shù)????都是在各自的GPU中計(jì)算的（假設(shè)是異步Batch Norm機(jī)制），和Batch Norm有關(guān)的知識(shí)可見(jiàn)之前博文[8]。而由于BN層的統(tǒng)計(jì)參數(shù)和all_gather機(jī)制，會(huì)導(dǎo)致在大尺度對(duì)比學(xué)習(xí)訓(xùn)練過(guò)程中的嚴(yán)重過(guò)擬合現(xiàn)象。然而B(niǎo)N的統(tǒng)計(jì)參數(shù)導(dǎo)致的過(guò)擬合問(wèn)題并不只在存在all_gather機(jī)制的對(duì)比學(xué)習(xí)模型中存在，注意到MoCo看成是維護(hù)了一個(gè)負(fù)樣本隊(duì)列[1]，因此可以視為不采用all_gather機(jī)制，也可在單卡上進(jìn)行超大batch size的訓(xùn)練。然而MoCo也會(huì)遇到BN層的統(tǒng)計(jì)參數(shù)泄露信息的問(wèn)題。且讓筆者慢慢道來(lái)。

一般來(lái)說(shuō)，提高負(fù)樣本數(shù)量的方法有以下幾種：

端到端，此時(shí)通過(guò)all_gather 機(jī)制可以擴(kuò)大batch size，進(jìn)而擴(kuò)大負(fù)樣本數(shù)量。 在這種方式下，負(fù)樣本數(shù)量和batch size耦合。MoCo，這種方法通過(guò)負(fù)樣本隊(duì)列和動(dòng)量更新保證了Query-Key編碼器的狀態(tài)一致性和足夠大的負(fù)樣本詞表。 在這種方式下，負(fù)樣本數(shù)量和batch size解耦。Memory Bank [10]，此時(shí)通過(guò)維護(hù)一個(gè)負(fù)樣本數(shù)據(jù)庫(kù)，稱之為memory bank進(jìn)行，然而此時(shí)的Query-Key編碼器不是一致?tīng)顟B(tài)的，Key編碼器永遠(yuǎn)落后于Query編碼器。

我們留意到，并不是所有方法都會(huì)收到BN層信息泄露問(wèn)題的，只有在（正樣本對(duì)）Query-Key編碼器一致更新的模型中才會(huì)遇到，而在Memory Bank中就不會(huì)遇到。其中，我們先討論端到端的形式中的BN層信息泄露。注意，我們這里說(shuō)的狀態(tài)一致，或者一致更新，并不是指的數(shù)值上的一致，而是假使存在一個(gè)訓(xùn)練狀態(tài)，這兩者是同步的。

端到端模式的對(duì)比學(xué)習(xí)

在端到端模式的對(duì)比學(xué)習(xí)過(guò)程中，Query-Key編碼器是一致更新的，簡(jiǎn)單來(lái)說(shuō)就是兩個(gè)塔的參數(shù)在同個(gè)step中進(jìn)行更新。此時(shí)如果采用了多卡進(jìn)行all_gather，并且采用的BN層是異步BN（也就是每張卡的統(tǒng)計(jì)參數(shù)和學(xué)習(xí)參數(shù)都是不同的，分別是獨(dú)立計(jì)算的），那么可想而知，通過(guò)all_gather之后，其形成的打分矩陣如Fig 3所示。Fig 3中的不同顏色塊表示來(lái)自于不同GPU上的正樣本Query-Key對(duì)，省略號(hào)表示的是通過(guò)不同卡匯聚得到的特征進(jìn)行打分。我們注意到由于異步BN的原因，不同顏色塊上的統(tǒng)計(jì)參數(shù)是不相同的，而正樣本顯然又位于打分矩陣的對(duì)角線上，正樣本都由同一個(gè)GPU進(jìn)行計(jì)算，此時(shí)由gather得到的諸多負(fù)樣本的統(tǒng)計(jì)參數(shù)會(huì)和同一個(gè)GPU下的正樣本的統(tǒng)計(jì)參數(shù)存在明顯差別。由BN的計(jì)算公式(1)可知，不僅通過(guò)學(xué)習(xí)表征，通過(guò)『學(xué)習(xí)』統(tǒng)計(jì)參數(shù)也可以『等價(jià)』于學(xué)習(xí)表征，讓模型『預(yù)測(cè)』正樣本的位置。然而，這種『等價(jià)』并不是真正的等價(jià)，而『預(yù)測(cè)』也不是通過(guò)真正學(xué)習(xí)表征得到的，因此表現(xiàn)為過(guò)擬合，嚴(yán)重影響模型的表征性能，這個(gè)情況筆者稱之為BN層統(tǒng)計(jì)參數(shù)泄露。在這種情況下，由于統(tǒng)計(jì)參數(shù)泄露了『正樣本所在于對(duì)角線』這個(gè)秘密，導(dǎo)致表征學(xué)習(xí)以失敗告終。

Fig 3. 端到端模式下的all_gather將會(huì)導(dǎo)致BN層統(tǒng)計(jì)參數(shù)的泄露。

從以上的分析來(lái)看，在端到端模式下導(dǎo)致統(tǒng)計(jì)參數(shù)泄露的本質(zhì)還是在于統(tǒng)計(jì)參數(shù)都是在各自的GPU中進(jìn)行計(jì)算的，那么解決方案自然要從這里著手。在simCLR[4]中，作者提出的方案是采用所謂的Global BN，其方法就是同樣gather不同GPU上的統(tǒng)計(jì)參數(shù)，然后計(jì)算出一個(gè)新的統(tǒng)計(jì)參數(shù)后分發(fā)到所有GPU上，此時(shí)所有GPU的統(tǒng)計(jì)參數(shù)都是相同，也就談不上泄露了。當(dāng)然你還可以用更簡(jiǎn)單的方法，比如在[5]中，作者采用Layer Norm取代了Batch Norm。從Fig 4.中可以看出，Layer Norm進(jìn)行統(tǒng)計(jì)參數(shù)計(jì)算的維度是[Feature, Channel]，而不涉及Batch維度，統(tǒng)計(jì)參數(shù)不會(huì)跨Batch使得統(tǒng)計(jì)參數(shù)不會(huì)泄露樣本之間的信息。

這個(gè)方法相當(dāng)?shù)刂庇^，因?yàn)樽罾硐氲那闆r下，我們應(yīng)該對(duì)所有的正負(fù)樣本一個(gè)個(gè)地送到編碼器中，以達(dá)到完全隔離不同樣本之間的目的，通過(guò)將BN替換成LN，達(dá)到了這個(gè)目的。

Fig 4. Layer Norm的統(tǒng)計(jì)參數(shù)不會(huì)跨batch內(nèi)的樣本，因此不會(huì)泄露統(tǒng)計(jì)參數(shù)。

MoCo還有一種非?；鸬膶?shí)踐是何凱明大佬的MoCo[3,1]，這種方式不僅需要維護(hù)一個(gè)大尺度的負(fù)樣本隊(duì)列，還需要用動(dòng)量更新的方式去一致更新Query-Key編碼器，如Fig 5.所示。此時(shí)的負(fù)樣本數(shù)量的提升不是由于all_gather機(jī)制導(dǎo)致的，并且負(fù)樣本數(shù)量和batch size也是解耦了。因此這種情況下，我們認(rèn)為即便是單GPU也可以跑很大的負(fù)樣本數(shù)量的對(duì)比學(xué)習(xí)。那么此時(shí)在上文所說(shuō)的BN層統(tǒng)計(jì)參數(shù)泄露問(wèn)題在MoCo中存在嗎？

Fig 5. MoCo的QK編碼器更新由于采用了動(dòng)量更新，因此是狀態(tài)一致的。

很遺憾，即便在MoCo中，BN層的統(tǒng)計(jì)參數(shù)泄露還是存在的，但是原因顯然不是由于多GPU的異步BN統(tǒng)計(jì)參數(shù)導(dǎo)致的，因?yàn)榧幢阒挥幸粡埧ㄒ部梢岳碚撋吓躆oCo。如Fig 6.所示，此時(shí)正樣本打分的計(jì)算如Code 1.所示，是通過(guò)對(duì)某個(gè)樣本進(jìn)行數(shù)據(jù)增廣后，將其視為正樣本，再進(jìn)行打分。這個(gè)部分也就是Fig 6.的藍(lán)色部分。而負(fù)樣本是直接采用Key-負(fù)樣本隊(duì)列中的特征，直接和Queryt特征進(jìn)行打分，如Fig 6.的綠色部分所示。注意到代碼中的k = k.detach(), 這意味著構(gòu)造出來(lái)的正樣本的梯度流只會(huì)更新Query編碼器，而Key編碼器是通過(guò)動(dòng)量更新的。

Code 1. MoCo中的正樣本打分計(jì)算。

f_k.params = f_q.params # initialize
for x in loader: # load a minibatch x with N samples
    x_q = aug(x) # a randomly augmented version
    x_k = aug(x) # another randomly augmented version
    q = f_q.forward(x_q) # queries: NxC
    k = f_k.forward(x_k) # keys: NxC
    k = k.detach() # no gradient to keys
    # positive logits: Nx1
    l_pos = bmm(q.view(N,1,C), k.view(N,C,1))
    # 正樣本打分計(jì)算
    l_neg = mm(q.view(N,C), queue.view(C,K))
    # 負(fù)樣本打分計(jì)算
    # logits: Nx(1+K)
    logits = cat([l_pos, l_neg], dim=1)
    labels = zeros(N) # positives are the 0-th
    loss = CrossEntropyLoss(logits/t, labels)
    loss.backward()
    update(f_q.params)
    enqueue(queue, k) # enqueue the current minibatch
    dequeue(queue) # dequeue the earliest minibatch

Fig 6. 在MoCo中采用組織超大規(guī)模負(fù)樣本隊(duì)列的方式進(jìn)行負(fù)樣本組織，因此batch size可以不需要很大，并且理論上可在單卡中完成計(jì)算。

即便在MoCo中不通過(guò)正樣本構(gòu)造的k進(jìn)行key編碼器的更新，但是同個(gè)step中，將會(huì)通過(guò)式子(2)進(jìn)行動(dòng)量更新，因此QK編碼器的更新也是一致的。

注意到此時(shí)的正樣本是在同一狀態(tài)下的Query-Key編碼器中被編碼的，如Fig 7.所示，其中的正樣本Query和正樣本Key的統(tǒng)計(jì)特征是整個(gè)batch內(nèi)共享的，分別是 和? ??????。這意味著整個(gè)batch內(nèi)的所有正樣本都共享著相同的統(tǒng)計(jì)特征，而且兩者的狀態(tài)保持一致，此時(shí)的統(tǒng)計(jì)特征將會(huì)很容易泄露正樣本的『位置』。筆者理解是，此時(shí)因?yàn)榻y(tǒng)計(jì)特征都一樣，而且正樣本的位置都在第一列，那么模型會(huì)嘗試『欺騙』整個(gè)訓(xùn)練任務(wù)，通過(guò)更簡(jiǎn)單地『學(xué)習(xí)』統(tǒng)計(jì)特征的分布進(jìn)而『預(yù)測(cè)』正樣本的位置。畢竟統(tǒng)計(jì)特征的維度更小，而表征的維度高得多，因此在某些數(shù)據(jù)集上可能就會(huì)出現(xiàn)BN層的統(tǒng)計(jì)參數(shù)泄露的問(wèn)題。

Fig 7. MoCo中的正樣本對(duì)打分。因此，雖然MoCo和端到端的方式中數(shù)據(jù)組成方式大相徑庭，但是BN層的統(tǒng)計(jì)特征在QK編碼器一致更新的過(guò)程中，都存在泄露正樣本的位置的可能。為了解決這個(gè)問(wèn)題，在MoCo中采用的是Shuffling BN，分為以下幾個(gè)步驟，注意Shuffling BN必須運(yùn)行在多卡環(huán)境下：

將輸入進(jìn)行all_gather，并且進(jìn)行隨機(jī)打亂，此時(shí)需要記下打亂后的索引unshuffle_idx，因?yàn)樽詈笮枰悍创騺y』回一開(kāi)始的樣本順序。將打亂后的樣本平均分發(fā)到M???個(gè)GPU上，每個(gè)GPU上的batch size大小為。通過(guò)Key編碼器計(jì)算特征，注意到Key編碼器中存在BN層，比如經(jīng)典的resnet結(jié)構(gòu)。將所有GPU上計(jì)算得到的特征進(jìn)行all_gather，并且通過(guò)unshuffle_idx進(jìn)行反打亂回原來(lái)的樣本順序。整個(gè)邏輯的示意圖如Fig 8所示。整個(gè)過(guò)程其實(shí)就是通過(guò)打亂和分發(fā)到多個(gè)GPU，實(shí)現(xiàn)統(tǒng)計(jì)參數(shù)的打亂，解耦統(tǒng)計(jì)參數(shù)和正樣本位置的關(guān)聯(lián)。Shuffling BN通過(guò)打亂實(shí)現(xiàn)隨機(jī)，而Global BN通過(guò)gather->分發(fā)的方式實(shí)現(xiàn)統(tǒng)計(jì)參數(shù)的全局統(tǒng)一。一個(gè)是『打亂』一個(gè)是統(tǒng)一，這些手段都保證了BN的統(tǒng)計(jì)參數(shù)不會(huì)帶有正樣本的位置信息。

Fig 8. MoCo中的Shuffling BN邏輯示意圖。Memory Bank

Fig 9. Memory Bank的示意圖。

然而，在memory bank [10, 12]中卻并不會(huì)出現(xiàn)BN層統(tǒng)計(jì)參數(shù)泄露的問(wèn)題，那是因?yàn)閙emory bank是通過(guò)『異步』的方式取正樣本key的。具體而言，如Fig 9.所示memory bank維護(hù)了一個(gè)負(fù)樣本隊(duì)列?，假如負(fù)樣本詞表大小為10000，而輸出的embedding特征維度為128，那么模型一開(kāi)始將會(huì)初始化一個(gè)?的隨機(jī)矩陣作為負(fù)樣本的初始值，記作。假如編碼器為將對(duì)第i ii??個(gè)樣本?進(jìn)行編碼，輸出??。在[10]這篇文章中，作者將參數(shù)化的交叉熵?fù)p失(3)替換成了非參數(shù)化形式的交叉熵?fù)p失(4)，原因很簡(jiǎn)單，在這種對(duì)比學(xué)習(xí)場(chǎng)景中，一個(gè)樣本除了和自己可以看成是正樣本外，其他的樣本都視為負(fù)樣本。這意味著當(dāng)訓(xùn)練數(shù)據(jù)足夠大時(shí)（通常對(duì)比學(xué)習(xí)的數(shù)據(jù)集都可以非常大），將有著數(shù)不勝數(shù)的負(fù)樣本，此時(shí)如果交叉熵?fù)p失是參數(shù)化的形式，那么權(quán)值矩陣?將變得大到無(wú)法計(jì)算。

而我們?nèi)绻捎孟嚓P(guān)性計(jì)算的方式計(jì)算樣本之間的距離，那么就可以形成如式子(4)所示的非參數(shù)化的交叉熵?fù)p失（其中的是溫度系數(shù)，通?？梢栽O(shè)置為常數(shù)），通過(guò)NCE（Noise Contrastive Estimation）可以進(jìn)一步減少負(fù)樣本太多造成的計(jì)算負(fù)擔(dān)。

因此，按照這個(gè)邏輯，此時(shí)我們計(jì)算正樣本的相關(guān)性時(shí)候，采用式子(5)

此時(shí)的 i+表示輸入樣本的索引i ii在V \mathcal{V}V 中取到的第i?個(gè)特征?。是的，在memory bank中，除了當(dāng)前輸入的特征之外，其他的正樣本Key亦或是負(fù)樣本Key都是從memory bank中取得的。而memory bank中的每一個(gè)樣本表征都是從之前的編碼器計(jì)算后插入得到的，這意味著無(wú)論是正樣本Key還是負(fù)樣本Key的更新?tīng)顟B(tài)都是落后于當(dāng)前的編碼器的。因此正樣本Key的統(tǒng)計(jì)特征早已落后于當(dāng)前編碼器中的BN層統(tǒng)計(jì)特征了，此時(shí)如果模型還想通過(guò)統(tǒng)計(jì)特征去『欺騙』訓(xùn)練過(guò)程將變得很困難。

Fig 10. 非參數(shù)化后的交叉熵?fù)p失的正樣本Key和負(fù)樣本Key都來(lái)自于memory bank，而memory bank的狀態(tài)是落后于當(dāng)前的編碼器的。

說(shuō)在最后

對(duì)于對(duì)比學(xué)習(xí)的研究，在學(xué)校中，我們沒(méi)有足夠的計(jì)算資源和超大型的數(shù)據(jù)量，因此比較難發(fā)現(xiàn)在大batch size甚至是超大batch size下才會(huì)出現(xiàn)的BN層統(tǒng)計(jì)參數(shù)泄露問(wèn)題。目前在學(xué)術(shù)界有報(bào)道這個(gè)問(wèn)題的文章?lián)P者了解也就[3,4,5]這幾篇。然而在公司的實(shí)踐中，在面對(duì)數(shù)以億計(jì)的大量數(shù)據(jù)時(shí)，簡(jiǎn)單粗暴地提高batch size將導(dǎo)致意料外的結(jié)果，因此筆者將其進(jìn)行筆記，希望對(duì)讀者有所幫助。此外，以上結(jié)論并不一定在所有數(shù)據(jù)集上都成立，我們發(fā)現(xiàn)數(shù)據(jù)集的特性也很重要，如果讀者在相同實(shí)踐中也遇到了對(duì)比學(xué)習(xí)的過(guò)擬合問(wèn)題，不妨也可以往著B(niǎo)N層統(tǒng)計(jì)參數(shù)泄露這方面考慮。

Reference

[1]. https://fesian.blog.csdn.net/article/details/119515146

[2]. https://fesian.blog.csdn.net/article/details/119516894

[3]. He, K., Fan, H., Wu, Y., Xie, S., & Girshick, R. (2020). Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 9729-9738).

[4]. Chen, T., Kornblith, S., Norouzi, M., & Hinton, G. (2020, November). A simple framework for contrastive learning of visual representations. In International conference on machine learning (pp. 1597-1607). PMLR.

[5]. Hénaff, O. J., Razavi, A., Doersch, C., Eslami, S., and Oord, A.v. d. Data-efficient image recognition with contrastive predictive coding. arXiv preprint arXiv:1905.09272, 2019.

[6]. https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_gather

[7]. https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/distributed/all_gather_cn.html

[8]. https://blog.csdn.net/LoseInVain/article/details/86476010

[9]. Memory Bank Code: https://github.com/zhirongw/lemniscate.pytorch

[10]. Wu, Z., Xiong, Y., Yu, S. X., & Lin, D. (2018). Unsupervised feature learning via non-parametric instance discrimination. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 3733-3742).

[11]. https://github.com/facebookresearch/moco/blob/master/moco/builder.py

[12]. https://github.com/zhirongw/lemniscate.pytorch