微軟提出MiniViT | 把DeiT壓縮9倍，性能依舊超越ResNet等卷積網(wǎng)絡(luò)

Vision Transformer由于其較高的模型性能在計算機視覺領(lǐng)域受到廣泛關(guān)注。然而，Vision Transformer受到大量參數(shù)的影響，限制了它們在內(nèi)存有限的設(shè)備上的適用性。為了緩解這一問題，本文提出了一種新的壓縮框架MiniViT，MiniViT能夠在保持相同性能的同時實現(xiàn)了Vision Transformer的參數(shù)縮減。

MiniViT的核心思想是將連續(xù)Vision TRansformer Block的權(quán)重相乘。更具體地說，使權(quán)重跨層共享，同時對權(quán)重進行轉(zhuǎn)換以增加多樣性。Weight distillation也被應(yīng)用于將知識從Large-scale ViT模型轉(zhuǎn)移到權(quán)重復用的緊湊模型。

綜合實驗證明了MiniViT的有效性，MiniViT可以將預(yù)訓練的Swin-B Transformer的尺寸減少48%，同時在ImageNet上Top-1準確率提高了1.0%。此外，使用單層參數(shù)，MiniViT能夠?qū)?code style="font-size: 14px;padding: 2px 4px;border-radius: 4px;margin-right: 2px;margin-left: 2px;color: rgb(30, 107, 184);background-color: rgba(27, 31, 35, 0.05);font-family: "Operator Mono", Consolas, Monaco, Menlo, monospace;word-break: break-all;">DeiT-B壓縮9.7倍，從86M到9M的參數(shù)，而不會嚴重影響性能。最后，通過MiniViT在下游基準上的性能來驗證其可遷移性。

1簡介

大規(guī)模預(yù)訓練的Vision TRansformer，如ViT, CvT和Swin，由于其高性能和下游任務(wù)的優(yōu)越性能，最近引起了極大的關(guān)注。然而，它們通常涉及巨大的模型尺寸和大量的訓練數(shù)據(jù)。例如，ViT需要使用3億張圖像來訓練一個帶有6.32億參數(shù)的巨大模型，才實現(xiàn)了圖像分類的最先進性能。同時，Swin使用2-3億個參數(shù)，并在ImageNet-22K上進行了預(yù)訓練，以在下游檢測和分割任務(wù)上取得良好的性能。

數(shù)以億計的參數(shù)消耗了相當大的存儲和內(nèi)存，這使得這些模型不適合涉及有限計算資源的應(yīng)用程序，如邊緣和物聯(lián)網(wǎng)設(shè)備，或者需要實時預(yù)測的任務(wù)。最近的研究表明，大規(guī)模的預(yù)訓練模型是過度參數(shù)化的。因此，在不影響這些預(yù)訓練模型性能的情況下，消除冗余參數(shù)和計算開銷是必要的。

權(quán)重共享是一種簡單且有效的減少模型尺寸的技術(shù)。神經(jīng)網(wǎng)絡(luò)中權(quán)重共享的最初想法是在20世紀90年代由LeCun和Hinton提出的，最近被重新發(fā)明用于自然語言處理(NLP)中的Transformer模型壓縮。最具代表性的工作是ALBERT，它引入了一種跨層權(quán)重共享的方法，以防止參數(shù)的數(shù)量隨著網(wǎng)絡(luò)深度的增加而增長。該技術(shù)可以在不嚴重影響模型性能的情況下顯著降低模型尺寸，從而提高參數(shù)效率。然而，Weight sharing在Vision Transformer壓縮中的有效性尚未得到很好的探索。

為了驗證這一點，作者在DeiT-S和Swin-B Transformer上執(zhí)行跨層權(quán)重共享。出乎意料的是，這種直接使用權(quán)重共享帶來了2個嚴重的問題:

1. 訓練不穩(wěn)定：作者觀察到，Weight sharing across transformer layers使訓練變得不穩(wěn)定，甚至隨著共享層數(shù)量的增加，導致訓練坍塌;

