同濟提出簡化 Transformer結構:在RTX3090上實現CLIP的輕量級訓練 !
作者丨小書僮
來源丨集智書僮
編輯丨極市平台
導讀
同濟大學提出的一種簡化Transformer結構的方法,通過在僅使用一塊Nvidia RTX3090 GPU和一兆字節存儲空間的情況下實現競爭性能。

對比語言圖像預訓練(CLIP)因其卓越的零樣本性能和優秀的遷移能力而受到了廣泛關注。然而,訓練這樣的大型模型通常需要大量的計算和存儲,這對一般用戶(擁有消費級計算機)來說是一個障礙。
為了應對這一觀察,本文探討了如何在僅使用一塊Nvidia RTX3090 GPU和一兆字節存儲空間的情況下實現競爭性能。
一方面,作者簡化了 Transformer 塊結構,並將權重繼承與多階段知識蒸餾(WIKD)相結合,從而減少了參數並提高了訓練和部署期間的推理速度。另一方面,面對小數據集帶來的收斂挑戰,作者為每個樣本生成合成字幕作為數據增強,並設計了一種新穎的配對匹配(PM)損失,以充分利用正負圖像文本對之間的區分。
大量實驗表明,作者的模型可以在數據量-參數-精度之間實現新的最先進權衡,這可能進一步在相關研究社區中普及CLIP模型。
1 Introduction
預訓練的大型圖像文本基礎模型,如對比語言圖像預訓練(CLIP)模型[28],最近在計算機視覺和自然語言處理領域引起了廣泛關注。這些模型在廣泛的下遊任務上表現出出色的零樣本性能和魯棒性,例如圖像文本檢索和分類(朱等,2023年)。然而,CLIP類模型的巨大計算和存儲成本阻礙了它們的進一步普及。例如,MobileCLIP [33]在256xA100 GPU上訓練,全局批量為65,536,相應的數據集 DataCompDR-1B需要140 TB的本地存儲空間。此外,巨大的參數大小(例如,CLIP-B/16模型[28]包含86.2M個圖像編碼器參數和63.4M個文本編碼器參數)導致了推理延遲的增加,這為部署在計算資源有限的設備上帶來了挑戰。這些缺陷為沒有足夠計算資源和數據集的一般用戶參與大規模模型的訓練和部署設置了障礙。
在實際應用中,消費級計算機的GPU內存通常不超過24GB(例如,Nvidia RTX 3090),存儲容量可能小於1TB。在這樣資源限制的背景下訓練CLIP類似的模型,需要解決兩個主要問題。首先,必須儘可能減少需要訓練的參數數量,同時保留儘可能多的現有模型知識。其次,小型數據集需要適當擴充,並需要開發更有效的方法,充分利用圖像文本對在有限樣本內的內部關聯。
在這篇論文中,作者研究了如何使用僅有一個RTX3090顯卡和1TB存儲的輕量級CLIP模型訓練方法,從而普及在消費級計算機上的CLIP類模型研究。為此,作者首先提出將傳統的Transformer塊簡化為SAS-P塊,並採用權重共享策略。然後,通過繼承現有模型的權重並提取知識,可以進一步減少訓練所需的參數數量。在數據集方面,作者選擇廣泛使用的CC12M [1]作為基礎。該數據集不僅規模較小,而且標籤質量低,這都為模型訓練過程的收斂帶來了困難。為瞭解決這個問題,作者對CC12M中的每個圖像樣本添加多個文本標籤,創建了新的CC12M-SYN。此外,為了從這樣的小數據集中提取有價值的信息,作者引入了Pair Matching (PM)損失,以幫助模型捕捉正負圖像-文本對之間的區分。這些方法在作者的廣泛實驗中顯著提高了模型訓練的收斂速度。最後,通過在38個數據集上的性能比較(如圖1所示),作者的提出的SiCLIP框架實現了新的數據量-參數-準確率權衡的最先進水平。

