當(dāng)前 Vision Transformers （ViT）領(lǐng)域有兩個(gè)主要的痛點(diǎn)：1、缺少對(duì) ViT 進(jìn)行設(shè)計(jì)和擴(kuò)展的有效方法；2、訓(xùn)練 ViT 的計(jì)算成本比卷積網(wǎng)絡(luò)要大得多。

為了解決這兩個(gè)問(wèn)題，來(lái)自得克薩斯大學(xué)奧斯汀分校、悉尼科技大學(xué)和谷歌的研究者提出了 As-ViT（Auto-scaling Vision Transformers），這是一個(gè)無(wú)需訓(xùn)練的 ViT 自動(dòng)擴(kuò)展框架，它能以高效且有原則的方式自動(dòng)設(shè)計(jì)和擴(kuò)展 ViT。

論文鏈接：https://arxiv.org/abs/2202.11921

具體來(lái)說(shuō)，研究人員首先利用無(wú)訓(xùn)練搜索過(guò)程設(shè)計(jì)了 ViT 拓?fù)涞摹阜N子」，這種極快的搜索是通過(guò)對(duì) ViT 網(wǎng)絡(luò)復(fù)雜性的全面研究來(lái)實(shí)現(xiàn)的，從而產(chǎn)生了與真實(shí)準(zhǔn)確度的強(qiáng) Kendall-tau 相關(guān)性。其次，從「種子」拓?fù)溟_(kāi)始，通過(guò)將寬度 / 深度增加到不同的 ViT 層來(lái)自動(dòng)化 ViT 的擴(kuò)展規(guī)則，實(shí)現(xiàn)了在一次運(yùn)行中具有不同數(shù)量參數(shù)的一系列架構(gòu)。最后，基于 ViT 在早期訓(xùn)練階段可以容忍粗粒度 tokenization 的經(jīng)驗(yàn)，該研究提出了一種漸進(jìn)式 tokenization 策略來(lái)更快、更節(jié)約地訓(xùn)練 ViT。

作為統(tǒng)一的框架，As-ViT 在分類（ImageNet-1k 上 83.5% 的 top1）和檢測(cè)（COCO 上 52.7% 的 mAP）任務(wù)上實(shí)現(xiàn)了強(qiáng)大的性能，無(wú)需任何手動(dòng)調(diào)整或擴(kuò)展 ViT 架構(gòu)，端到端模型設(shè)計(jì)和擴(kuò)展過(guò)程在一塊 V100 GPU 上只需 12 小時(shí)。

具有網(wǎng)絡(luò)復(fù)雜度的 ViT 自動(dòng)設(shè)計(jì)和擴(kuò)展

為加快 ViT 設(shè)計(jì)并避免繁瑣的手動(dòng)工作，該研究希望以高效、自動(dòng)化和有原則的 ViT 搜索和擴(kuò)展為目標(biāo)。具體來(lái)說(shuō)有兩個(gè)問(wèn)題需要解決：1）在訓(xùn)練成本最小甚至為零的情況下，如何高效地找到最優(yōu)的 ViT 架構(gòu)拓?fù)洌?）如何擴(kuò)大 ViT 拓?fù)涞纳疃群蛯挾纫詽M足模型尺寸的不同需求？

擴(kuò)展 ViT 的拓?fù)淇臻g

在設(shè)計(jì)和擴(kuò)展之前，首先是為 As-ViT 擴(kuò)展的拓?fù)渌阉骺臻g：首先將輸入圖像嵌入到 1/4 尺度分辨率的塊中，并采用逐級(jí)空間縮減和通道加倍策略。這是為了方便密集預(yù)測(cè)任務(wù)，例如需要多尺度特征的檢測(cè)。

通過(guò)流形傳播評(píng)估初始化時(shí)的 ViT 復(fù)雜性

