機器之心報道
編輯：小舟
來自德克薩斯大學奧斯汀分校、悉尼科技大學和谷歌的研究者提出了一個無需訓練就能自動擴充架構 As-ViT，其能以高效和有原則的方式自動發現和擴充 ViT。
目前 Vision Transformers （ViT）領域有兩個主要的痛點：1、缺少對 ViT 進行設計和擴充的有效方法；2、訓練 ViT 的計算成本比卷積網絡要大得多。
為了解決這兩個問題，來自得克薩斯大學奧斯汀分校、悉尼科技大學和谷歌的研究者提出了 As-ViT（Auto-scaling Vision Transformers），這是一個無需訓練的 ViT 自動擴充架構，它能以高效且有原則的方式自動設計和擴充 ViT。
論文連結：https://arxiv.org/abs/2202.11921
具體來說，研究人員首先利用無訓練搜尋過程設計了 ViT 拓撲的「種子」，這種極快的搜尋是通過對 ViT 網絡複雜性的全面研究來實作的，進而産生了與真實準确度的強 Kendall-tau 相關性。其次，從「種子」拓撲開始，通過将寬度 / 深度增加到不同的 ViT 層來自動化 ViT 的擴充規則，實作了在一次運作中具有不同數量參數的一系列架構。最後，基于 ViT 在早期訓練階段可以容忍粗粒度 tokenization 的經驗，該研究提出了一種漸進式 tokenization 政策來更快、更節約地訓練 ViT。
作為統一的架構，As-ViT 在分類（ImageNet-1k 上 83.5% 的 top1）和檢測（COCO 上 52.7% 的 mAP）任務上實作了強大的性能，無需任何手動調整或擴充 ViT 架構，端到端模型設計和擴充過程在一塊 V100 GPU 上隻需 12 小時。
具有網絡複雜度的 ViT 自動設計和擴充
為加快 ViT 設計并避免繁瑣的手動工作，該研究希望以高效、自動化和有原則的 ViT 搜尋和擴充為目标。具體來說有兩個問題需要解決：1）在訓練成本最小甚至為零的情況下，如何高效地找到最優的 ViT 架構拓撲？2）如何擴大 ViT 拓撲的深度和寬度以滿足模型尺寸的不同需求？
擴充 ViT 的拓撲空間
在設計和擴充之前，首先是為 As-ViT 擴充的拓撲搜尋空間：首先将輸入圖像嵌入到 1/4 尺度分辨率的塊中，并采用逐級空間縮減和通道加倍政策。這是為了友善密集預測任務，例如需要多尺度特征的檢測。
通過流形傳播評估初始化時的 ViT 複雜性
ViT 訓練速度很慢，是以，通過評估訓練模型的準确率來進行架構搜尋的成本将高得讓人難以承受。最近學界出現很多用基于 ReLU 的 CNN 的免訓練神經架構搜尋方法，利用局部線性圖 (Mellor et al., 2020)、梯度敏感性 (Abdelfattah et al., 2021)、線性區域數量 (Chen et al., 2021e;f) 或網絡拓撲（Bhardwaj et al., 2021）等方式。
然而 ViT 配備了更複雜的非線性函數如 self-attention、softmax 和 GeLU。是以需要以更一般的方式衡量其學習能力。在新研究中，研究者考慮通過 ViT 測量流形傳播的複雜性，以估計複雜函數可以如何被 ViT 逼近。直覺地說，一個複雜的網絡可以在其輸出層将一個簡單的輸入傳播到一個複雜的流形中，是以可能具有很強的學習能力。在 UT Austin 的工作中，他們通過 ViT 映射簡單圓輸入的多種複雜性：h(θ) = √ N [u^0 cos(θ) + u^1 sin(θ)]。這裡，N 是 ViT 輸入的次元（例如，對于 ImageNet 圖像，N = 3 × 224 × 224），u^0 和 u^1 形成了圓所在的 R^N 的二維子空間的标準正交基。
搜尋 ViT 拓撲獎勵
研究者提出了基于 L^E 的免訓練搜尋（算法 1），大多數 NAS（神經架構搜尋）方法将單路徑或超級網絡的準确率或損失值評估為代理推理。當應用于 ViT 時，這種基于訓練的搜尋将需要更多的計算成本。對于采樣的每個架構，這裡不是訓練 ViT，而是計算 L^E 并将其視為指導搜尋過程的獎勵。
除了 L^E，還包括 NTK 條件數 κΘ = λ_max/λ_min ，以訓示 ViT 的可訓練性（Chen et al., 2021e; Xiao et al., 2019; Yang, 2020; Hron et al., 2020）。λ_max 和 λ_min 是 NTK 矩陣 Θ 的最大和最小特征值。
搜尋使用強化學習方法，政策被定為聯合分類分布，并通過政策梯度進行更新，該研究将政策更新為 500 step，觀察到足以使政策收斂（熵從 15.3 下降到 5.7）。搜尋過程非常快：在 ImageNet-1k 資料集上隻有七個 GPU 小時 (V100)，這要歸功于繞過 ViT 訓練的 L^E 的簡單計算。為了解決 L^E 和 κΘ 的不同大小，該研究通過它們的相對值範圍對它們進行歸一化（算法 1 中的第 5 行）。
表 3 總結了新搜尋方法的 ViT 拓撲統計資料。我們可以看到 L^E 和 κΘ 高度偏好：（1）具有重疊的 token （K_1 K_4 都大于 stride ），以及（2）在更深層中更大的 FFN 擴充率（E_1
ViT 自主的原則型擴充
得到最優拓撲後，接下來要解決的一個問題是：如何平衡網絡的深度和寬度？
目前，對于 ViT 擴充沒有這樣的經驗法則。最近的工作試圖擴大或增長不同大小的卷積網絡以滿足各種資源限制（Liu et al., 2019a; Tan & Le, 2019）。然而，為了自動找到一個有原則的擴充規則，訓練 ViT 将花費巨大的計算成本。也可以搜尋不同的 ViT 變體（如第 3.3 節中所述），但這需要多次運作。相反，「向上擴充，scaling-up」是在一個實驗中生成多個模型變體的更自然的方式。是以，該研究試圖以一種免訓練且有原則的有效方法将搜尋到的基本「種子」ViT 擴充到更大的模型。算法 2 中描述了這種自動擴充方法：
初始架構的每個階段都有一個注意力塊，初始隐藏次元 C = 32。每次疊代找出最佳深度和寬度，以進行進一步向上擴充。對于深度，該研究嘗試找出要加深哪個階段（即，在哪個階段添加一個注意力塊）；對于寬度，該研究嘗試發現最佳擴充比（即，将通道數擴大到什麼程度）。
擴充軌迹如下圖 3 所示。比較自主擴充和随機擴充，研究者發現擴充原則更喜歡舍棄深度來換取更多寬度，使用更淺但更寬的網絡。這種擴充更類似于 Zhai et al. (2021) 開發的規則。相比之下，ResNet 和 Swin Transformer (Liu et al., 2021) 選擇更窄更深。
通過漸進靈活的 re-tokenization 進行高效的 ViT 訓練
該研究通過提出漸進靈活的 re-tokenization 訓練政策來提供肯定的答案。為了在訓練期間更新 token 的數量而不影響線性投影中權重的形狀，該研究在第一個線性投影層中采用不同的采樣粒度。以第一個投影核 K_1 = 4 且 stride = 4 為例：訓練時研究者逐漸将第一個投影核的 (stride, dilation) 對逐漸變為 (16, 5), (8, 2) 和 (4 , 1)，保持權重的形狀和架構不變。
這種 re-tokenization 的政策激發了 ViT 的課程學習（curriculum learning）：訓練開始時引入粗采樣以顯着減少 token 的數量。換句話說，As-ViT 在早期訓練階段以極低的計算成本（僅全分辨率訓練的 13.2% FLOPs）快速從圖像中學習粗略資訊。在訓練的後期階段，該研究逐漸切換到細粒度采樣，恢複完整的 token 分辨率，并保持有競争力的準确率。如圖 4 所示，當在早期訓練階段使用粗采樣訓練 ViT 時，它仍然可以獲得很高的準确率，同時需要極低的計算成本。不同采樣粒度之間的轉換引入了性能的跳躍，最終網絡恢複了具有競争力的最終性能。
如圖 4 所示，當 ViT 在早期訓練階段使用粗采樣訓練 ViT 時，它仍然可以獲得很高的準确率，同時需要極低的計算成本。不同采樣粒度之間的轉換引入了性能的跳躍，最終網絡恢複了具有競争力的最終性能。
實驗
AS-VIT：自動擴充 VIT
該研究在表 4 中展示了搜尋到的 As-ViT 拓撲。這種架構在第一個投影（tokenization）step 和三個重新嵌入 step 中，促進了 token 之間的強烈重疊。FFN 擴充比首先變窄，然後在更深的層變寬。利用少量注意力拆分來更好地聚合全局資訊。
圖像分類
下表 5 展示了 As-ViT 與其他模型的比較。與之前基于 Transformer 和基于 CNN 的架構相比，As-ViT 以相當數量的參數和 FLOP 實作了 SOTA 性能。
高效訓練
研究者調整了表 6 中為每個 token 減少階段的時期，并将結果顯示在表 6 中。标準訓練需要 42.8 TPU 天，而高效訓練可節省高達 56.2% 的訓練 FLOP 和 41.1% 的訓練 TPU 天，仍然達到很高的準确率。
拓撲和擴充的貢獻
為了更好地驗證搜尋型拓撲和擴充規則的貢獻，該研究進行了更多的消融研究（表 7）。首先，在擴充之前直接訓練搜尋到的拓撲。該研究搜尋的種子拓撲優于圖 2 中 87 個随機拓撲中的最佳拓撲。
第二，該研究将基于複雜度的規則與「随機擴充 + As-ViT 拓撲」進行比較。在不同的擴充下，該研究的自動擴充也優于随機擴充。
COCO 資料集上的目标檢測
該研究将 As-ViT 與标準 CNN 和之前的 Transformer 網絡進行了比較。比較是通過僅更改主幹而其他設定未更改來進行的。從下表 8 的結果可以看出，As-ViT 也可以捕獲多尺度特征并實作最先進的檢測性能，盡管它是在 ImageNet 上設計的，并且它的複雜性是為分類而測量的。