如何一次性訓(xùn)練100,000+個Vision Transformers？
本文轉(zhuǎn)自雷鋒網(wǎng)，如需轉(zhuǎn)載請至雷鋒網(wǎng)官網(wǎng)申請授權(quán)。
經(jīng)過漫長的等待，ICCV 2021終于迎來放榜時刻！
ICCV官方在推特上公布了這一消息，并表示今年共有6236篇投稿，最終1617篇論文被接收，接收率為25.9%，相比于2017年（約29%），保持了和2019年相當(dāng)?shù)妮^低水平。
而投稿量則依舊逐年大幅增長，從2017年的2143篇，到2109年的4328篇，再到如今的6236篇，相比上一屆多了50%左右。
你看郵件的時候是這表情嗎？
不得不說，官方皮起來也是接地氣、真扎心、沒誰了哈哈~
論文ID地址：https://docs.google.com/spreadsheets/u/1/d/e/2PACX-1vRfaTmsNweuaA0Gjyu58H_Cx56pGwFhcTYII0u1pg0U7MbhlgY0R6Y-BbK3xFhAiwGZ26u3TAtN5MnS/pubhtml
也就在今天，AI科技評論發(fā)現(xiàn)了一項(xiàng)非常厲害的研究，號稱可一次性訓(xùn)練10萬個ViT，論文也剛剛喜提ICCV accepted！
近來，Vision Transformer (ViT) 模型在諸多視覺任務(wù)中展現(xiàn)出了強(qiáng)大的表達(dá)能力和潛力。
紐約州立大學(xué)石溪分校與微軟亞洲研究院的研究人員提出了一種新的網(wǎng)絡(luò)結(jié)構(gòu)搜索方法AutoFormer，用來自動探索最優(yōu)的ViT模型結(jié)構(gòu)。
AutoFormer能一次性訓(xùn)練大量的不同結(jié)構(gòu)的ViT模型，并使得它們的性能達(dá)到收斂。
其搜索出來的結(jié)構(gòu)對比手工設(shè)計的ViT模型有較明顯的性能提升。
方法亮點(diǎn)：
-
同時訓(xùn)練大量Vision Transformers模型，使其性能接近單獨(dú)訓(xùn)練；
-
簡單有效，能夠靈活應(yīng)用于Vision Transformer的變種搜索；
-
性能較ViT, DeiT等模型有較明顯提升。
論文地址：https://arxiv.org/abs/2107.00651
代碼地址：https://github.com/microsoft/AutoML/tree/main/AutoFormer
1. 引言
最近的研究發(fā)現(xiàn)，ViT能夠從圖像中學(xué)習(xí)強(qiáng)大的視覺表示，并已經(jīng)在多個視覺任務(wù)（分類，檢測，分割等）上展現(xiàn)出了不俗的能力。
然而，Vision Transformer 模型的結(jié)構(gòu)設(shè)計仍然比較困難。例如，如何選擇最佳的網(wǎng)絡(luò)深度、寬度和多頭注意力中的頭部數(shù)量？
作者的實(shí)驗(yàn)發(fā)現(xiàn)這些因素都和模型的最終性能息息相關(guān)。然而，由于搜索空間非常龐大，我們很難人為地找到它們的最佳組合。
圖1: 不同搜索維度的變化會極大地影響模型的表現(xiàn)能力
本文的作者提出了一種專門針對Vision Transformer 結(jié)構(gòu)的新的Neural Architecture Search (NAS) 方法 AutoFormer。AutoFormer大幅節(jié)省了人為設(shè)計結(jié)構(gòu)的成本，并能夠自動地快速搜索不同計算限制條件下ViT模型各個維度的最佳組合，這使得不同部署場景下的模型設(shè)計變得更加簡單。
圖2: AutoFormer的結(jié)構(gòu)示意圖，在每一個訓(xùn)練迭代中，超網(wǎng)會動態(tài)變化并更新相應(yīng)的部分權(quán)重
2. 方法
常見的One-shot NAS 方法[1, 2, 3]通常采取權(quán)重共享的方式來節(jié)省計算開銷，搜索空間被編碼進(jìn)一個權(quán)重共享的超網(wǎng) (supernet) 中，并運(yùn)用超網(wǎng)權(quán)重作為搜索空間中結(jié)構(gòu)權(quán)重的一個估計。其具體搜索過程可分為兩個步驟，第一步是更新超網(wǎng)的權(quán)重，如下公式所示。
第二步是利用訓(xùn)練好的超網(wǎng)權(quán)重來對搜索空間中結(jié)構(gòu)進(jìn)行搜索。
在實(shí)驗(yàn)的過程中，作者發(fā)現(xiàn)經(jīng)典One-shot NAS方法的權(quán)重共享方式很難被有效地運(yùn)用到Vision Transformer的結(jié)構(gòu)搜索中。這是因?yàn)橹暗姆椒ㄍǔH僅共享結(jié)構(gòu)之間的權(quán)重，而解耦同一層中不同算子的權(quán)重。
如圖3所示，在Vision Transformer的搜索空間中，這種經(jīng)典的策略會遇到收斂緩慢和性能較低的困難。
圖3 權(quán)重糾纏和經(jīng)典權(quán)重共享的訓(xùn)練以及測試對比
受到OFA [4], BigNAS [5] 以及Slimmable networks [6, 7] 等工作的啟發(fā)，作者提出了一種新的權(quán)重共享方式——權(quán)重糾纏 (Weight Entanglement)。
如圖4所示，權(quán)重糾纏進(jìn)一步共享不同結(jié)構(gòu)之間的權(quán)重，使得同一層中不同算子之間能夠互相影響和更新，實(shí)驗(yàn)證明權(quán)重糾纏對比經(jīng)典的權(quán)重共享方式，擁有占用顯存少，超網(wǎng)收斂快和超網(wǎng)性能高的優(yōu)勢。
同時，由于權(quán)重糾纏，不同算子能夠得到更加充分的訓(xùn)練，這使得AutoFormer能夠一次性訓(xùn)練大量的ViT模型，且使其接近收斂。（詳情見實(shí)驗(yàn)部分）
圖4 權(quán)重糾纏和權(quán)重共享的對比示意圖
3. 實(shí)驗(yàn)
作者設(shè)計了一個擁有超過1.7x10^17備選結(jié)構(gòu)的巨大搜索空間，其搜索維度包括ViT模型中的五個主要的可變因素：寬度 (embedding dim)、Q-K-V 維度 (Q-K-V dimension)、頭部數(shù)量 (head number)、MLP 比率 (MLP ratio) 和網(wǎng)絡(luò)深度 (network depth)，詳見表1。
表1：AutoFormer的搜索空間
為了驗(yàn)證方法的有效性，作者將AutoFormer搜索得到的結(jié)構(gòu)和近期提出的ViT模型以及經(jīng)典的CNN模型在ImageNet上進(jìn)行了比較。
對于訓(xùn)練過程，作者采取了DeiT [8]類似的數(shù)據(jù)增強(qiáng)方法，如 Mixup, Cutmix, RandAugment等, 超網(wǎng)的具體訓(xùn)練參數(shù)如表2所示。所有模型都是在 16塊Tesla V100 GPU上進(jìn)行訓(xùn)練和測試的。
表2 超網(wǎng)的訓(xùn)練參數(shù)
如圖5 和表3所示，搜索得到的結(jié)構(gòu)在ImageNet數(shù)據(jù)集上明顯優(yōu)于已有的ViT模型。
表3：各個模型在ImageNet 測試集上的結(jié)果
從表4中可以看出，在下游任務(wù)中，AutoFormer依然表現(xiàn)出色，利用僅僅25%的計算量就超越了已有的ViT和DeiT模型，展現(xiàn)了其強(qiáng)大的泛化性能力。
表4：下游分類任務(wù)遷移學(xué)習(xí)的結(jié)果
同時，如圖5所示，利用權(quán)重糾纏，AutoFormer能夠同時使得成千上萬個Vision Transformers模型得到很好的訓(xùn)練（藍(lán)色的點(diǎn)代表從搜索空間中選出的1000個較好的結(jié)構(gòu)）。
不僅僅使得其在搜索后不再需要重新訓(xùn)練（retraining）結(jié)構(gòu)，節(jié)約了搜索時間，也使得其能在各種不同的計算資源限制下快速搜索最優(yōu)結(jié)構(gòu)。
圖5：左：AutoFormer能夠同時訓(xùn)練大量結(jié)構(gòu)，并使得其接近收斂。藍(lán)色的點(diǎn)代表從搜索空間中選出的1000個較好的結(jié)構(gòu)。右：ImageNet上各模型對比
4. 結(jié)語
本文提出了一種新的專用于Vision Transformer結(jié)構(gòu)搜索的One-shot NAS方法—— AutoFormer。AutoFormer 配備了新的權(quán)重共享機(jī)制，即權(quán)重糾纏 （Weight Engtanglement）。在這種機(jī)制下，搜索空間的網(wǎng)絡(luò)結(jié)構(gòu)幾乎都能被充分訓(xùn)練，省去了結(jié)構(gòu)搜索后重新訓(xùn)練（Retraining）的時間。大量實(shí)驗(yàn)表明所提出的算法可以提高超網(wǎng)的排序能力并找到高性能的結(jié)構(gòu)。在文章的最后，作者希望通過本文給手工ViT結(jié)構(gòu)設(shè)計和NAS+Vision Transformer提供一些靈感。在未來工作，作者將嘗試進(jìn)一步豐富搜索空間，以及給出權(quán)重糾纏的理論分析。