2. 性能下降：權(quán)重共享 Vision Transformer的性能與Vision Transformer相比有明顯下降。例如，雖然權(quán)重共享可以將模型參數(shù)的數(shù)量減少4倍，但是它還是帶來了Swin-s Transformer精度下降5.6%，。

為了調(diào)查這些結(jié)果的潛在原因，作者分析了訓練過程中梯度的?2-范數(shù)以及模型權(quán)重前后的中間特征表征之間的相似性。作者發(fā)現(xiàn)，在不同層上完全相同的權(quán)重是問題的主要原因。特別是，在權(quán)重共享過程中，不同Transformer Block中的層歸一化不應(yīng)該完全相同，因為不同層的特征具有不同的尺度和統(tǒng)計量。同時，梯度的?2-范數(shù)變大，并在不同層間波動，導致訓練不穩(wěn)定。

最后，Central Kernel Alignment(CKA)值（一個流行的相似性度量）在最后幾層顯著下降，表明模型在權(quán)重共享前后生成的特征圖相關(guān)性較小，這可能是性能下降的原因。

在本文中提出了一種新的技術(shù)，稱為Weight Multiplexing，來解決上述問題。Weight Multiplexing由Weight Transformation和Weight Distillation兩個組件組成，共同壓縮預(yù)訓練好的Vision Transformer。

Weight Transformation的關(guān)鍵思想是對共享的權(quán)值進行轉(zhuǎn)換，使不同層的權(quán)值略有不同，如圖2所示。該操作不僅可以促進參數(shù)的多樣性，而且還可以提高訓練的穩(wěn)定性。

更具體地說，對每個權(quán)重共享Transformer Layer的多頭自注意力(MSA)模塊和多層感知器(MLP)模塊進行了簡單的線性變換。每一層都包含單獨的變換矩陣，因此MLP對應(yīng)的注意力權(quán)重和輸出在不同層間是不同的。與共享相同的參數(shù)相比，不同層的層歸一化也是分開的。因此，可以讓Weight Sharing Transformer網(wǎng)絡(luò)的優(yōu)化變得更加穩(wěn)定。

為了防止性能下降，作者進一步用Weight Distillation設(shè)計了Weight Multiplexing，這樣嵌入在預(yù)訓練的模型中的信息可以轉(zhuǎn)移到權(quán)重共享的小模型中，這就可以產(chǎn)生更緊湊和更輕的模型。與之前僅依賴于Prediction-Level蒸餾的工作相比，本文的方法同時考慮了Attention-Level和Hidden-State蒸餾，允許較小的模型更好地模擬原始預(yù)訓練的大型教師模型的行為。

實驗表明，Weight Multiplexing方法在Baseline上實現(xiàn)了明顯的精度提高，并將預(yù)訓練好的Vision Transformer壓縮了2倍。例如，通過提出的Weight Multiplexing，12層的Mini-Swin-B模型比24層Swin-B高0.8%。此外，具有9M參數(shù)的MiniDeiT-B在ImageNet上達到了79.8%的Top-1位精度，比DeiT-B小9.7倍。用本文的方法壓縮得到的12M微型模型可以很好地遷移到下游目標檢測，在COCO驗證集上實現(xiàn)了48.6的AP，這與使用28M參數(shù)的原始Swin-T相當。

主要貢獻

1. 系統(tǒng)地研究了權(quán)重共享在Vision Transformer中的有效性，并分析了權(quán)重共享帶來問題的原因;

2. 提出了一種新的通用Vision Transformer壓縮框架MiniViT。實驗結(jié)果表明，MiniViT可以在不損失精度的前提下獲得較大的壓縮比。此外，MiniViT的性能也可以很好地遷移到下游任務(wù)。

2師從何處？

2.1 Vision Transformer

Transformer雖然最初是為NLP設(shè)計的，但最近在計算機視覺方面也顯示出了巨大的潛力。Vision Transformer首先將輸入圖像分割成一系列被稱為Token的2D Patch。然后，使用線性投影或堆疊的CNN層將這些Patch展開并轉(zhuǎn)換為d維向量(也稱為Patch Embeddings)。為了保留位置信息，Positional Embeddings被添加到Patch Embeddings中。然后將組合的Embeddings輸入到Transformer編碼器。最后，使用一個線性層來產(chǎn)生最終的分類。