作者的貢獻:本工作的貢獻可以概括如下:
作者提出了一個系統性的框架,用於在消費級計算機上訓練輕量級CLIP模型,包括數據集構建和相應的訓練過程,稱為SiCLIP。在SiCLIP中,計算和存儲成本得到了降低,同時保持了與其他大規模模型競爭的性能。
作者通過在SAS-P塊之間共享權重來簡化CLIP模型的結構,並將權重繼承與多階段知識蒸餾(WIKD)相結合,從而在訓練和部署過程中降低了內存需求。
一種名為PM損失的新損失函數被設計出來,該函數在訓練過程中預測圖像-文本對是否匹配。結合作者增強的CC12M-SYN數據集,PM損失可以利用正負圖像-文本對之間的區分。實驗結果顯示,新的數據集和PM損失都可以顯著提高訓練效率,同時略微增加數據集大小。
2 Related Work
Efficient Training for CLIP
自從CLIP作為一種大規模圖像文本基礎模型,在各種下遊任務上展現出驚人的零樣本性能以來,旨在提高其訓練效率和模型大小的相關研究已經有很多。例如,包括細粒度圖像文本對齊 [42],數據增強 [23, 20, 33],單模態自監督 [23, 20],以及圖像文本標籤空間中的對比學習 [39]。此外,Zhai等人 [45] 提出了一種對偶的Sigmoid損失作為對比損失的簡單替代方案,在小型批量訓練時證明了其有效性。然而,由於需要計算所有正負圖像文本對之間的匹配對數,其可能導致二次計算複雜度。Li等人 [19] 使用了細粒度圖像文本匹配(ITM)損失作為對比損失的補充,但ITM需要一個基於多層 Transformer 的編碼器來編碼多模態細粒度特徵,這並不適合輕量級模型。
基於權重繼承(WI)和知識蒸餾(KD)[13]的方法也被採用以實現高效的訓練。TinyCLIP[36]通過跨模態親和力模擬和WI訓練緊湊的CLIP模型。楊等人[38]探討了不同的KD方法對CLIP訓練的有效性。
高質量的數據集對於有效的訓練也非常重要。方等人在[6]中利用過濾方法去除噪聲樣本,Gadre等人[7]也提出了同樣的方法。然而,賸餘的描述性仍然不夠。最近的研究[40; 18]表明,從預訓練的描述生成模型中合成的合成描述可以提高數據集質量。
Simplifying the Transformer Architecture
近年來,隨著 Transformer 在各種任務中的顯著成功,許多努力都致力於簡化 Transformer 架構以提高其訓練和推理效率。餘等人[44]證明了 Transformer 塊的一般結構對其性能更為關鍵,因此可以消除基於注意力的 Token 混合器,這些混合器通常由於多頭自注意力(MHSA)在長序列表示上的平方複雜度而變得 prohibitively 昂貴。此外,CNN和 Transformer 的前期研究都表明,淺層主要關注局部模式,而更深層則傾向於捕捉高級語義或全局關係[14; 37; 5],因此,在早期階段通常不需要通過MHSA建模全局關係。基於這些事實,劉等人[22]提出了一種分層的 Transformer ,並採用移位窗口來限制自注意力計算限於非重疊的局部窗口,同時允許跨窗口連接,從而帶來更高的效率。在另一項工作線中,潘等人[26]和郭等人[9]將卷積層引入到 Transformer 的前幾層。在這些工作的基礎上,瓦蘇等人[32]提出了RepMixer作為 Token 混合器,它使用結構重參化通過在網絡中刪除 Short-Cut 來降低內存訪問成本。
作為一種簡單但有效輕量級方法,權重共享策略已在許多基於 Transformer 的模型中得到應用。Dehghani等人[3]首次提出了利用 Transformer 層進行自然語言處理任務的不同動機:他們將重覆的網絡層視為 Transformer 引入循環歸納偏置的一種互補方式,並觀察到他們的方法在多個任務上優於原始 Transformer 。Jaegle等人[17]在多模態預訓練中採用了交叉注意力層權重共享。Hernandez等人[12]探討了在不同粒度 Level 上共享符合器[8]的不同部分,受模型大小硬約束。最近,He等人[10]通過信號傳播理論研究了標準Pre-LN Transformer 塊[34],並提出了一個配備形狀注意力[25]的簡化並行結構 Transformer 塊作為 Token 混合器,稱為簡化注意力子塊並行(SAS-P),在降低模型參數數量的同時,不損失語言下遊任務的表現。作者的工作是首次嘗試將SAS-P擴展到多模態領域,並通過在相鄰塊之間共享 Token 混合器的權重進一步簡化。
3 Methods
在本節中,作者首先介紹了作者的簡化模型結構,該結構通過共享SAS-P塊的權重來實現。然後,作者介紹了一種名為WIKD的高效訓練方法。接下來,作者介紹了一種新的損失函數,稱為配對匹配(PM)損失,以進一步提高訓練性能。最後,作者還通過在訓練模型所使用的CC12M數據集中添加合成字幕,以提高數據多樣性和數據質量,同時最小化附加存儲空間。新的數據集稱為CC12M-SYN。圖2顯示了作者的方法的整體框架。