ViT 訓(xùn)練速度很慢，因此，通過(guò)評(píng)估訓(xùn)練模型的準(zhǔn)確率來(lái)進(jìn)行架構(gòu)搜索的成本將高得讓人難以承受。最近學(xué)界出現(xiàn)很多用基于 ReLU 的 CNN 的免訓(xùn)練神經(jīng)架構(gòu)搜索方法，利用局部線性圖 (Mellor et al., 2020)、梯度敏感性 (Abdelfattah et al., 2021)、線性區(qū)域數(shù)量 (Chen et al., 2021e;f) 或網(wǎng)絡(luò)拓?fù)洌˙hardwaj et al., 2021）等方式。

然而 ViT 配備了更復(fù)雜的非線性函數(shù)如 self-attention、softmax 和 GeLU。因此需要以更一般的方式衡量其學(xué)習(xí)能力。在新研究中，研究者考慮通過(guò) ViT 測(cè)量流形傳播的復(fù)雜性，以估計(jì)復(fù)雜函數(shù)可以如何被 ViT 逼近。直觀地說(shuō)，一個(gè)復(fù)雜的網(wǎng)絡(luò)可以在其輸出層將一個(gè)簡(jiǎn)單的輸入傳播到一個(gè)復(fù)雜的流形中，因此可能具有很強(qiáng)的學(xué)習(xí)能力。在 UT Austin 的工作中，他們通過(guò) ViT 映射簡(jiǎn)單圓輸入的多種復(fù)雜性：h(θ) = √ N [u^0 cos(θ) + u^1 sin(θ)]。這里，N 是 ViT 輸入的維度（例如，對(duì)于 ImageNet 圖像，N = 3 × 224 × 224），u^0 和 u^1 形成了圓所在的 R^N 的二維子空間的標(biāo)準(zhǔn)正交基。

搜索 ViT 拓?fù)洫?jiǎng)勵(lì)

研究者提出了基于 L^E 的免訓(xùn)練搜索（算法 1），大多數(shù) NAS（神經(jīng)架構(gòu)搜索）方法將單路徑或超級(jí)網(wǎng)絡(luò)的準(zhǔn)確率或損失值評(píng)估為代理推理。當(dāng)應(yīng)用于 ViT 時(shí)，這種基于訓(xùn)練的搜索將需要更多的計(jì)算成本。對(duì)于采樣的每個(gè)架構(gòu)，這里不是訓(xùn)練 ViT，而是計(jì)算 L^E 并將其視為指導(dǎo)搜索過(guò)程的獎(jiǎng)勵(lì)。

除了 L^E，還包括 NTK 條件數(shù) κΘ = λ_max/λ_min ，以指示 ViT 的可訓(xùn)練性（Chen et al., 2021e; Xiao et al., 2019; Yang, 2020; Hron et al., 2020）。λ_max 和 λ_min 是 NTK 矩陣 Θ 的最大和最小特征值。

搜索使用強(qiáng)化學(xué)習(xí)方法，策略被定為聯(lián)合分類分布，并通過(guò)策略梯度進(jìn)行更新，該研究將策略更新為 500 step，觀察到足以使策略收斂（熵從 15.3 下降到 5.7）。搜索過(guò)程非?？欤涸?ImageNet-1k 數(shù)據(jù)集上只有七個(gè) GPU 小時(shí) (V100)，這要?dú)w功于繞過(guò) ViT 訓(xùn)練的 L^E 的簡(jiǎn)單計(jì)算。為了解決 L^E 和 κΘ 的不同大小，該研究通過(guò)它們的相對(duì)值范圍對(duì)它們進(jìn)行歸一化（算法 1 中的第 5 行）。

表 3 總結(jié)了新搜索方法的 ViT 拓?fù)浣y(tǒng)計(jì)數(shù)據(jù)。我們可以看到 L^E 和 κΘ 高度偏好：（1）具有重疊的 token （K_1～K_4 都大于 stride ），以及（2）在更深層中更大的 FFN 擴(kuò)展率（E_1 < E_2 < E_3 < E_4）。在注意力分裂和正面數(shù)量上沒(méi)有發(fā)現(xiàn) L^E 和 κΘ 的明顯偏好。