Transformer編碼器由MSA和MLP的交替組成。在每個塊前后分別應(yīng)用層歸一化(LN)和殘差連接。詳細說明MSA和MLP塊如下。

1、MSA

設(shè)M為Head的數(shù)量，也稱為自注意力模塊。給定輸入序列，在第k個Head中，通過線性投影生成Query、Key和Value，分別用、和表示，其中N是Token的數(shù)量。D和d分別是Patch Embeddings和Q-K-V矩陣的維數(shù)。然后，計算序列中每個位置的所有值的加權(quán)和。這些權(quán)重被稱為注意力權(quán)重，用表示，是基于序列中2個元素之間的成對相似性，即

其中，對輸入矩陣的每一行進行softmax(·)操作。最后，將一個全連接層應(yīng)用于所有Head的輸出的連接。

2、MLP

MLP塊由2個FC層組成，其激活函數(shù)用σ(·)表示，通常為GELU。設(shè)為MLP的輸入。MLP的輸出可以表示為

其中，、、、分別為第1層和第2層的權(quán)重和偏差。通常設(shè)置d'>d。

2.2 Weight Sharing

權(quán)重共享是一種簡單而有效的提高參數(shù)效率的方法。其核心思想是跨層共享參數(shù)，如圖2(a)所示從數(shù)學上講，權(quán)重共享可以表述為一個Transformer Block f（即一個共享層）的遞歸更新：

其中為序列在第i層中嵌入的特征，L為層總數(shù)，θ為各層間Transformer Block的共享權(quán)值。

很多研究工作在自然語言Transformer模型中探索并驗證了權(quán)重共享的有效性。它可以防止參數(shù)的數(shù)量隨著網(wǎng)絡(luò)深度的增加而增加，而不會嚴重影響性能，從而提高參數(shù)的效率。

3青出于藍

3.1 Weight Multiplexing

權(quán)重共享的潛力已經(jīng)在自然語言處理中得到了證實;然而，其在Vision Transformer中的作用尚不清楚。為了檢驗這一點，作者直接將跨層權(quán)重共享應(yīng)用于DeiT-S和Swin-B模型，并觀察2個問題:

• 訓練不穩(wěn)定
• 性能下降

根據(jù)作者的實驗分析，不同層之間權(quán)值的嚴格一致性是問題的主要原因。其中，權(quán)重共享后的梯度?2-范數(shù)較大，在不同的Transformer Block內(nèi)出現(xiàn)波動，如圖4所示。

如圖5所示，從CKA值可以看出，權(quán)重共享后模型生成的Feature map與原模型的相關(guān)性較小。為了解決這些問題，受到電信領(lǐng)域多路復用技術(shù)的啟發(fā)，提出了一種新的Transformer Compression技術(shù)，Weight Multiplexing。它將多層權(quán)重組合為共享部分的單個權(quán)重，同時涉及轉(zhuǎn)換和蒸餾，以增加權(quán)重的多樣性。

更具體地說，如圖2(b)所示，本文提出的權(quán)重復用方法包括:

1. 在多個Transformer Block之間共享權(quán)重，可以認為在復用時是一個組合過程;

2. 在每一層引入轉(zhuǎn)換來模擬解復用;

3. 運用知識蒸餾增加壓縮前后模型之間特征表示的相似性。

根據(jù)Eq.(4)，可以將權(quán)重復用重新表述如下:

其中，表示第個Transformer Layer中Transformer Block的權(quán)值。請注意，中的參數(shù)數(shù)量遠少于。

1、Weight Transformation

Weight Transformation是強加于注意力矩陣和前饋網(wǎng)絡(luò)。這種轉(zhuǎn)換允許每一層都是不同的，從而提高了參數(shù)多樣性和模型表示能力。

如圖3(右)所示，各層之間不共享Transformation Kernels的參數(shù)，而原始Transformer中除了LayerNorm以外的所有其他層都是共享的。由于共享層占據(jù)了模型參數(shù)的絕大部分，權(quán)重復用后模型大小僅略有增加。

