對(duì)比學(xué)習(xí)(Contrastive Learning)是最近在表征學(xué)習(xí)(Representation Learning)中非常流行的學(xué)習(xí)范式。對(duì)比學(xué)習(xí)不同于一般的分類模型，考慮對(duì)數(shù)據(jù)進(jìn)行多個(gè)類別的分類，在對(duì)比學(xué)習(xí)中，我們通常需要考慮的只有兩類：正樣本和負(fù)樣本。舉個(gè)實(shí)際應(yīng)用系統(tǒng)的例子，在搜索系統(tǒng)中的QU相關(guān)性中[2]，我們通常需要判斷一個(gè)Query和一個(gè)Doc之間的相關(guān)性，如果相關(guān)的我們認(rèn)為是正樣本，而不相關(guān)的我們認(rèn)為是負(fù)樣本。目前很多模型都可以基于對(duì)比學(xué)習(xí)設(shè)計(jì)預(yù)訓(xùn)練任務(wù)，并且得到非?？捎^的性能提升，比如CLIP模型[3]。對(duì)于存有著海量用戶行為數(shù)據(jù)（比如用戶搜索的點(diǎn)擊數(shù)據(jù)，瀏覽數(shù)據(jù)，電商的購(gòu)買數(shù)據(jù)，停留時(shí)間等）的大型互聯(lián)網(wǎng)公司而言，在對(duì)比學(xué)習(xí)這個(gè)方向有著更大的探索需求。

“正樣本很多是正得相似，但是負(fù)樣本卻負(fù)得各有不同”，這句話雖然不太準(zhǔn)確，但是描述了對(duì)比學(xué)習(xí)中的本質(zhì)困難：如何去尋找合適的，更多的負(fù)樣本去進(jìn)行訓(xùn)練，從而讓模型學(xué)習(xí)出足夠能區(qū)分正負(fù)樣本的特征？在Kaiming大佬的MoCo[1]這篇工作中，大佬給出了他的解決方案。

對(duì)比學(xué)習(xí)中的海量負(fù)樣本需求在對(duì)比學(xué)習(xí)中，對(duì)負(fù)樣本有著海量的需求。在搜索場(chǎng)景中，我們會(huì)想盡辦法去構(gòu)造盡可能多的負(fù)樣本。比如將展現(xiàn)了但是并沒有被用戶點(diǎn)擊的樣本當(dāng)作是負(fù)樣本，比如將同為點(diǎn)擊了的樣本混搭組成batch negative負(fù)樣本。如Fig 1.1所示，通過對(duì)Query的特征和Doc的特征進(jìn)行組合，如式子(1.1)所示，可以構(gòu)建出打分矩陣S ，其中N 為batch size，D為特征大小。

那么可知道只有對(duì)角線的打分才是正樣本，而其他綠色部分的打分則是負(fù)樣本打分，如果負(fù)樣本的打分越高，那么就可以認(rèn)為是難負(fù)樣本。通過batch negative從一個(gè)batch內(nèi)構(gòu)造的負(fù)樣本，因?yàn)檫@些Doc都是被用戶點(diǎn)擊過的，而且Query都是被搜索過的，通過對(duì)他們進(jìn)行混搭構(gòu)成的負(fù)樣本足夠“真實(shí)”。而通過無點(diǎn)數(shù)據(jù)構(gòu)成的負(fù)樣本則可能因?yàn)槭钦宫F(xiàn)的Doc過多，已經(jīng)滿足了用戶的需求而沒有被點(diǎn)擊，因此可能出現(xiàn)將正樣本誤分為負(fù)樣本的情況。當(dāng)然，這個(gè)和具體的業(yè)務(wù)有關(guān)。

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

不管怎么說，通過以上構(gòu)造負(fù)樣本的方式，我們有能力構(gòu)造足夠多的負(fù)樣本。特別是通過batch negative的方式，在batch size為N的情況下我們能構(gòu)造出的負(fù)樣本，其中很可能就會(huì)出現(xiàn)足夠“好”的負(fù)樣本讓模型學(xué)習(xí)的好的表征。但是問題在于，我們受限于顯存的大小，我們無法無限增大batch size，即便通過多機(jī)多卡分布式的方法，每張卡上的batch size大小也是有極限的，即便通過某些機(jī)制去匯總多卡之間的負(fù)樣本打分，在匯總階段也會(huì)遇到瓶頸。因此目前能看到的公開工作中，據(jù)筆者了解，batch size最大的就是CLIP[3]，開到了32,768。然而如果繼續(xù)增大呢？可能就會(huì)遇到無法逾越的硬件瓶頸了。

