訓(xùn)練Backbone你還用EMA？ViT訓(xùn)練的大殺器EWA升級(jí)來襲
點(diǎn)擊下方卡片，關(guān)注「集智書童」公眾號(hào)
結(jié)構(gòu)重參化是一種適用于卷積神經(jīng)網(wǎng)絡(luò)（CNNs）的通用訓(xùn)練方案，它在不增加推理成本的情況下實(shí)現(xiàn)了性能提升。隨著視覺Transformer（ViTs）在各種視覺任務(wù)中逐漸超越CNNs，人們可能會(huì)提出一個(gè)問題：是否存在一種專門針對(duì)ViTs的訓(xùn)練方案，也能在不增加推理成本的情況下實(shí)現(xiàn)性能提升？
最近，由于多個(gè)分支結(jié)構(gòu)可以有效地?cái)U(kuò)展Transformer的容量，Mixture-of-Experts（MoE）引起了越來越多的關(guān)注?？紤]到MoE也可以被看作是一個(gè)多分支結(jié)構(gòu)，我們是否可以利用MoE來實(shí)現(xiàn)類似于結(jié)構(gòu)重參化的ViT訓(xùn)練方案？
在本文中，作者肯定地回答了這些問題，并提出了一種新的ViT通用訓(xùn)練策略。具體來說，作者將ViT的一些前饋網(wǎng)絡(luò)（FFNs）替換為經(jīng)過特殊設(shè)計(jì)的、更高效的MoE，這些MoE通過隨機(jī)均勻分區(qū)將Token分配給 Expert，并在每次迭代結(jié)束時(shí)對(duì)這些MoE進(jìn)行 Experts Weights Averaging（EWA）。訓(xùn)練后，通過對(duì) Expert進(jìn)行平均，將每個(gè)MoE轉(zhuǎn)化為一個(gè)FFN，將模型轉(zhuǎn)換回原始的ViT以進(jìn)行推理。
作者還提供了理論分析，以展示為什么以及如何運(yùn)作。對(duì)各種2D和3D視覺任務(wù)、ViT體系結(jié)構(gòu)和數(shù)據(jù)集進(jìn)行的綜合實(shí)驗(yàn)驗(yàn)證了所提出的訓(xùn)練方案的有效性和通用性。此外，當(dāng)對(duì)ViTs進(jìn)行微調(diào)時(shí)，作者的訓(xùn)練方案也可以應(yīng)用于提高性能。
最后但同樣重要的是，所提出的EWA技術(shù)可以在各種2D視覺小數(shù)據(jù)集和3D視覺任務(wù)中顯著提高樸素MoE的效果。
1、簡介
近年來，通用訓(xùn)練方案在深度學(xué)習(xí)中發(fā)揮了關(guān)鍵作用。這些方法在不修改模型架構(gòu)的情況下增強(qiáng)了現(xiàn)有模型，并且不需要對(duì)特定架構(gòu)做出任何假設(shè)。因此，它們可以應(yīng)用于各種模型，并提供均勻的好處。
對(duì)于CNN架構(gòu)，廣泛使用的通用訓(xùn)練方案是結(jié)構(gòu)重參化（structural re-parameterization）（圖1(b)）。這種方法通過在訓(xùn)練過程中使用多分支結(jié)構(gòu)，將CNN模型的訓(xùn)練和推理階段解耦，并在推理過程中等效地轉(zhuǎn)換回原始的單分支結(jié)構(gòu)。結(jié)構(gòu)重參化可以在不增加推理成本的情況下實(shí)現(xiàn)性能提升。
然而，在訓(xùn)練過程中使用的多分支結(jié)構(gòu)可能會(huì)降低訓(xùn)練速度（如Tab. 1的第2列所示，其訓(xùn)練延遲高達(dá)普通訓(xùn)練的2.18倍）。此外，結(jié)構(gòu)重參化僅適用于CNN，而不適用于其他像ViTs（如Tab. 1的第6列所示）這樣的最近更強(qiáng)大的架構(gòu)。
對(duì)于Transformer架構(gòu)，最近的通用訓(xùn)練方法涉及將前饋網(wǎng)絡(luò)（FFN）層替換為稀疏的 Expert混合（MoE）（圖1(a)）。MoE是一種新穎的網(wǎng)絡(luò)組件，包括多個(gè)具有唯一權(quán)重和可訓(xùn)練Routing網(wǎng)絡(luò)的 Expert（FFN）。
在訓(xùn)練和推理階段，Routing網(wǎng)絡(luò)為每個(gè)輸入選擇一種稀疏的 Expert組合，通過稀疏計(jì)算實(shí)現(xiàn)了Transformer模型容量的有效擴(kuò)展。在自然語言處理領(lǐng)域，許多稀疏的MoE模型已經(jīng)證明它們的性能優(yōu)于密集模型。
然而，MoE仍然具有一些限制，限制了它在Vision Transformers（ViTs）中的廣泛使用。
首先，V-MoE和SwinV2-MoE是在大規(guī)模數(shù)據(jù)集（JFT-300M，ImageNet-22K）上進(jìn)行預(yù)訓(xùn)練的，而當(dāng)在ImageNet-1K上進(jìn)行預(yù)訓(xùn)練時(shí)，V-MoE的性能低于ViT。此外，還沒有研究研究MoE Vision Transformer在小數(shù)據(jù)集（如CIFAR-100和Tiny-ImageNet）上的情況。
其次，當(dāng)前的MoE ViTs專注于2D視覺任務(wù)，沒有用于3D視覺任務(wù)的MoE ViTs（見Tab. 1的第4行）。
第三，MoE本身的前kRouting機(jī)制并不高效和有效。如Tab. 1的第4行所示，由于前kRouting需要額外的計(jì)算來為每個(gè)輸入做出Routing決策，訓(xùn)練和推理的延遲都增加到了1.12倍，而訓(xùn)練和推理的參數(shù)都大幅增加到了3.25倍，給模型帶來了負(fù)擔(dān)。
此外，前kRouting容易導(dǎo)致 Expert之間的負(fù)載不平衡，因此許多研究致力于通過各種設(shè)計(jì)來解決這個(gè)問題，如添加輔助負(fù)載平衡損失，在全局范圍內(nèi)做出Routing決策，或重新設(shè)計(jì)Routing算法。
因此，以下問題產(chǎn)生了：
是否有一種新的訓(xùn)練方法，可以在訓(xùn)練階段高效地訓(xùn)練稀疏激活的 Expert，而在推理階段可以將稀疏模型轉(zhuǎn)化為稠密模型而不需要額外的計(jì)算成本？是否可以進(jìn)一步將MoE的有效性擴(kuò)展到2D視覺數(shù)據(jù)集和3D視覺任務(wù)中？
在本文中設(shè)計(jì)了一種名為"Experts Weights Averaging"（EWA）訓(xùn)練的新型通用訓(xùn)練策略，用于ViTs。本文的方法實(shí)現(xiàn)了性能提升，同時(shí)不增加推理延遲和參數(shù)，如Tab. 1的第2行和最后一行所示。
具體而言，作者在ViT的訓(xùn)練和推理階段進(jìn)行了解耦，如圖1(c)所示。在訓(xùn)練過程中，用經(jīng)過特殊設(shè)計(jì)的、更高效的MoEs替換了普通ViT中的一些FFN，這些MoEs通過隨機(jī)均勻劃分（RUP）將Token分配給 Expert，不需要額外的參數(shù)、特殊設(shè)計(jì)或輔助損失。
此外，每次迭代結(jié)束時(shí)，對(duì)每個(gè)MoE執(zhí)行"Experts Weights Averaging"（EWA）。在推理過程中，通過簡單地將每個(gè)MoE的 Expert平均為單個(gè)FFN，可以將稀疏的MoE模型轉(zhuǎn)換為普通的ViT模型，而不會(huì)有任何性能損失。
作者對(duì)各種2D和3D視覺任務(wù)、ViT架構(gòu)和數(shù)據(jù)集進(jìn)行了全面的實(shí)驗(yàn)，以驗(yàn)證所提出的訓(xùn)練方案的有效性和泛化性。例如，作者提出的訓(xùn)練方法在圖像分類任務(wù)上實(shí)現(xiàn)了1.72%的性能提升，在點(diǎn)云語義分割任務(wù)上實(shí)現(xiàn)了1.74%的平均IoU提升。
此外，EWA訓(xùn)練還可以用于微調(diào)預(yù)訓(xùn)練的ViT模型。本文的訓(xùn)練方法將CIFAR100上的ViT-B微調(diào)結(jié)果從90.71%提高到91.42%（↑ 0.71%）。最后，作者發(fā)現(xiàn)EWA訓(xùn)練中的"Experts Weights Averaging"關(guān)鍵技術(shù)可以顯著提高2D小型視覺數(shù)據(jù)集和3D視覺任務(wù)中原始MoE的效果。
主要貢獻(xiàn)總結(jié)如下：
-
提出了一種新的通用訓(xùn)練方案，用于ViTs，實(shí)現(xiàn)了性能提升，同時(shí)不增加推理延遲和參數(shù)。在訓(xùn)練過程中，將ViT架構(gòu)中的某些FFN替換為使用隨機(jī)均勻劃分的MoE，并在每次模型權(quán)重更新后對(duì)這些MoE層執(zhí)行"Experts Weights Averaging"。訓(xùn)練結(jié)束后，通過將模型中的每個(gè)MoE層的 Expert進(jìn)行平均，將模型轉(zhuǎn)換回原始的ViT架構(gòu)，用于推理。
-
驗(yàn)證了EWA訓(xùn)練可以為不同的2D/3D視覺任務(wù)和數(shù)據(jù)集上的各種Vision Transformer架構(gòu)帶來統(tǒng)一的性能提升。此外，作者發(fā)現(xiàn)EWA訓(xùn)練還可以用于微調(diào)預(yù)訓(xùn)練的ViT模型，實(shí)現(xiàn)更進(jìn)一步的性能提升。作者提供了理論分析，展示了為什么以及如何有效。
-
探索了在小型2D視覺數(shù)據(jù)集和3D視覺任務(wù)中原始MoE的有效性。作者發(fā)現(xiàn)使用原始MoE的性能低于密集模型，但所提出的"Experts Weights Averaging"技術(shù)可以顯著提高原始MoE的效果。
2、相關(guān)工作
2.1、MoE
稀疏門控混合 Expert模型（Mixture-of-Experts，MoE）在訓(xùn)練時(shí)間、模型容量和性能方面顯示出顯著的改進(jìn)。此外，當(dāng)MoE被引入基于 Transformer的語言模型中時(shí)，它受到越來越多的關(guān)注。迄今為止，許多工作都改進(jìn)了MoE的關(guān)鍵組件，即Routing網(wǎng)絡(luò)。
一方面，對(duì)于傳統(tǒng)的可學(xué)習(xí)Routing網(wǎng)絡(luò)，Switch Transformer簡化了Routing，僅為每個(gè)Token選擇前1個(gè) Expert，Base Layer通過處理線性分配問題來執(zhí)行Routing，Expert Choice將前k個(gè)TokenRouting到每個(gè) Expert。
另一方面，Hash Layer為TokenRouting設(shè)計(jì)了確定性哈希，THOR隨機(jī)激活每個(gè)輸入的 Expert。除了Routing網(wǎng)絡(luò)的改進(jìn)，許多工作還探索了MoE在現(xiàn)代硬件、擴(kuò)展性和應(yīng)用方面的實(shí)現(xiàn)。
然而，大多數(shù)MoE工作都是基于 Transformer的語言模型，計(jì)算機(jī)視覺領(lǐng)域中只有少數(shù)MoE Vision Transformer的工作存在。
Riquelme等人通過將ViT的一些FFN層替換為MoE層，對(duì)JFT-300M或JFT-3B進(jìn)行預(yù)訓(xùn)練，取得了在圖像分類方面的巨大成功。Hwang等人在Swin-transformer V2架構(gòu)中添加MoE層，以進(jìn)行ImageNet-22K的預(yù)訓(xùn)練，用于圖像分類和目標(biāo)檢測(cè)任務(wù)。Xue等人通過將FFN替換為MoE并在 Transformer塊之間共享MoE層，提出了一種更加參數(shù)高效和有效的框架WideNet，在ImageNet-1K預(yù)訓(xùn)練中實(shí)現(xiàn)了性能提升。Komatsuzaki等人使用預(yù)訓(xùn)練的密集 Transformer檢查點(diǎn)來啟動(dòng)MoE模型的訓(xùn)練，并在JFT-300M預(yù)訓(xùn)練和ImageNet-1K微調(diào)方面取得了有希望的結(jié)果。Xue等人通過從預(yù)訓(xùn)練的MoE模型中進(jìn)行知識(shí)蒸餾，來學(xué)習(xí)一個(gè)密集模型，并在ImageNet-1K上進(jìn)行了驗(yàn)證。
在本文中提出了一種新的通用訓(xùn)練方法，以提高ViT的性能，而不增加推理成本，其目標(biāo)和實(shí)現(xiàn)方法與MoE不同。此外，以前的MoE-ViT研究主要集中在更大的2D視覺數(shù)據(jù)集上，在較小的2D視覺數(shù)據(jù)集和3D視覺方面進(jìn)行了較少的探索。
本文表明，在ViT中添加原始的MoE會(huì)降低在小型2D視覺數(shù)據(jù)集和3D視覺上的性能，并引入了一種新穎的"Early EMA"訓(xùn)練方法，可以極大地提高它們的性能。
2.2、結(jié)構(gòu)重參化
結(jié)構(gòu)重參數(shù)化的本質(zhì)是在訓(xùn)練過程中使用多分支結(jié)構(gòu)代替卷積或全連接層，以增強(qiáng)模型。訓(xùn)練完成后，使用等效的參數(shù)轉(zhuǎn)換或特定算法融合多個(gè)分支以進(jìn)行推理。典型的例子有ACNet、DBB和RepVGG。
具體而言，ACNet在訓(xùn)練過程中將常規(guī)的K × K卷積替換為具有K × 1、1 × K和K × K卷積并行結(jié)構(gòu)的ACBlocks。訓(xùn)練完成后，ACBlocks等效地轉(zhuǎn)換為單個(gè)K × K卷積層。DBB在訓(xùn)練過程中使用1 × 1卷積、K × K卷積、1 × 1-K × K卷積和1 × 1卷積-平均池化的多分支結(jié)構(gòu)來替代原始的K × K卷積，推理過程中等效地轉(zhuǎn)換為單個(gè)K × K卷積。RepVGG在訓(xùn)練過程中在每個(gè)3 × 3卷積層上添加并行的1 × 1卷積分支和恒等分支，在訓(xùn)練后等效地轉(zhuǎn)換為單個(gè)3 × 3卷積層。
由于能夠在不增加推理成本的情況下提高性能，結(jié)構(gòu)重參化在各種計(jì)算機(jī)視覺任務(wù)中受到了越來越多的關(guān)注和應(yīng)用。例如，ExpandNets使用它來設(shè)計(jì)緊湊型模型，RepNAS將其用于神經(jīng)架構(gòu)搜索，ResRep將其與剪枝相結(jié)合。然而，以前對(duì)于結(jié)構(gòu)重參化的研究主要集中在CNN架構(gòu)上，這使得它們與ViT架構(gòu)不兼容。
本文引入了一種新穎的方法，將ViT架構(gòu)的訓(xùn)練和推理階段解耦。在訓(xùn)練過程中，使用高效的MoE稀疏計(jì)算，采用隨機(jī)均勻劃分和 Experts Weights Averaging機(jī)制。在推理過程中，將模型轉(zhuǎn)換回原始的ViT架構(gòu)，從而在不增加推理參數(shù)和延遲的情況下實(shí)現(xiàn)性能改進(jìn)。這種方法是首次將MoE稀疏計(jì)算應(yīng)用于改善ViT架構(gòu)性能的方法。
2.3、Weight Averaging
權(quán)重平均是深度學(xué)習(xí)中常用的技術(shù)，在各種任務(wù)中得到了廣泛應(yīng)用。例如，在自監(jiān)督和半監(jiān)督學(xué)習(xí)中，有方法使用學(xué)生模型的指數(shù)移動(dòng)平均（EMA）權(quán)重作為教師，為學(xué)生提供更平滑的目標(biāo)。同樣，在在線知識(shí)蒸餾中，PCL將每個(gè)分支的EMA權(quán)重作為對(duì)等的平均教師，協(xié)同地在分支之間傳遞知識(shí)。
此外，許多研究使用權(quán)重平均來增強(qiáng)模型的泛化能力。例如，隨機(jī)權(quán)重平均（SWA）對(duì)SGD優(yōu)化軌跡上的多個(gè)點(diǎn)的權(quán)重進(jìn)行平均，實(shí)現(xiàn)了比傳統(tǒng)訓(xùn)練更好的泛化能力。Wortsman等人對(duì)原始的零樣本模型和微調(diào)模型進(jìn)行平均，從而在分布內(nèi)和分布外的泛化中都取得了改進(jìn)。Model Soup在微調(diào)期間對(duì)具有相同預(yù)訓(xùn)練初始化但不同超參數(shù)配置的模型進(jìn)行平均，以獲得更好的模型。DiWA對(duì)權(quán)重平均在OOD方面的成功提供了理論分析，并提出通過對(duì)多個(gè)獨(dú)立訓(xùn)練的模型的權(quán)重進(jìn)行平均來減少預(yù)測(cè)之間的協(xié)方差，從而提高OOD的泛化性能。
與上述在模型級(jí)別上進(jìn)行權(quán)重平均以獲得單個(gè)模型的研究不同，本文的方法在MoE的訓(xùn)練過程中在 Expert級(jí)別上執(zhí)行權(quán)重平均。通過將其他 Expert的權(quán)重平均到每個(gè) Expert上，獲得了多個(gè) Expert。
2.4、Dropout and Other variants for ViT
Dropout是一種常見的技術(shù)，通過在訓(xùn)練過程中隨機(jī)地丟棄神經(jīng)元及其對(duì)應(yīng)的連接來防止過擬合，從而提高網(wǎng)絡(luò)的泛化能力。對(duì)于ViTs，大多數(shù)方法在自注意力或FFN層中使用固定的丟棄率進(jìn)行dropout。與此同時(shí)，隨機(jī)深度（Stochastic Depth）通常用于ViTs，隨機(jī)地丟棄一些塊，并對(duì)更深的塊應(yīng)用較高的丟棄率以提高泛化性能。最近，一種針對(duì)ViT的新型dropout變體被引入，稱為dropkey。Dropkey通過在注意力計(jì)算過程中隨機(jī)丟棄一些Key為每個(gè)注意力塊分配自適應(yīng)的操作符。這種策略約束了注意力分布，使其更加平滑，更有效地減輕了過擬合問題。此外，dropkey會(huì)隨著塊數(shù)的增加逐漸降低丟棄率，避免了在訓(xùn)練過程中使用恒定丟棄率的傳統(tǒng)dropout方法可能導(dǎo)致的不穩(wěn)定性。
在本文的方法中，對(duì)于每個(gè)MoE層，在訓(xùn)練過程中隨機(jī)且均勻地將所有Token分配給所有 Expert。這等效于在每個(gè) Expert處隨機(jī)丟棄大量的Token。這樣的操作可能有助于防止過擬合，同時(shí)確保訓(xùn)練延遲幾乎與原始訓(xùn)練相當(dāng)。
3、本文方法
3.1、PRELIMINARY
1、MoE
如圖1(a)所示，標(biāo)準(zhǔn)的MoE層包括一組N個(gè) Expert（即N個(gè)FFN），表示為 ，以及具有權(quán)重 的Routing網(wǎng)絡(luò)G。給定輸入示例 ，訓(xùn)練時(shí)和推斷時(shí)的MoE層的輸出可以寫成：
其中， 表示到第 個(gè) Expert的Routing分?jǐn)?shù)， 代表第 個(gè) Expert的輸出， 表示選擇前 個(gè) Expert，并將其他 Expert的 設(shè)置為0。通常情況下， ，這意味著 是一個(gè)稀疏的N維向量。當(dāng) 時(shí)，不需要計(jì)算 。
2、結(jié)構(gòu)重參化
如圖1(b)所示，結(jié)構(gòu)重參化在訓(xùn)練期間使用多分支CNN結(jié)構(gòu)替代常規(guī)卷積層。用 表示總共的M個(gè)分支（通常是具有兼容尺寸的M個(gè)不同運(yùn)算符）。對(duì)于給定的輸入 ，訓(xùn)練時(shí)結(jié)構(gòu)的輸出可以表示為：
訓(xùn)練完成后，這些M個(gè)分支將被等效地轉(zhuǎn)換為單個(gè)卷積層 進(jìn)行推斷。對(duì)于推斷時(shí)的輸入 ，輸出變?yōu)?span style="cursor:pointer;"> 。
3.2、 A New General Training Scheme
1、動(dòng)機(jī)
本論文旨在設(shè)計(jì)一種新的通用訓(xùn)練方案，用于ViT，在不增加推斷成本的情況下實(shí)現(xiàn)性能提升。如公式3所示，經(jīng)典的結(jié)構(gòu)重參化僅適用于由卷積組成的CNN，由于多分支結(jié)構(gòu)，它顯著增加了訓(xùn)練成本。如公式2所示，最近的MoE由N個(gè) Expert和一個(gè)Routing網(wǎng)絡(luò)組成。
-
首先，Routing網(wǎng)絡(luò)引入了需要學(xué)習(xí)的附加參數(shù)，并容易導(dǎo)致負(fù)載不均衡問題。
-
其次，N個(gè) Expert大大增加了部署模型的負(fù)擔(dān)。
-
第三，Routing分?jǐn)?shù) 在將MoE轉(zhuǎn)換為FFN時(shí)成為瓶頸。
考慮到這些問題，首先提出了一種特別設(shè)計(jì)的、更高效的MoE，用于替換ViT的一些FFN，通過隨機(jī)均勻劃分將Token分配給 Expert。然后，提出了一種簡單但有效的技術(shù)，稱為 Experts Weights Averaging，它在每個(gè)訓(xùn)練迭代結(jié)束時(shí)將其他 Expert的權(quán)重平均到每個(gè) Expert。因此，在訓(xùn)練完成后，可以將每個(gè)MoE轉(zhuǎn)換為單個(gè)FFN進(jìn)行推斷。
2、EWA Training
Overall
在將EWA訓(xùn)練應(yīng)用于從頭開始訓(xùn)練ViT模型時(shí)，首先通過用基于隨機(jī)均勻劃分的MoE層替換一些FFN來創(chuàng)建一個(gè)隨機(jī)初始化的MoE ViT。在訓(xùn)練過程中，在每次權(quán)重更新后對(duì)每個(gè)MoE層進(jìn)行 Experts Weights Averaging，共享率為 ，并且 會(huì)在訓(xùn)練時(shí)期從0線性增加到超參數(shù)共享率。在訓(xùn)練完成后，通過對(duì) Expert進(jìn)行平均將每個(gè)MoE轉(zhuǎn)換為一個(gè)FFN。
MoE layer with Random Uniform Partition
考慮一個(gè)由 個(gè) Expert（即N個(gè)相同的FFN）組成的MoE層，表示為 ，假設(shè)有T個(gè)輸入 ，其中 可以被 整除。具有隨機(jī)均勻劃分（RUP）的MoE層的輸出可以表示為：
如第4式所示， 個(gè)輸入被隨機(jī)均勻地劃分為 部分，其中 表示分配給第 個(gè) Expert的 個(gè)輸入。進(jìn)一步地，如第5式所示，對(duì)于 中的每個(gè)輸入，其輸出為 。
Experts Weights Averaging
給定 個(gè) Expert的權(quán)重 和共享率 ，對(duì)每個(gè)MoE層執(zhí)行的Experts Weights Averaging（EWA）可以表述為：
其中 表示第 個(gè) Expert的新權(quán)重。簡而言之，將其他 Expert的權(quán)重平均分配給每個(gè) Expert，從而獲得多個(gè) Expert。
Convert MoE into FFN
訓(xùn)練結(jié)束后，每個(gè)MoE層將通過簡單地平均 Expert的權(quán)重來轉(zhuǎn)換為一個(gè)FFN層。對(duì)于由N個(gè) Expert 組成的MoE層，相應(yīng)的推斷時(shí)FFN可以表示為：
通過這種方式，訓(xùn)練時(shí)的MoE模型將在推斷時(shí)轉(zhuǎn)換為原始的ViT模型。
3、EWA Fine-tuning
在將EWA訓(xùn)練應(yīng)用于微調(diào)預(yù)訓(xùn)練的ViT模型時(shí)，與從頭開始訓(xùn)練ViT的唯一區(qū)別是創(chuàng)建的MoE ViT模型的權(quán)重初始化。不同于隨機(jī)初始化，使用給定的ViT模型檢查點(diǎn)來初始化MoE ViT。在用MoE層替換某些FFN層時(shí)，特定MoE層的每個(gè) Expert被初始化為預(yù)訓(xùn)練ViT模型中對(duì)應(yīng)FFN層的副本。對(duì)于保持不變的其他層，它們的權(quán)重直接繼承自原始的ViT檢查點(diǎn)。
4、理論分析
本節(jié)從理論上分析了EWA訓(xùn)練為何有效以及其工作原理。為了方便分析，關(guān)注一個(gè)MoE層和m個(gè)訓(xùn)練步驟。假設(shè)當(dāng)前訓(xùn)練步驟為 ，在根據(jù)公式6計(jì)算后，經(jīng)過 Experts Weights Averaging化后的第 個(gè) Expert的新權(quán)重可以表示為：
接下來，開始第 個(gè)訓(xùn)練步驟，MoE的每個(gè) Expert的權(quán)重更新為，其中 ， 表示學(xué)習(xí)率。進(jìn)一步進(jìn)行 Experts Weights Averaging化，如公式9所示。根據(jù)公式8和 ，可以將公式9進(jìn)一步重寫為公式10。
類似地，可以得到 ， 和 。通過數(shù)學(xué)歸納，得到公式12。
根據(jù)上述推導(dǎo)，有兩個(gè)發(fā)現(xiàn)：
-
觀察公式9和公式11的第一個(gè)項(xiàng)，進(jìn)行 Experts Weights Averaging化時(shí)會(huì)出現(xiàn)逐層的權(quán)重衰減；
-
觀察公式10和公式12的最后一個(gè)項(xiàng)，會(huì)在EWA訓(xùn)練的訓(xùn)練迭代過程中持續(xù)聚合多 Expert歷史指數(shù)平均權(quán)重。
5、實(shí)驗(yàn)
5.1、EWA Training
1、在各種架構(gòu)和數(shù)據(jù)集上的2D分類
如表2所示，與普通訓(xùn)練相比，EWA訓(xùn)練可以為不同的ViT架構(gòu)在不同的數(shù)據(jù)集上帶來一致的性能提升。對(duì)于標(biāo)準(zhǔn)的ViT-S，EWA訓(xùn)練可以在CIFAR-100和Tiny-ImageNet數(shù)據(jù)集上分別獲得{1.72％，2.09％}的顯著提升。
即使對(duì)于專門設(shè)計(jì)的ViT架構(gòu)及其增強(qiáng)版本，添加EWA訓(xùn)練方案仍然可以提高它們的性能，這進(jìn)一步證實(shí)了其有效性和泛化性。例如，增強(qiáng)版Swin-T在CIFAR-100和Tiny-ImageNet數(shù)據(jù)集上的普通訓(xùn)練已經(jīng)取得了相對(duì)較高的top-1準(zhǔn)確率，分別為80.31％和64.34％。即使對(duì)于每個(gè)數(shù)據(jù)集，它仍然可以從EWA訓(xùn)練中獲得{0.21％，0.79％}的性能提升。
2、3D視覺的任務(wù)
如表3所示，提出的EWA訓(xùn)練方案在點(diǎn)云分類和語義分割任務(wù)中獲得了一致的性能提升。值得注意的是，本文的方法在分割任務(wù)中實(shí)現(xiàn)了0.52％的整體準(zhǔn)確度（OA）提升，1.40％的平均準(zhǔn)確度（mAcc）提升和1.74％的平均交并比（mIoU）提升。
5.2、EWA Fine-tuning
如表4所示，展示了標(biāo)準(zhǔn)微調(diào)和EWA微調(diào)之間的性能比較。盡管標(biāo)準(zhǔn)微調(diào)的性能已經(jīng)相當(dāng)高，但EWA微調(diào)仍然可以帶來進(jìn)一步的改進(jìn)。具體而言，在經(jīng)過EWA微調(diào)后，ViT-S和ViT-B在{Food-101，CIFAR-100}數(shù)據(jù)集上分別獲得了{(lán)0.38％，0.45％}和{0.54％，0.71％}的性能提升。
5.3、消融研究
Share schedule and Share rate
在上述內(nèi)容中，默認(rèn)情況下共享計(jì)劃設(shè)置為線性增加，在這里將其與常數(shù)共享計(jì)劃進(jìn)行比較。對(duì)于共享率，在{0.1，0.2，0.3，0.4，0.5}范圍內(nèi)進(jìn)行調(diào)整。使用不同的共享計(jì)劃和共享率在CIFAR100上訓(xùn)練ViTS，并進(jìn)行EWA訓(xùn)練。
如圖2所示，不同的共享計(jì)劃和共享率的EWA訓(xùn)練始終可以優(yōu)于香草（即基準(zhǔn)）訓(xùn)練，其中線性增加的共享計(jì)劃在大多數(shù)情況下表現(xiàn)更好。
5.4、更多討論和細(xì)節(jié)
更有趣的是，作者發(fā)現(xiàn)在小型2D視覺數(shù)據(jù)集和3D視覺任務(wù)上，標(biāo)準(zhǔn)的MoE相對(duì)于其原始的密集對(duì)應(yīng)模型表現(xiàn)較差，而EWA技術(shù)可以無縫地應(yīng)用于MoE并幫助其取得更好的學(xué)習(xí)效果。
1、naive MoE
按照在CIFAR-100、Tiny-ImageNet、ModelNet40和S3DIS上訓(xùn)練ViT-S的設(shè)置，進(jìn)一步訓(xùn)練了V-MoE-S。V-MoE-S將ViT-S的部分FFN替換為naive top-1RoutingMoE，每隔一個(gè)塊進(jìn)行一次。將 Expert數(shù)量設(shè)置為4，負(fù)載平衡損失λ為0.01，2D數(shù)據(jù)集的容量比C為1.05，3D數(shù)據(jù)集的C為1.2。如表5所示，在使用naive訓(xùn)練時(shí)，naive V-MoE-S在不同數(shù)據(jù)集上的性能較ViT-S有所下降。
2、naive MoE + Early EWA
當(dāng)將Experts Weights Averaging（EWA）引入naive MoE時(shí)，情況發(fā)生了變化?？紤]到在整個(gè)訓(xùn)練過程中始終使用EWA最終會(huì)導(dǎo)致各個(gè) Expert之間的權(quán)重幾乎相同，只在訓(xùn)練的早期階段對(duì)naive MoE執(zhí)行EWA。默認(rèn)情況下，使用恒定的共享計(jì)劃，并且只在總訓(xùn)練輪數(shù)的前半部分執(zhí)行早期EWA。通過早期EWA訓(xùn)練，V-MoE-S的性能得到了極大的改善，在CIFAR-100和Tiny-ImageNet上分別獲得了{(lán)2.14%，3.63%}的top-1準(zhǔn)確率提升。
此外，早期EWA訓(xùn)練還可以提高V-MoE-S在各種3D視覺任務(wù)上的性能，ModelNet40分類任務(wù)上準(zhǔn)確率提升了0.36%，S3DIS Area5分割任務(wù)上OA提升了0.30%，mAcc提升了1.40%，mIoU提升了1.31%。
7、總結(jié)
總之，本文為ViTs提出的訓(xùn)練方案在不增加推理延遲和參數(shù)的情況下實(shí)現(xiàn)了性能提升。通過在訓(xùn)練過程中設(shè)計(jì)高效的MoEs和EWA，并在推理過程中通過對(duì) Expert的平均來將每個(gè)MoE轉(zhuǎn)化為FFN，將ViTs的訓(xùn)練和推理階段解耦。在各種2D和3D視覺任務(wù)、ViT架構(gòu)和數(shù)據(jù)集上進(jìn)行的全面實(shí)驗(yàn)表明了其有效性和通用性。
理論分析進(jìn)一步支持了本文的方法，作者的訓(xùn)練方案還可以應(yīng)用于fine-tuning ViTs，并在各種視覺任務(wù)中改善naive MoE的效果。
8、參考
[1].Experts Weights Averaging: A New General Training Scheme for Vision Transformers.
9、推薦閱讀
ADAS落地 | 自動(dòng)駕駛的硬件加速
Fast-BEV的CUDA落地 | 5.9ms即可實(shí)現(xiàn)環(huán)視BEV 3D檢測(cè)落地！代碼開源
南京大學(xué)提出量化特征蒸餾方法QFD | 完美結(jié)合量化與蒸餾，讓AI落地更進(jìn)一步！??！
掃碼加入??「集智書童」交流群
（備注：方向+學(xué)校/公司+昵稱）
前沿AI視覺感知全棧知識(shí)??「分類、檢測(cè)、分割、關(guān)鍵點(diǎn)、車道線檢測(cè)、3D視覺（分割、檢測(cè)）、多模態(tài)、目標(biāo)跟蹤、NerF」
歡迎掃描上方二維碼，加入「集智書童-知識(shí)星球」，日常分享論文、學(xué)習(xí)筆記、問題解決方案、部署方案以及全棧式答疑，期待交流！