Transformation for MSA

為了提高參數(shù)的多樣性，分別在Softmax模塊前后插入了2個線性變換。

其中，分別為Softmax前后的線性變換。這種線性變換可以使每個注意力矩陣不同，同時結(jié)合注意力Head之間的信息來增加參數(shù)方差。

Transformation for MLP

另一方面，進一步對MLP進行了輕量級轉(zhuǎn)換，以提高參數(shù)的多樣性。特別地，設(shè)輸入為，其中表示所有Token的嵌入向量的第1個位置。然后引入線性變換將Y轉(zhuǎn)換為，其中是線性層的獨立權(quán)重矩陣。然后是等式3被重新表述為:

為了減少參數(shù)的數(shù)量并在變換中引入局域性，本文采用Depth-wise Convolutional來稀疏和共享每個權(quán)重矩陣中的權(quán)值，并且參數(shù)量只有比(K<

理論上，通過這些轉(zhuǎn)換，權(quán)重共享層可以恢復預(yù)訓練模型的行為，類似于解復用過程。這樣可以緩解訓練不穩(wěn)定性和性能下降問題，因為這些問題在原始模型中沒有觀察到。類似的轉(zhuǎn)換也被應(yīng)用于提高沒有權(quán)重共享的Transformer的性能，如Talking-heads Attention和CeiT。

2、Weight Distillation

為了壓縮大型預(yù)訓練模型并解決權(quán)重共享導致的性能下降問題，作者進一步采用權(quán)重蒸餾法將知識從大型模型轉(zhuǎn)移到小型且緊湊的模型。考慮了Transformer Block的3種蒸餾方法：

• Prediction-Logit Distillation
• Self-Attention Distillation
• Hidden-State Distillation

Prediction-Logit Distillation

Hinton等人證明了深度學習模型可以通過模仿訓練過程中表現(xiàn)良好的教師模型的輸出行為來獲得更好的表現(xiàn)。作者利用這個想法來引入一個預(yù)測損失，如下所示：

其中，和分別是學生模型和教師模型預(yù)測的對數(shù)，T是控制對數(shù)平滑度的溫度值。在實驗中，設(shè)置T=1。CE表示交叉熵損失。

Self-Attention Distillation

最近有研究表明，利用Transformer Layer中的注意力圖來指導學生模型的訓練是有效。為了解決學生模型和教師模型之間由于Head num不同所導致的維度不一致問題，受Minilmv2啟發(fā)，在MSA中對Query、Key和Value之間的關(guān)系應(yīng)用了交叉熵損失。

首先在所有的Head上附加矩陣。例如，定義，以同樣的方式定義。為了簡化符號，分別用、和來分別表示Q、K和V。然后，可以生成由定義的9個不同的關(guān)系矩陣。注意，是注意力矩陣a，Self-Attention Distillation損失可以表示為：

其中，表示的第n行。

Hidden-State Distillation

類似地，可以生成隱藏狀態(tài)的關(guān)系矩陣，即由MLP輸出的特征。用表示Transformer Layer的隱藏狀態(tài)，將基于關(guān)系矩陣的隱藏狀態(tài)蒸餾損失定義為:

其中表示第n行，計算公式為。

總的蒸餾損失

根據(jù)作者的觀察，僅使用prediction soft labels比同時使用prediction soft labels + ground truth labels能夠產(chǎn)生更好的性能，因此最終蒸餾目標函數(shù)表示為：

其中，β和γ分別為超參數(shù)，默認值分別為1和0.1。

3.2 Compression Pipeline

階段1：通過權(quán)重轉(zhuǎn)換生成緊湊的架構(gòu)

給定一個預(yù)訓練的Vision Transformer模型，首先在共享每個相鄰Transformer Layer中除了LayerNorm以外的參數(shù)。然后通過在Softmax層之前和之后插入一個線性層對每一層進行權(quán)值變換。此外，還為MLP引入了一個Depth-wise Convolutional。這些線性層和轉(zhuǎn)換塊的參數(shù)不共享。

階段2：用Weight Distillation訓練壓縮后的模型

在這一步中使用所提出的Weight Distillation方法，將知識從大的預(yù)訓練模型轉(zhuǎn)移到小的模型。Transformer Block內(nèi)部的這種精餾使得學生網(wǎng)絡(luò)能夠再現(xiàn)教師網(wǎng)絡(luò)的行為，從而從大規(guī)模的預(yù)訓練模型中提取出更多有用的知識。

請注意，只有當教師模型和學生模型都是Transformer架構(gòu)時，才會執(zhí)行此操作。在其他情況下，學生和教師的結(jié)構(gòu)是異質(zhì)的，只保留Prediction-Logit Distillation。

4實驗

4.1 消融實驗

1、Weight Sharing

如圖4所示，DeiT-S和Swin-B 權(quán)重共享后梯度的?2-范數(shù)較大，說明權(quán)值大小變化較快。此外，權(quán)重共享還導致了不同層間的梯度范數(shù)的波動。這可能會導致不同的層的優(yōu)化空間。特別是，一些層更新得很快，而其他部分幾乎沒有優(yōu)化，這使得模型很可能收斂到一個糟糕的局部最優(yōu)，甚至在訓練中出現(xiàn)分歧。

因此，不同層間共享的嚴格相同的權(quán)重會導致訓練的不穩(wěn)定性。然而，權(quán)值復用方法可以通過引入變換來提高參數(shù)多樣性，從而降低梯度范數(shù)和提高層間的平滑度，促進更穩(wěn)定的訓練過程。

在性能分析方面，在圖5中對權(quán)值共享和權(quán)重復用之間的特征相似性與CKA進行了比較。CKA值越高，表示兩個模型之間的特征表示越相似，從而獲得相似的性能。可以觀察到，DeiT和Swin在應(yīng)用權(quán)重共享后都存在很大的特征表示偏差，特別是在最后幾層，這可能是權(quán)重共享導致性能下降的原因之一。然而提出的權(quán)值復用方法可以提高相似性。

2、Component-wise

采用權(quán)值共享方法后，Swin-T和DeiT-S模型的參數(shù)數(shù)量減少了一半，精度也下降了2%。性能可以通過應(yīng)用權(quán)重精餾或權(quán)重轉(zhuǎn)換來提高，這表明了它們各自的有效性。值得注意的是權(quán)值變換只引入了少量的參數(shù)。

此外，當結(jié)合這3個組件時，MiniViT可以達到最好的精度，甚至超過原始的Swin-T和DeiT-S。

3、Number of Sharing Blocks

在表2中，在Swin-T或DeiT-B的每個階段共享2個Block可以顯著將參數(shù)量從28M減少到16M，86M減少到44M，而Top-1的精度提高了1%。在極端情況下，每個階段的所有塊都共享，Mini-Swin-T仍然可以以43%的參數(shù)優(yōu)于原始模型。Mini-DeiT-B可以實現(xiàn)90%的參數(shù)降低，性能僅下降2%。

結(jié)果表明，Vision Transformer存在參數(shù)冗余，所提出的MiniViT可以有效地提高參數(shù)效率。此外，MiniViT是可配置的，以滿足模型大小和性能的各種要求。

4、Distillation Losses

如表3所示，與僅使用預(yù)測損失相比，額外的GT標簽導致Swin的性能下降了0.3%，這是由于權(quán)值共享導致學習能力下降。此外，還觀察到，在應(yīng)用Self-Attention Distillation和Hidden-State Distillation后，精確度提高了約0.2%，表明提出的Distillation方法是有效性的。

4.2 SoTA分類

4.3 遷移學習

5參考

[1].MiniViT: Compressing Vision Transformers with Weight Multiplexing

6推薦閱讀

Sparse R-CNN升級版 | Dynamic Sparse R-CNN使用ResNet50也能達到47.2AP

CVPR2022 Oral | CosFace、ArcFace的大統(tǒng)一升級,AdaFace解決低質(zhì)量圖像人臉識

掃描下方二維碼關(guān)注【集智書童】公眾號，獲取更多實踐項目源碼和論文解讀，非常期待你我的相遇，讓我們以夢為馬，砥礪前行！