Simplifying Model Structure by Sharing Weights Among SAS-P Blocks
作者基於最新的MobileCLIP-S0模型[33]構建了作者的架構,並通過多種方式對其進行了增強。MobileCLIP-S0框架具有圖像編碼器(image encoder)和文本編碼器(text encoder)的混合結構,其中包含基於卷積(convolution-based)和MHSA(Mobile Hyperspectral-to-Spatial Augmentation)的塊的協同排列。然而,對於每個基於MHSA的塊,MobileCLIP-S0僅採用標準的Pre-LN塊,其中MHSA作為混合器(token mixer)[34],如圖3(左)所示。

作者首先從每個預連接塊內的 Shortcut 參數開始減少。這些連接在內存訪問和推理速度上造成了 Bottleneck ,因此,使用輕量級的基於MHSA的塊設計變得至關重要。此外,已經證明,注意力模塊的饋送層可以無縫集成到 Transformer 層中,而不會降低 Transformer 層的性能[31; 24]。
如圖3右側說明了SAS-P(He和Hofmann,2024年),這是一種簡化的並行 Transformer 模塊,在消除跳接連接的同時,還消除了值和投影參數。它使用形狀注意力[25]作為其 Token 混合器,以防止在消除跳接連接後信號降解,使注意力矩陣更具標識性,從而保持良好的信號傳播。形狀注意力的注意力矩陣由:

其中 表示 SAS-P 輸入, 和 分別是 Query 和鍵矩陣, 是模型維數, 是單位矩陣, 是可學習的參數。 是中心矩陣,其中每個元素都設置為 是輸入 Token的數量。在初始化時, 設置為 分別設置為 1 ,導致 和 ,這對於良好的信號傳播是有效的。
SAS-P 在多個語言任務上表現出令人印象深刻的性能,同時比 Pre-LN 更快地推理速度,並使用較少的參數。為了進一步簡化模型結構,作者評估了相鄰基於MHSA的塊之間的Jensen-Shannon(JS)離差(見圖4)。低JS離差意味著在這些矩陣之間進行權重共享不會降低性能。因此,在訓練期間使用KD之前,作者的”學生”模型用SAS-P塊替換了所有的預LN塊,並在這些塊之間進行權重共享。結果,作者模型的圖像編碼器與MobileCLIP-S0相比大約有14%更少的參數,而在OpenAI-B/16[28]中只有11%。
為了從小型數據集中受益,一個廣泛使用的方法是使用與任務相關的預訓練 Backbone 網絡,並添加一些任務特定的層[15]。受到使用 Backbone 網絡思想的啟發,作者採用WI[36]來訓練CLIP在小型數據集上。在實際應用中,由於作者修改了MobileCLIP-S0結構的基於MHSA的模塊,而保持基於RepMixer的模塊不變(這些模塊已經高效),作者可以直接從已經在大型數據集上預訓練得很好的MobileCLIP-S0繼承這些模塊的權重。在這種情況下,繼承的模塊可以被視為「 Backbone 」。
然後,作者凍結這些繼承的層,只在新添加的SAS-P塊上在一個非常小的數據集上進行訓練。應用上述方法可以減少梯度存儲,從而允許作者使用更大的批量大小來保持對比學習的性能。此外,作者認為作者的模型是MobileCLIP-S0的移動模型,並在訓練過程中進行多階段知識蒸餾,從而進一步提高性能。具體而言,作者在單模態特徵空間(階段1)、對比關係空間(階段2)和交互對比關係空間(階段3)上應用知識蒸餾。對於一個圖像文本對齊的批處理,學生模型首先通過優化特徵蒸餾損失()模擬教師圖像和文本特徵分佈:

