模態編碼器 | CLIP改進之SigLIP,采用sigmoid損失的圖文預訓練
DeepMind對CLIP改進的一篇工作--SigLIP,發表在2023CVPR。
簡單看下研究動機:傳統的對比學習方法如CLIP等依賴于 softmax 歸一化,這需要一個全局視角來計算成對相似度,從而限制了批處理大小的擴展能力,并且在小批處理大小下表現不佳。因此本文提出了一個簡單的成對 Sigmoid 損失函數用于語言-圖像預訓練(SigLIP)。
01、方法介紹
用于語言圖像預訓練的Softmax損失
CLIP等經典對比學習方法通常使用基于 softmax 的損失函數來對齊圖像和文本的表示。具體來說,給定一個包含圖像-文本對的小批量,對比學習的目標是使匹配對
的嵌入對齊,同時將不匹配對
的嵌入推遠。通常通過訓練一個圖像模型f(?) 和一個文本模型g(?)實現。損失函數如下:
其中:
由于 softmax 損失的不對稱性,規范化分別在圖像和文本上獨立進行。這種損失函數需要計算所有成對相似度的全局歸一化因子,在大規模批處理中會導致內存和計算開銷。
用于語言圖像預訓練的sigmoid損失
為了解決 softmax 損失的局限性,本文提出了一種基于 Sigmoid 的損失函數。Sigmoid 損失函數獨立處理每個圖像-文本對,避免了全局歸一化的需要。具體來說,Sigmoid 損失函數將學習問題轉化為標準的二分類問題,其中匹配對標記為正樣本,所有其他對
的嵌入推遠,標記為負樣本。Sigmoid 損失函數的形式如下:
其中:是標簽,
如果是匹配對,則
,否則
。
是 Sigmoid 函數,t和b是可學習的溫度和偏置參數,初始化為t=log10 和b=?10。
初始化時,由于負樣本數量遠多于正樣本,損失函數會被負樣本主導,導致初始優化步驟較大。為了緩解這一問題,引入了一個額外的可學習偏置項b,以確保訓練從接近先驗的地方開始,避免需要巨大的矯正步驟。
高效“分塊”實現
在對比訓練中,通常利用數據并行性來加速計算。當數據被分割到D個設備上時,計算損失函數需要收集所有嵌入向量,這通常涉及到昂貴的全聚集操作,并且需要在內存中構建一個大規模的成對相似性矩陣,這在內存消耗上是非常密集的。
- Sigmoid損失的優勢:sigmoid損失函數不需要計算全局歸一化因子,這使得它特別適合于內存高效、快速且數值穩定的實現。這種方法可以減少數據并行處理中的挑戰。
- 分塊實現方法:作者提出了一種“分塊”方法,避免了全聚集操作,并且減少了內存使用。具體來說,每個設備的批量大小表示為
,損失函數被重新表述為:
D為設備數,是sigmoid損失函數
02、計算過程
計算過程簡單來說,首先計算與正對對應的損失分量,然后對設備間的表示進行排列,使得每個設備從其鄰近設備中獲取負樣本。每個設備獨立地計算其本地批量b的損失,并將所有設備的損失簡單地求和。
著重介紹下SigLIP是如何進行分布式訓練的,假設全局的batch size為B,一共有個GPU,那么每個GPU上的batch size為
,可以將Sigmoid 損失函數表示公式的損失拆解為上述公式所示,在Fig 1展示了整個過程的示意圖,在初始化階段,以第一個GPU為例子,其所包含的樣本為:
此時GPU 1可以完成一次上述公式中C的計算,然后,交換GPU 1和GPU 2的文本編碼器特征向量,即:
此時GPU 2完成一次上述公式中的B計算,以此類推,直到GPU 1遍歷完所有樣本為止,其他GPU也是如此操作的,最終把所有卡上的損失聚集即可,也就是A計算。這個輪流交換不同GPU之間數據的操作,可以稱之為排列。不難發現,整個過程的通信成本來自于排列,一共需要D?1次聚集操作即可完成一次排列。
而在CLIP中需要對圖文的編碼器特征都進行聚集,因此需要2次全聚集
操作。如果全聚集
采用ring方法的話,那么一個全聚集
就是D?1次聚集
操作。由此得出一個SigLIP和CLIP的性能復雜度對比:
這種方法顯著減少了內存消耗,從減少到
,并且由于b通常是固定的,通過增加加速器的數量來擴展|B|。這種方法使得在相對較少的設備上訓練超過一百萬個批量大小的模型成為可能。
03、實驗結果
batch_size影響
實驗結果顯示,在batch size小于32k的時候,采用sigmoid的SigLIP的性能都會優于采用softmax的CLIP。在batch size足夠大的時候,CLIP能夠追上,甚至超越SigLIP的表現,但是最佳性能仍然是SigLIP@32k情況下得到,從實際應用來看,采用SigLIP能用更少的資源更快的訓練速度得到更好的性能。從SigLiT的實驗來看,隨著batch size的尺度放大性能將會在32k batch size的情況下達到飽和,同時能觀察到SigLiT在不同batch size下持續優于LiT。繼續提高batch size將不能帶來更好的收益,甚至會有細微的性能下降。
訓練時間影響
另外,作者在SigLiT的設置下,訓練了更長時間(也即是見了更多數據量),結果顯示,在超大batch size,如262k的情況下,對比較小batch size(如8k)提供更長的訓練時間的確能觀察到性能的較大提升。并且也能觀察到在超大batch size下,采用sigmoid和采用softmax的區別很小,但是在較小batch size(如8k)的情況下,差距明顯。
因此,在資源受限的情況下,采用SigLIP更劃算,所需資源少而且性能更強。同時,這個實驗也說明了,超大的batch size并不意味著訓練得更快,反而還需要更長的訓練時間。
多語言數據集表現
在多語言數據集上,性能同樣在32k batch size上達到了飽和
有限資源場景
在僅有四塊 TPU-v4 芯片的情況下使用 SigLiT 模型的訓練效果如下,在四個芯片上訓練一天的結果顯示,模型在ImageNet上的零樣本分類準確率達到了79.7%,這在資源有限的情況下非常有競爭力。使用ViT-g/14模型作為視覺塔和大型文本塔,可以在兩天內以20k的批量大小進行107k步的訓練,進一步將零樣本分類準確率提高到84.5%。
另外還探討了在資源有限的情況下,特別是只有少量TPUv4芯片時,如何有效地訓練SigLIP模型。結果顯示,使用16個TPUv4芯片,以16k的批量大小訓練了2.4B個樣本,僅靠微調并不能立即提高模型的性能。這與之前的研究一致,即微調圖像模型會降低視覺表示的質量。然而,通過凍結預訓練權重的權重衰減,可以獲得更好的結果。
兩個微調經驗
微調圖像模型會降低視覺表示的質量。然而,通過凍結預訓練權重的weight-decay,可以獲得更好的結果
隨著批量大小的增加,訓練過程變得越來越不穩定。實驗發現降低Adam和AdaFactor中的β2參數值(從默認的0.999降低到0.95)可以穩定訓練過程。
sigmoid損失函數中正負樣本比例的影響
對于sigmoid來說,它的loss是以pair為粒度計算的,positive和negative非常不平衡。以batch_size=16k為例,positive和negative的比率約為1:16k
文中嘗試了4種掩碼策略來保留最難的負樣本(損失最高)或最易的負樣本(損失最低)
- 隨機(Random):隨機選擇負對進行掩碼。
- 難(Hard):保留最難的負對。
- 易(Easy):保留最易的負對。
- 難+匹配總對數(Hard + matching total pairs seen):在固定數量的步驟中掩碼樣本會減少訓練期間看到的總對數。因此,在匹配對的設置中,作者增加了掩碼比例的步數,以保持看到的對數不變。
不做matched pair的情況下,用3種mask方式均會造成精度下降。影響程度:easy>random>hard。
結果顯示,隨機移除負樣本以重新平衡會惡化性能。保留最易的樣本根本不起作用,而保留最難的負樣本幾乎可以保持質量。通過保留最難的負樣本,模型的性能得到了維持,這表明在負樣本中,較難的樣本對學習更為重要。
另外,還觀察了學習到的偏置項的最終值以及正負樣本的平均logit值,發現當正負樣本的imbalance減弱時,learnable bias和pair的logit都在上升,說明了預設的learnable bias起到了積極的作用。
數據噪聲對模型魯棒性的影響
通過以下五種方法“污染”訓練噪聲:
- Image:以概率p將圖文對的圖片用均勻噪聲替換;
- Text:以概率p將圖文對的文本token序列用隨機采樣的等長token序列替換;
- Batch alignment: 隨機將batch中的p%的sample的圖文pair進行shuffle;
- Image & text: 同時進行1.和2.
- Image, text & batch: 同時進行3和4
從結果可見,sigmoid loss在“污染”數據的性能更好。
04、總結
總的來說,SigLIP是一個很棒的工作,作者采用sigmoid損失去取代對比學習中的softmax函數,以更小的資源開銷帶來了更好的模型表現,目前也被很多多模態大模型所采用,作為視覺端的編碼器。