ViT 自主的原則型擴(kuò)展

得到最優(yōu)拓?fù)浜?，接下?lái)要解決的一個(gè)問(wèn)題是：如何平衡網(wǎng)絡(luò)的深度和寬度？

目前，對(duì)于 ViT 擴(kuò)展沒(méi)有這樣的經(jīng)驗(yàn)法則。最近的工作試圖擴(kuò)大或增長(zhǎng)不同大小的卷積網(wǎng)絡(luò)以滿足各種資源限制（Liu et al., 2019a; Tan & Le, 2019）。然而，為了自動(dòng)找到一個(gè)有原則的擴(kuò)展規(guī)則，訓(xùn)練 ViT 將花費(fèi)巨大的計(jì)算成本。也可以搜索不同的 ViT 變體（如第 3.3 節(jié)中所述），但這需要多次運(yùn)行。相反，「向上擴(kuò)展，scaling-up」是在一個(gè)實(shí)驗(yàn)中生成多個(gè)模型變體的更自然的方式。因此，該研究試圖以一種免訓(xùn)練且有原則的有效方法將搜索到的基本「種子」ViT 擴(kuò)展到更大的模型。算法 2 中描述了這種自動(dòng)擴(kuò)展方法：

初始架構(gòu)的每個(gè)階段都有一個(gè)注意力塊，初始隱藏維度 C = 32。每次迭代找出最佳深度和寬度，以進(jìn)行進(jìn)一步向上擴(kuò)展。對(duì)于深度，該研究嘗試找出要加深哪個(gè)階段（即，在哪個(gè)階段添加一個(gè)注意力塊）；對(duì)于寬度，該研究嘗試發(fā)現(xiàn)最佳擴(kuò)展比（即，將通道數(shù)擴(kuò)大到什么程度）。

擴(kuò)展軌跡如下圖 3 所示。比較自主擴(kuò)展和隨機(jī)擴(kuò)展，研究者發(fā)現(xiàn)擴(kuò)展原則更喜歡舍棄深度來(lái)?yè)Q取更多寬度，使用更淺但更寬的網(wǎng)絡(luò)。這種擴(kuò)展更類似于 Zhai et al. (2021) 開(kāi)發(fā)的規(guī)則。相比之下，ResNet 和 Swin Transformer (Liu et al., 2021) 選擇更窄更深。

通過(guò)漸進(jìn)靈活的 re-tokenization 進(jìn)行高效的 ViT 訓(xùn)練

該研究通過(guò)提出漸進(jìn)靈活的 re-tokenization 訓(xùn)練策略來(lái)提供肯定的答案。為了在訓(xùn)練期間更新 token 的數(shù)量而不影響線性投影中權(quán)重的形狀，該研究在第一個(gè)線性投影層中采用不同的采樣粒度。以第一個(gè)投影核 K_1 = 4 且 stride = 4 為例：訓(xùn)練時(shí)研究者逐漸將第一個(gè)投影核的 (stride, dilation) 對(duì)逐漸變?yōu)?(16, 5), (8, 2) 和 (4 , 1)，保持權(quán)重的形狀和架構(gòu)不變。