其中 和 分別表示教師模型和學生模型的(圖像,文本)特徵,其中 為批量大小。然後,它計算對比關係蒸餾損失 和交互對比關係損失 來模擬對比關係空間和交互對比關係空間的圖像-文本相似矩陣分佈,定義如下:

作者的最終蒸餾損失定義為:
其中是一個可學習的溫度參數,Sim 表示圖像特徵和文本特徵之間的相似矩陣。

λ1, λ2, λ3 均為超參數。
Pair Matching (PM) Loss
小型數據集訓練的CLIP模型通常在零樣本性能方面表現不佳[38]。作者認為這種現象的一個可能原因是,使用較少數據訓練的模型在區分語義上相似的圖像-文本對方面存在更大的困難。因此,作者提出構建一個輔助超平面來幫助模型判斷圖像-文本對是否匹配。請參閱圖5。

特別地,作者增加了一個額外的二進製匹配任務。給定一組圖像-文本對,作者首先提取正面的圖像-文本對,並計算它們的匹配對數 和 ,計算公式如下:

關於負樣本對,對於每個圖像,作者根據圖像到文本相似性矩陣(如果其嵌入相似,則選擇一個負樣本文本的概率較高)選擇一個負文本。同樣地,這個過程也應用於每個文本。因此,負匹配logits被定義為:

因此,作者提出的PM損失函數可以表示為:

CC12M-SYN Dataset
用於訓練CLIP模型的圖像文本數據集大多來源於互聯網,其中包含固有的噪聲樣本,這些樣本的描述性不足。當使用小型數據集時,數據樣本的多樣性和質量變得尤為重要。增加合成字幕是一種既便宜又有效的方法,可以提高多樣性和質量。作者採用廣泛使用的數據集CC12M [1],並使用coca [43]為該數據集中的每個圖像生成多個合成字幕,得到CC12M-SYN。圖6展示了CC12M-SYN中的一些帶有合成字幕的示例。在訓練過程中,作者隨機選擇一組原始和合成字幕中的一個文本。因此,CC12M-SYN中的一個樣本包括一個圖像和一個合成或原始的描述。

4 Experiments
Implementation Details
作者在第一萬次訓練迭代中採用了一個 Warm up 策略。作者使用了AdamW優化器,並將批量大小設置為1536,權重衰減設置為0.1。作者在Nvidia RTX3090上訓練模型32個週期,學習率設置為0.001。在消融研究中,作者將週期數設置為9。作者採用MobileCLIP-S0作為WIKD的教師。對於超參數,作者設置λ1=4000,λ2=λ3=1,λ4=0.1。其他設置遵循CLIP-KD [38]。
作者在多個數據集上評估了零樣本性能。具體來說,作者使用ImageNet-1k [4],ImageNet-V2 [29],ImageNet-R [11]和ImageNet-S [35]來評估零樣本圖像分類性能。對於零樣本圖像文本檢索,作者使用MSCOCO [21]和Flickr30k [27]。預設為圖像分類中的Top-1準確率(acc1)和圖像文本檢索中的R@1。
4.1.1 Data Augmentation.
作者應用了隨機縮放(RandomResizedCrop)和隨機裁剪(RandAugment)進行圖像增強。在RandomResizedCrop中,作者將縮放比例設置為(0.08, 1.0),以在原始圖像上進行強大的增強。然後,作者在處理後的圖像上應用RandAugment,通過隨機採用預設的31種增強方法([2])進一步增加圖像的多樣性。
Main Results
4.2.1 Zero-shot image-text retrieval.
表1報告了在MSCOCO和Flickr30k上的零樣本圖像文本檢索性能。與相似規模的訓練數據集(最多2000萬樣本)上的模型相比,作者的模型在MSCOCO上超過了所有其他工作。對於Flickr30k,作者的模型在使用了較少的參數的同時,也實現了當前最先進模型TinyCLIP的性能。與在大規模數據集(400M-1B)上訓練的模型相比,作者的模型在競爭性能上取得了良好的表現,並超越了許多現有工作。例如,與最先進的MobileCLIP-S0相比,作者的模型在使用了大約3%的訓練樣本和14%更少的圖像編碼器參數的情況下,T2I性能僅低約1%。此外,作者的模型在兩個數據集的I2T和T2I指標上都超過了DataComp-B/32、OpenAI-X(除B/16)和LAION-B/32。