Kaiming大佬在MoCo中，認(rèn)為對(duì)比學(xué)習(xí)可以從某種形式看成是“字典查詢”（Dictionary look-up），這一點(diǎn)不難理解，我們的正樣本可以認(rèn)為是分布在諸多的負(fù)樣本之中，而輸入一個(gè)Query，那么我們期望能匹配到正確的Doc（也是字典中的Key）。一次理想的匹配應(yīng)該是：query盡可能相似于正樣本的Key，而和其他負(fù)樣本的Key盡可能不相似。從這個(gè)角度上看，我們能發(fā)現(xiàn)對(duì)比學(xué)習(xí)的一個(gè)根本難題在于：

如何提高這個(gè)字典的詞表大小。如何提高字典的連續(xù)性（consistency）。在NLP任務(wù)中，特別是word2vec這種在所有有限詞表中選取負(fù)樣本的方法，因?yàn)檎Z言的詞表顯然是個(gè)離散有限集合，可能也就兩萬來個(gè)Key，因此對(duì)比學(xué)習(xí)在NLP中可能更好地進(jìn)行。而對(duì)于視覺任務(wù)而言，因?yàn)橐曈X信息是連續(xù)的，很難將其劃分為離散的有限Key集合將其覆蓋。因此視覺任務(wù)的對(duì)比學(xué)習(xí)會(huì)比NLP更難。

讓我們?cè)倏偨Y(jié)下，對(duì)比學(xué)習(xí)的兩個(gè)難題，第一個(gè)其實(shí)就可以任務(wù)是如何去增大負(fù)樣本的數(shù)量；第二個(gè)其實(shí)可以認(rèn)為是如何穩(wěn)定Query和Key的編碼器的一致性 （為了表述的通用性，我們將會(huì)用Key表示匹配任務(wù)中的Doc，Item端）。要知道，Query端和Key端的編碼器可以是一樣的，也可以是不一樣的；可以是共享參數(shù)的，也可以是獨(dú)立參數(shù)的。因此在參數(shù)更新過程中，要保證兩端編碼器狀態(tài)的一致性。 （本文考慮的是雙塔模型，因此交互式模型暫時(shí)不考慮。）

我們下面對(duì)已有的一些解決負(fù)樣本數(shù)量的方法進(jìn)行概述，以便于了解提出MoCo的歷史淵源。

端到端方式最為傳統(tǒng)的提高負(fù)樣本（也就是詞典的詞表）的方式就是提高Batch size大小，然后在更新參數(shù)階段同時(shí)對(duì)Query和Key的編碼器進(jìn)行更新。這種方式可以最大程度上保證Query和Key編碼器的一致性，但是顯然會(huì)遇到我們之前討論過的硬件瓶頸。

Fig 1.2 通過端到端的方式進(jìn)行對(duì)比學(xué)習(xí)。

Memory Bank一種直接的想法是將詞典詞表大小和batch大小進(jìn)行解耦，從而使得詞表不受限于batch大小。在[4]中，作者提出Memory Bank框架，該框架的主要特點(diǎn)在于維護(hù)一個(gè)隊(duì)列作為負(fù)樣本集合。在Memory Bank框架中，首先需要對(duì)數(shù)據(jù)集所有數(shù)據(jù)進(jìn)行刷特征向量，然后在每個(gè)mini batch中從中隨機(jī)挑選負(fù)樣本作為詞典，但是并不進(jìn)行梯度回傳，也就是不會(huì)對(duì)Key編碼器進(jìn)行參數(shù)更新。當(dāng)然，為了優(yōu)化特征，會(huì)對(duì)Query編碼器進(jìn)行參數(shù)更新，然后將當(dāng)前Query編碼器的特征入到負(fù)樣本隊(duì)列，并且將負(fù)樣本隊(duì)列中的舊特征出隊(duì)，從而實(shí)現(xiàn)對(duì)負(fù)樣本隊(duì)列的更新。這個(gè)負(fù)樣本隊(duì)列稱之為Memory Bank。