這種 re-tokenization 的策略激發(fā)了 ViT 的課程學(xué)習(xí)（curriculum learning）：訓(xùn)練開(kāi)始時(shí)引入粗采樣以顯著減少 token 的數(shù)量。換句話說(shuō)，As-ViT 在早期訓(xùn)練階段以極低的計(jì)算成本（僅全分辨率訓(xùn)練的 13.2% FLOPs）快速?gòu)膱D像中學(xué)習(xí)粗略信息。在訓(xùn)練的后期階段，該研究逐漸切換到細(xì)粒度采樣，恢復(fù)完整的 token 分辨率，并保持有競(jìng)爭(zhēng)力的準(zhǔn)確率。如圖 4 所示，當(dāng)在早期訓(xùn)練階段使用粗采樣訓(xùn)練 ViT 時(shí)，它仍然可以獲得很高的準(zhǔn)確率，同時(shí)需要極低的計(jì)算成本。不同采樣粒度之間的轉(zhuǎn)換引入了性能的跳躍，最終網(wǎng)絡(luò)恢復(fù)了具有競(jìng)爭(zhēng)力的最終性能。

如圖 4 所示，當(dāng) ViT 在早期訓(xùn)練階段使用粗采樣訓(xùn)練 ViT 時(shí)，它仍然可以獲得很高的準(zhǔn)確率，同時(shí)需要極低的計(jì)算成本。不同采樣粒度之間的轉(zhuǎn)換引入了性能的跳躍，最終網(wǎng)絡(luò)恢復(fù)了具有競(jìng)爭(zhēng)力的最終性能。

實(shí)驗(yàn)

AS-VIT：自動(dòng)擴(kuò)展 VIT

該研究在表 4 中展示了搜索到的 As-ViT 拓?fù)?。這種架構(gòu)在第一個(gè)投影（tokenization）step 和三個(gè)重新嵌入 step 中，促進(jìn)了 token 之間的強(qiáng)烈重疊。FFN 擴(kuò)展比首先變窄，然后在更深的層變寬。利用少量注意力拆分來(lái)更好地聚合全局信息。

圖像分類

下表 5 展示了 As-ViT 與其他模型的比較。與之前基于 Transformer 和基于 CNN 的架構(gòu)相比，As-ViT 以相當(dāng)數(shù)量的參數(shù)和 FLOP 實(shí)現(xiàn)了 SOTA 性能。

高效訓(xùn)練

研究者調(diào)整了表 6 中為每個(gè) token 減少階段的時(shí)期，并將結(jié)果顯示在表 6 中。標(biāo)準(zhǔn)訓(xùn)練需要 42.8 TPU 天，而高效訓(xùn)練可節(jié)省高達(dá) 56.2% 的訓(xùn)練 FLOP 和 41.1% 的訓(xùn)練 TPU 天，仍然達(dá)到很高的準(zhǔn)確率。

拓?fù)浜蛿U(kuò)展的貢獻(xiàn)

為了更好地驗(yàn)證搜索型拓?fù)浜蛿U(kuò)展規(guī)則的貢獻(xiàn)，該研究進(jìn)行了更多的消融研究（表 7）。首先，在擴(kuò)展之前直接訓(xùn)練搜索到的拓?fù)洹Ｔ撗芯克阉鞯姆N子拓?fù)鋬?yōu)于圖 2 中 87 個(gè)隨機(jī)拓?fù)渲械淖罴淹負(fù)洹?/p>

第二，該研究將基于復(fù)雜度的規(guī)則與「隨機(jī)擴(kuò)展 + As-ViT 拓?fù)洹惯M(jìn)行比較。在不同的擴(kuò)展下，該研究的自動(dòng)擴(kuò)展也優(yōu)于隨機(jī)擴(kuò)展。

COCO 數(shù)據(jù)集上的目標(biāo)檢測(cè)

該研究將 As-ViT 與標(biāo)準(zhǔn) CNN 和之前的 Transformer 網(wǎng)絡(luò)進(jìn)行了比較。比較是通過(guò)僅更改主干而其他設(shè)置未更改來(lái)進(jìn)行的。從下表 8 的結(jié)果可以看出，As-ViT 也可以捕獲多尺度特征并實(shí)現(xiàn)最先進(jìn)的檢測(cè)性能，盡管它是在 ImageNet 上設(shè)計(jì)的，并且它的復(fù)雜性是為分類而測(cè)量的。

? ?