4.2.2 Zero-shot image classification on ImageNet.
表2報告了零樣本分類性能。與其他在相似規模的數據集上進行訓練的模型相比,作者的模型在所有報告的數據集上都優於其他工作,這證明了作者的方法的有效性。對於大規模數據集,儘管與最新的state-of-the-art DataComp-B/16相比並非最佳,但作者與幾項現有工作相比仍取得了一些具有競爭力的結果。

4.2.3 Inference Speed.
為了評估推理速度,作者在CPU(Intel(R)-Xeon(R)-Silver-4314-CPU@2.40GHz)上進行了模擬實驗,並比較了作者的模型與最先進的MobileCLIP系列[33]的平均推理速度。如表3所示,當輸入序列包含1000張圖像時,作者模型的處理速度達到39.5張/秒,略高於最先進的MobileCLIP-S0(38.2張/秒)。這表明採用SAS-P塊的好處。
Ablation Studies
4.3.1 Training Efficiency of CC12M-SYN.
為了證明CC12M-SYN的訓練效率提升,作者在CC12M-SYN和CC12M上分別訓練作者的模型20個週期。作者報告了前9個週期的平均損失曲線,以及最後週期的IN-1k和Flickr30k上的零樣本性能。圖7報告了CC12M和CC12M-SYN的損失曲線,顯示在CC12M-SYN上訓練導致損失更快減少。表4表明,在CC12M-SYN上訓練的模型在零樣本分類和零樣本圖像文本檢索方面具有更好的性能。這些結果表明合成標籤對於數據多樣性和質量提升具有益處。


4.3.2 Analysis of WIKD and PM loss.
作者通過比較有無WIKD和PM損失的訓練( Baseline )以及僅WI、僅WIKD和同時WIKD和PM損失的訓練效果,來探索WIKD和PM損失的有效性。結果見表5。它表明,僅WI的訓練可以提高零樣本分類(+13.0和+6.1/+0.1在分類準確率acc1上)和圖像文本檢索(R@1的檢索率+15.9/+15.2)。而當訓練僅WIKD時,性能更高(+25.4和+15.9/+15.2,分別)。當同時訓練WIKD和PM損失時,模型達到最高性能。這些結果明顯證明了WIKD和PM損失的有效性。

5 Conclusion
在這項工作中,作者提出了一系列技術,使得在消費級計算機上進行CLIP模型的訓練和推理成為可能,同時實現了具有競爭力的性能。這對於將基礎模型的卓越成果帶到邊緣設備上至關重要。
作者減少了模型結構,提高了推理速度。
此外,作者提出了WIKD和PM損失函數,這些損失函數對性能提升做出了貢獻,並且可以用於簡化其他領域中的模型。
最後,在增強的CC12M-SYN數據集上訓練後,作者的模型在參數更少、訓練數據集更小的情況下,與現有工作相比仍實現了具有競爭力的性能。
參考文獻
[0]. Simplifying CLIP: Unleashing the Power of Large-Scale Models on Consumer-level Computers.