然而這個(gè)過程中，雖然這個(gè)方法滿足了“詞典詞表足夠大”的條件，但是“QK編碼器一致性”沒有保證，因?yàn)镵ey編碼總是落后于Query編碼的，這個(gè)過程中無法保證Query-Key匹配的一致性?？梢哉J(rèn)為Memory Bank提出的編碼器更新方式是“異步”逐漸更新的。

Fig 1.3 Memory Bank框架示意圖，主要通過維護(hù)一個(gè)負(fù)樣本隊(duì)列實(shí)現(xiàn)大尺度范圍的負(fù)樣本訓(xùn)練。

MoCokaiming大佬提出MoCo解決了以上的問題。首先進(jìn)行符號(hào)化表達(dá)，考慮一個(gè)Query q 和一系列編碼后的樣本，這些樣本是詞典的key，假設(shè)這個(gè)詞典中只有唯一的和q匹配，那么在對(duì)比學(xué)習(xí)中，只有q和匹配的時(shí)候，損失函數(shù)才是最低的。利用infoNCE (info Noise Contrastive Estimate)損失，我們可以對(duì)這個(gè)過程建模為(1.2)所示。

其中的是溫度系數(shù)，而K KK是負(fù)樣本個(gè)數(shù)。我們用表示query編碼器，表示key編碼器，那么有，用分別表示這兩個(gè)編碼器的參數(shù)。

Fig 1.4 MoCo框架示意圖。

為了保持大規(guī)模的詞典詞表大小，和Memory Bank一樣需要保持一個(gè)負(fù)樣本隊(duì)列，不同點(diǎn)在于對(duì)Key編碼器的更新方式不再采用“異步”的方式，而是采用了“同步”的方式。一種最為樸素的同步更新方式是：直接將當(dāng)前step更新好的Query編碼器的參數(shù)拷貝到Key編碼器中（當(dāng)然前提是QK編碼器結(jié)構(gòu)一致），如式子(1.3)所示。然而這種方式將會(huì)導(dǎo)致Key編碼器訓(xùn)練狀態(tài)的完全紊亂，因此無法訓(xùn)練將會(huì)一直震蕩，那么通過引入動(dòng)量進(jìn)行QK狀態(tài)協(xié)調(diào)就是一個(gè)自然而然的方案，如式子(1.4)所示。

引入動(dòng)量，將Key編碼器的參數(shù)更新，以一種合適地方式引入了Query編碼器的狀態(tài)。此處的m mm通常是一個(gè)比較大的值，比如0.99。從Fig 1.5中，我們可以看到不同大小的動(dòng)量對(duì)結(jié)果的影響，其中0.999是最為合適的。這也告訴我們，保持QK一致性很重要，同時(shí)保持Key編碼器訓(xùn)練的一致性也很重要。

Fig 1.5 不同大小的動(dòng)量大小對(duì)結(jié)果的影響。

作者在文章中還以偽代碼的形式舉了一個(gè)應(yīng)用例子，如Alg 1所示。

Alg 1. 以PyTorch的方式寫的MoCo流程偽代碼。

# f_q, f_k: encoder networks for query and key
# queue: dictionary as a queue of K keys (CxK)
# m: momentum
# t: temperature
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))
    # negative logits: NxK
    l_neg = mm(q.view(N,C), queue.view(C,K))
    # logits: Nx(1+K)
    logits = cat([l_pos, l_neg], dim=1)
    # contrastive loss, Eqn.(1)
    labels = zeros(N) # positives are the 0-th
    loss = CrossEntropyLoss(logits/t, labels)
    # SGD update: query network
    loss.backward()
    update(f_q.params)
    # momentum update: key network
    f_k.params = m*f_k.params+(1-m)*f_q.params
    # update dictionary
    enqueue(queue, k) # enqueue the current minibatch
    dequeue(queue) # dequeue the earliest minibatch

在這個(gè)例子中，作者將同一個(gè)圖片的不同數(shù)據(jù)增廣樣本當(dāng)成是正樣本，而其他所有的圖片都是負(fù)樣本。在代碼最后，保持了對(duì)隊(duì)列的更新（通過將新的minibatch入隊(duì)列和將舊的出隊(duì)列實(shí)現(xiàn)），并且通過動(dòng)量更新實(shí)現(xiàn)了對(duì)QK編碼器的更新。

作者對(duì)比了端到端，Memory Bank和MoCo框架的效果，如Fig 1.6所示，我們可以看到幾個(gè)值得注意的現(xiàn)象：

端到端的方式無法一直提高負(fù)樣本大小，可以發(fā)現(xiàn)MoCo和端到端在256-1024的負(fù)樣本大小階段的結(jié)果可以匹配（track）上的。MoCo和Memory Bank都可以將負(fù)樣本大小提高到很大。但是MoCo的效果一直都比Memory Bank高約2%點(diǎn)。這個(gè)現(xiàn)象也和之前的分析保持一致。

Fig 1.6 三種框架的實(shí)驗(yàn)結(jié)果對(duì)比。

Shuffling BN作者在文中提到了一嘴“Shuffling BN”，而這似乎是在本文才引出來的概念，我們?cè)谶@兒討論一下。在實(shí)踐中，研究者發(fā)現(xiàn)在對(duì)比學(xué)習(xí)中的編碼器使用Batch Normalization可能會(huì)使得模型學(xué)習(xí)不到好的表征[5]，之前的研究會(huì)采取不采用BN層的方法。在本工作中，作者認(rèn)為編碼器中的傳統(tǒng)BN層可能會(huì)影響預(yù)訓(xùn)練任務(wù)，進(jìn)而讓預(yù)訓(xùn)練很容易就找到了類似于平凡解的“低損失解”，而這將會(huì)使得預(yù)訓(xùn)練失效。

這個(gè)過程在“同步”更新的對(duì)比學(xué)習(xí)框架中容易出現(xiàn)，比如如Fig 1.7所示，其中BN層將會(huì)在N-F面上進(jìn)行標(biāo)準(zhǔn)化，我們知道一個(gè)batch中的Query和Key是一一匹配的，都是正樣本，也即是是匹配的。但是時(shí)卻不是匹配的，因此通過BN可能會(huì)將本是負(fù)樣本的信息“泄漏”到正樣本中去。而在Memory Bank中不會(huì)出現(xiàn)這個(gè)問題，因?yàn)镸emory Bank的正樣本Key是來自于過去的mini batch的，因?yàn)槭?ldquo;異步”過程所以不會(huì)影響當(dāng)前結(jié)果。

Fig 1.7 在對(duì)比學(xué)習(xí)中，BN層導(dǎo)致預(yù)訓(xùn)練任務(wù)失效的原因猜想。

為了解決這個(gè)問題，作者提出了Shuffling BN的概念，而這個(gè)層必須在多GPU卡上才能進(jìn)行。首先作者將不同GPU的sub mini batch進(jìn)行匯聚，然后對(duì)匯聚后的batch進(jìn)行打亂，并且記下打亂后的索引（便于恢復(fù)shuffle結(jié)果），在提取了Key編碼特征后，通過之前記下的索引，恢復(fù)shuffle結(jié)果。注意到這個(gè)過程只對(duì)Key編碼器進(jìn)行，而不考慮Query編碼器。正如作者提供的代碼code 1.1所示。

 # compute key features
    with torch.no_grad():  # no gradient to keys
        self._momentum_update_key_encoder()  # update the key encoder

        # shuffle for making use of BN
        im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)

        k = self.encoder_k(im_k)  # keys: NxC
        k = nn.functional.normalize(k, dim=1)

        # undo shuffle
        k = self._batch_unshuffle_ddp(k, idx_unshuffle)

總結(jié)

本文簡(jiǎn)單對(duì)MoCo進(jìn)行了討論，從實(shí)驗(yàn)上看，的確MoCo能有很大程度的提高，但是也正如作者文章最后討論的，并沒有完全利用到數(shù)據(jù)。如Fig 1.8所示，MoCo在100萬的數(shù)據(jù)和在十億級(jí)別的數(shù)據(jù)上性能差別似乎不是很大。而且，目前的MoCo動(dòng)量更新框架似乎要求QK編碼器的結(jié)構(gòu)需要一致？這一點(diǎn)不知道是否筆者理解正確了。

Fig 1.8 MoCo在IN-1M和IG-1B上的結(jié)果并沒有差別很大。

Reference

[1]. 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).

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

[3]. Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., … & Sutskever, I. (2021). Learning transferable visual models from natural language supervision. arXiv preprint arXiv:2103.00020.

[4]. 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).

[5]. Henaff, O. (2020, November). Data-efficient image recognition with contrastive predictive coding. In International Conference on Machine Learning (pp. 4182-4192). PMLR.