NAS新視角:圖神經網路驅動的通用架構空間,卷積與Transformer混合架構性能躍升!

圖片

作者提出通用神經網路架構空間(UniNAS),一個用於神經網路架構搜尋(NAS)的通用搜尋空間,它在一個單一、靈活的框架下統一了卷積網路、Transformer及其混合架構。UniNAS能夠在通用框架下發現新穎的架構以及分析現有架構。

作者還提出了一種新的搜尋演算法來遍歷所提出的搜尋空間,並證明了該空間包含有趣的架構,在採用相同的訓練設定時,這些架構的性能優於最先進的手工設計架構。

最後,作者引入了一個統一的工具包,其中包含標準化的訓練和評估協議,以促進可重現性並能在NAS研究中實現公平比較。總而言之,這項工作為透過一個統一的、基於圖的NAS視角來系統地探索神經網路架構的全譜系開闢了一條道路。

1. 引言

儘管神經網路架構搜尋(NAS)在為預定義架構[37, 51]識別最佳超參數配置或在改善邊緣設備[15, 36]上的推論延遲方面取得了不可否認的成功,但據作者所知,它至今未能產生一種能夠顯著優於基於ResNet[13, 34, 35]或Vision Transformer[7, 10, 39]的流行但手工設計網路架構的新型網路架構。

與電腦視覺研究的許多領域一樣,NAS的進展主要由可用的基準測試驅動,因此近年來,大多數NAS方法[11, 17, 31, 42]都集中於NAS-Bench系列[9, 46]的表化NAS基準測試。表化基準測試是一個固定的網路架構資料集,其中每個網路的精度和其他參數都已被預先計算,因此NAS演算法無需訓練網路即可獲得其精度。儘管表化資料集透過降低所需的計算成本來促進NAS研究,但根據定義,它們也阻止了NAS方法發現新的、更好的架構——搜尋空間中的最佳架構已是已知的,並且只需查閱預先計算的精度表就能找到它。此外,這些表化資料集的搜尋空間通常相對較小,並且局限於多次重複相同的建構單元,這是因為在建立基準測試期間,需要對每個網路至少訓練一次。

得益於免訓練NAS方法的出現[5, 19, 22, 23, 40, 48],現在可以在不訓練網路架構的情況下估計其準確性,從而消除了NAS的計算約束,並能夠探索更大的搜尋空間和新型架構。為此目的引入的搜尋空間——即Zen-NAS [23]和AutoFormer [3]——是透過改變所選手工設計架構塊的超參數(如深度、寬度或擴展比率)創建的,這些架構塊分別是MobileNetV2 [33]和Vision Transformer [10]。由於每個搜尋空間都是透過改變給定架構塊的超參數創建的,因此使用這些搜尋空間仍然無法找到具有不同拓撲結構或建構塊的網路,隨後,手工設計的架構(如CoAtNet [7])仍然優於AutoFormer搜尋空間中已知的最佳架構[2]。

在本文中,作者旨在透過引入一種名為通用神經網路架構空間(UniNAS)的新穎搜尋空間來解決上述局限性。該搜尋空間以通用方式設計,不偏向於任何特定的現有架構,同時確保所有最先進但手工設計的架構都包含在該搜尋空間內。因此,該搜尋空間不僅允許探索現有架構的各種組合和全新的架構,還允許在一個統一的框架內對最先進的網路、其拓撲結構和設計選擇進行系統性分析。

此外,作者提出了一種新的架構搜尋演算法來遍歷所提出的UniNAS搜尋空間。當與最先進的免訓練NAS Agent [40]結合時,在遍歷搜尋空間的少量步驟後,作者發現了一種新穎的網路架構——UniNAS-A,該架構在使用相同的訓練協議和大小約束下,優於當前最先進的架構,這表明該搜尋空間包含有趣的架構,值得進一步探索。

最後但同樣重要的是,作者提供了一個工具包(作為一個易於安裝的Python套件),該工具包允許遍歷所提出的搜尋空間,為搜尋空間中的任何點創建PyTorch網路模組,最重要的是提供了一致且明確的訓練協議和訓練程式碼來創建和評估最終模型。這對於可重現性以及NAS方法的公平比較極其重要,因為儘管作者通常會報告在標準資料集(如ImageNet)上的最終準確度,但他們使用不同的訓練計畫、不同的超參數設定和資料增強技術,或者他們甚至會使用一個在顯著更多資料上預訓練的現有更大模型作為教師網路——這使得所提出架構之間的公平比較幾乎不可能。

綜上所述,作者對社群做出了如下貢獻:

1. 作者提出了一種名為UniNAS的新型通用NAS搜尋空間。該搜尋空間以通用方式設計,因此包含了許多新穎的網路拓撲結構,與此同時,它也包含了已知的人工設計架構,從而允許在統一框架內對網路拓撲結構進行系統分析。

2. 作者引入了一種新的搜尋演算法來遍歷所提出的設計空間,並表明該空間包含有趣的新穎架構——例如UniNAS-A——它在多種任務(分類、檢測和分割)上優於手工設計的最先進架構。

3. 作者提供了一個統一的工具包,包括標準化的訓練和評估協議,以幫助採用所提出的通用搜尋空間,增強NAS方法的可重現性,並促進未來的NAS研究。

2. UNINAS

本節介紹了UniNAS搜尋空間,首先介紹作為其核心組件的UniNAS塊,然後是整體網路架構,最後透過與現有搜尋空間的比較,突出關鍵差異和改進。

A. UniNAS塊

a) 基本操作:作者的目標是創建一個跨越多種不同架構的搜尋空間,因此作者在UniNAS塊的定義中採用通用設定,在架構搜尋過程中實現最大靈活性,同時保留與現代分層網路兼容的實際約束。具體而言,UniNAS塊被定義為具有一個輸入節點和一個輸出節點的任意有向無環圖(DAG),其中每個中間節點代表1個基本操作。基本操作包括卷積層(深度可分離、逐點、標準等)、池化、Mask、正規化元素、各種形式的矩陣乘法和點積、具有多輸入或多輸出邊的操作以及非線性函數(見表1)。這使得該塊能夠表示多樣化的局部計算模式,同時保持一致的接口屬性。

雖然表1中的大多數基本操作確實很基礎,因此無需進一步解釋,但作者將具體介紹其中兩個較為複雜的操作:對於張量 ,作者定義

對於 和 以及對於 ,

對於 , 和 。Matmul1和Matmul2表示一種乘法關係,但與逐元素乘法不同,它結合了所有通道資訊,並與Softmax結合,這些可以組合產生注意力機制。然而,作者再次注意到,可變性遠不止於此,因為任何圖都是可能的,只有維度才是重要的(!)。

表1: UniNAS 塊基本操作 作者報告了輸入大小為 的單次前向傳播的參數數量 (Params) 和浮點運算次數 (FLOPs)。某些節點會改變形狀和/或在多輸入或多輸出模式下運行, 表示 個相同形狀的張量。

b) 塊計算圖:作者對這個通用公式施加的唯一約束是:1) 塊的輸入和輸出維度保持相同,以及 2) DAG 內相鄰節點之間的維度匹配,以確保張量在圖中正確傳播。這一約束簡化了在更大架構中的整合,其中跨塊的一致特徵圖形狀顯著減少了前向傳播過程中動態形狀處理的複雜性,並能夠在各種搜尋到的塊上實現穩定訓練。透過強制執行一致的維度,作者避免了自動插入額外投影的需要,否則這些投影可能會干擾對搜尋到的拓撲結構的分析。

作者強調,現代深度學習架構中的許多常用模組,例如ResNet [13]及其變體的殘差塊、基於Transformer的模型[7, 10]中帶或不帶相對位置偏置的自注意力層、用於自適應通道重新校準的squeeze-and-excitation模組[16],以及EfficientNets [34, 35]中使用的inverted mobile bottleneck結構,都可以表示為作者的UniNAS塊的實例。關於如何將這些模組製作成UniNAS塊格式的示例,包括精確的圖結構和節點操作,可以在圖3中看到。

圖片

然而,作者鼓勵讀者不要局限於這些經典例子;人們當然可以想像出更多樣化的網路。搜尋空間不僅包括簡單地在鏈中堆疊卷積和非線性,還涵蓋了樹狀結構、具有選擇性注意力合併的並行路徑,以及卷積和自注意力層的混合組合,所有這些都封裝為一個UniNAS塊。這展示了塊設計的表達能力,同時在整個搜尋空間中保持統一和連貫的表示。

B. UniNAS網路

最終的網路結構將不同的UniNAS塊順序堆疊成單鏈(見圖2)。這遵循了當前最先進網路的一般設計,其分層骨幹網路(Backbone)類似於[7, 13, 35, 39]。莖階段(S0)使用卷積層對輸入進行降取樣(downsampling),其後是多個階段,每個階段包含多個不同的UniNAS塊,最後是一個由全局平均池化和全連接層組成的分類頭。由於UniNAS塊保持維度不變,在每個階段的開始,空間維度透過步長為2的標準最大池化減小,而通道數透過通道投影1增加。如前所述,UniNAS塊的結構(即圖表示)在整個網路中是變化的,這增強了網路的拓撲可變性。階段數、塊數以及空間和通道維度作為縮放超參數,允許探索不同的網路規模模式。

圖片

表2:NAS搜尋空間比較

UniNAS涵蓋了最多种類的網路,它不限制塊拓撲結構,並允許探索新穎的架構。作者還報告了每個空間中已知最佳網路在ImageNet-1k上的分類準確度——在未找到可比訓練設定下的報告結果時,作者使用與表3相同的訓練協議訓練給定搜尋空間中已知最佳架構(用†表示)。

圖片

C. 與現有搜尋空間的比較

現有的搜尋空間存在兩個主要限制:受限的拓撲可變性和較差的可擴展性。因此,儘管該領域已有廣泛的前期研究,但在這些空間中發現的性能最佳的模型,與使用 UniNAS 獲得的架構相比,其性能仍然不及(見表2)。

諸如[26, 49]中的DARTS方法依賴於權重共享超網路,這些網路計算成本高昂且產生有偏的梯度估計,導致不可靠的架構排名。當擴展到Transformer [4]時,自注意力層的成本迫使這些搜尋空間受到嚴重限制——通常將搜尋簡化為「使用注意力或不使用注意力」的二元選擇。即使有這些限制,基於DARTS的網路仍然局限於小規模設定,ImageNet-1k準確度低於(表2)。

諸如NAS-Bench [9, 46]等基準測試面臨一個根本的可擴展性問題:可能網路的數量隨操作數量呈指數級增長,使得窮舉訓練僅在玩具空間上可行。因此,這些基準測試僅限於在CIFAR [18]或ImageNet16-120 [6]等小型資料集上評估的簡化卷積網路。在ImageNet-1k上,要麼不存在結果,要麼報告的性能遠低於最先進水平,由於架構空間是有限的且已被完全探索,因此沒有進一步改進的空間。

Zen-NAS [23] 和 AutoFormer (V1/V2,也稱為 S3) [3, 4] 將搜尋空間限制在具有重複 MobileNetV2 塊 [33] 或 Vision Transformer 塊 [10, 27] 的網路上。生成的架構僅在超參數上有所不同,如擴展比例、通道數量或頭數。然而,目前無法公平比較這些空間中報告的頂級模型,因為它們是透過從一個更大的教師模型進行知識蒸餾來訓練的,並且使用了 substantially 更多資料 [19],因此尚不清楚報告的性能提升是由於搜尋本身還是蒸餾過程。事實上,當應用蒸餾時,[12] 使用標準的 EfficientNet-B2 [34] 在相同的 FLOPs 和訓練預算下實現了更高的準確度。類似地,當在相同的參數預算下訓練時,[43] 中的手工設計的Transformer在 ImageNet-1k 上的性能超過了 AutoFormer 的最佳 Transformer [2, 3] 超過 。

總結來說,有兩個主要的改進領域:網路多樣性和公平可重現的比較。這兩個問題都由UniNAS解決。1) UniNAS的靈活性允許進行真正拓撲感知的架構搜尋,其中注意力機制可以插入到選定位置,與卷積結合,或在某些分支中完全替換,這種靈活性在之前受限於單一操作選擇的框架中是不可行的;2) UniNAS允許對網路拓撲進行公平比較,因為它涵蓋了所有提到的空間以及諸如ResNet [13]、EfficientNet [34]、ViT [10]和CoAtNet [7](見圖3)等拓撲結構,因此它能夠在一個框架內直接比較它們的特性和準確性,以及相互之間和與新型架構的比較。

與之前的搜尋空間不同,在作者的UniNAS中,在每個階段使用不同的塊是可行的(參見[28]對NAS-Bench框架的批評,該批評源於以下事實:簡單的模型參數數量對於估計最終網路精度具有最高的預測強度,因此阻礙了研究人員探索新的有趣拓撲結構)。

3. 架構搜尋

在本節中,作者提出了一種基於給定標準在UniNAS搜尋空間中查找網路的演算法。具體而言,作者假設設計選擇如階段數、塊數和基礎通道維度是已知的,並且以Params和FLOPs邊界形式的約束是給定的。然後作者的目標是在上述約束條件下找到最大化特定目標(如最高準確度)的網路。

a) 搜尋步驟:更正式地說,任何UniNAS網路由一系列圖 標識,每個圖對應一個UniNAS塊,其中 是總塊數。這使作者能夠將UniNAS搜尋表述為一種基於圖的演算法,涉及節點的添加和消除。為了保證搜尋能夠高效地遍歷1)僅在網路大小約束範圍內的網路,並且

1. 可行圖(就節點維度與奇偶性而言),作者採用以下方法。

作者將每個節點 (其中 )與可訓練參數數量Params和近似浮點運算數量FLOPs相關聯。這些值可以很容易地作為每個基本操作節點的輸入張量形狀的函數獲得(可以在表1中找到),這使作者能夠在添加/刪除給定節點後高效計算整體網路成本。

作者注意到,每個基本操作的Params值計算起來非常直接。然而,估計FLOPs並非易事,作者選擇將每個操作的所有乘法和加法都包含在其最終的FLOPs值中。這樣做是為了給所有操作賦予非零成本,以防止出現不可控的複雜度發散。因此,這個值可能與執行時FLOPs估算器(如PyTorch profiler)返回的值不同,而且,該值還可能因底層硬體的不同而有所差異。

1. 作者為搜尋空間定義了可行的節點添加和消除操作。這些操作將確保UniNAS塊在一次搜尋步驟後仍然是一個可行的計算圖,同時允許作者在大型UniNAS空間中自由探索。請記住,圖的可行性是由表1中的輸入/輸出節點形狀決定的。只需確保:a) 僅當輸入空間維度具有整數平方根時才添加RelPosBias,b) 僅當通道維度可被2或3整除時才分別添加Chunk2、Chunk3,c) 僅當通道維度可被4整除時才添加ConvRed4,d) 改變維度及/或具有多個輸出的節點與其耦合節點一起添加。當消除時,這兩個節點之間的分支也會被消除。

演算法1 UniNAS搜尋步驟

圖片

要求:UniNAS網路,即一系列圖 ,可能的節點類型列表(參見表1,包含節點成本FLOPs,Params),搜尋邊界 , , Paramsmax, Params ),消除機率 peliminate ,最大嘗試次數 。 1: 當 時執行

2: 隨機選擇一個 graph 和一個 node

3:採樣

4: 如果 則

5: {節點 消除}

6: 確定在 v 和它的耦合節點之間的最小子圖 (對於只有一個輸出且不改變維度的節點則為空)

7: 計算在潛在消除 後 FLOPs 和 Params 的變化

8: 如果變化相對於 和 Params 是可信的,那麼

9: 消除

10: 休息

11: end if

12: 否則

13: {在 之後的節點添加}

14: 計算在潛在添加 及其耦合節點 之後 FLOPs 和 Params 的變化(對於不改變維度的單輸出節點,此項為空)

15: 如果變化相對於 和 Params 是合理的,則

17: 中斷

18: end if

19: 結束如果

20:

21: end while最後,作者建構了(單個)搜尋步驟。在每個步驟中,作者選擇一個將被調整的UniNAS塊,決定是添加還是消除一個節點。如果作者添加,則在隨機選擇的節點後添加,前提是這根據上述1)和2)是可行的。確定添加/消除選項的參數 通常設定為低於0.5,因為作者希望先添加然後消除節點,因為對於某些類型的節點,消除它們也意味著消除該節點與其配對對應節點之間的整個分支,這可能導致網路大小的大幅削減。在實驗中,作者選擇 的值,如圖1所示,透過這種選擇,隨機遊走避免了網路大小向無限制增長或縮小的退化。更正式的描述請參見演算法1。

圖片

b) 訓練與評估協議:當免訓練NAS演算法完成時,其識別為最有希望的候選架構仍需進行訓練以獲得最終(真實)準確度。遺憾的是,NAS文獻中一直缺乏精確的訓練方案,不同作者使用不同的訓練資料、Epoch數量、訓練批次大小等來報告結果,這使得不同NAS方法的直接比較變得不可能。在UniNAS搜尋空間中,因此作者也提供了訓練最終網路的詳細訓練協議(見表3),以便不同架構和NAS方法(包括未來的工作)可以以公平且可重現的方式進行比較。

c) UniNAS工具包:作者提供了一個可透過pip安裝的軟體套件uninas,使研究人員能夠輕鬆訪問使用所提出的UniNAS空間所需的所有組件。具體而言,作者提供了將任何UniNAS網路生成為PyTorch模型的工具,以及建構自定義UniNAS網路的直觀界面。此外,作者還包括一個用於計算圖形視覺化的模組,該模組專門為複雜的網路結構設計。作者還提供了演算法1中描述的搜尋演算法的實現,該實現考慮了計算預算約束,封裝在單個函數調用中,該函數調用可以觸發在UniNAS空間中的隨機遊走或針對可從PyTorch網路計算的任何目標的優化演算法。最後,作者還包括了表3中描述的訓練協議,並可選支援分散式訓練,以促進任何未來UniNAS網路的可重現性和公平比較。

4. 結果

a) UniNAS中的隨機遊走:作者使用作者的搜尋步驟演算法1和參數con執行了隨機遊走

表3:UniNAS評估協議

作者透過首先在ImageNet-1k上訓練圖像分類,然後使用上述超參數(hyper-parameters)在COCO和ADE20K資料集上進行微調(fine-tuning),從而在UniNAS空間中對分類、檢測和分割任務評估NAS方法。批次大小(batch size)的選擇以適合單個A100 GPU為準,並且是按每個GPU報告的,因此當使用更多GPU時,需要相應地調整學習率(learning rate)。

將參數量限制在22-28M,FLOPs限制在6-20G,以評估作者的架構搜尋在巨大的UniNAS空間中的導航能力。在圖1中,作者可以看到50萬個採樣網路的Params和FLOPs,表明作者的搜尋可以輕鬆跨越不同的網路大小和配置,還涵蓋了基於Transformer和卷積的非常不同的網路架構。在圖4中,作者提供了UniNAS基本操作的進一步細分,顯示網路大小在特定操作之間高度可變,表明作者的探索有效地覆蓋了非常多樣化的配置。

圖片

b) 架構搜尋:在作者的UniNAS中搜尋最佳架構時,作者使用免訓練的NAS Agent VKDNW[40]來評估網路性能,而不是對其進行訓練。作者定義

其中, 表示在隨機輸入批次下初始化時,Empirical Fisher Information Matrix (FIM) 特徵值的第k個十分位數,作為FIM譜的一種表示。具體而言,作者分別為每個塊 計算公式(3),然後對所有這些塊的結果取平均,由此得到一個單一純量值,作者在搜尋中對該值進行最大化。作者之所以選擇這個免訓練 Agent,是因為它與網路規模正交(見[40]中的圖3),而作者的目標是在特定的網路規模預算內搜尋最佳拓撲結構。

然後作者在UniNAS中疊代搜尋:作者從初始網路開始,並執行演算法1中的1024步,

表4:UniNAS搜尋空間中的模型在ImageNet-1k上的分類準確度。所有模型具有相似數量的參數,並採用表3中的訓練協議以相同方式進行訓練。訓練在8塊A100 GPU上耗時2天。

圖片

其中只保留前64個網路(族群大小)。在搜尋過程中,網路規模被限制為27M參數和20G FLOPs,在圖2中作者採用了4個階段,初始輸出大小為64,各階段分別有2、3、5和2個UniNAS塊,隱藏維度分別為96、192、384、768,與現代架構3一致。搜尋在單個A100 GPU上耗時12小時,最終作者根據VKDNW分數選擇了最佳網路——作者將其表示為UniNAS-A(見圖5)。

圖片

c) 圖像分類:首先,作者將UniNAS-A與UniNAS搜尋空間中的其他網路進行比較:EfficientNet [34]、ResNet [13]、CoAtNet和具有相對位置偏置的ViT [7]。為了在公平的設定下比較這些網路,作者將網路擴展到相同數量的階段和塊,並透過調整通道數量,作者也使網路規模保持相同。每個網路都在ImageNet1k [8]上進行訓練,使用與表3相同的訓練協議,訓練在8個A100 GPU上花費了2天時間。如表4所示,UniNAS-A以顯著優勢優於標準手工設計的網路。

d) 下游任務

接下來,作者對上一步的每個網路在兩個下游任務上進行了微調——在 MS-COCO [24] 上的目標檢測和在 ADE20K [50] 上的語義分割,使用與表3 相同的設定。在4個 A100 GPU 上,分割任務的訓練耗時約10小時,檢測任務的訓練耗時21小時4。再一次,UniNAS-A 的表現顯著優於現有網路(見表5)。

e) Agent 消融

最後,作者透過使用不同的 Agent 運行搜尋演算法1,並評估每個 Agent 所找到的最佳架構,來對選擇 VKDNW [40] 作為 Agent 進行消融。在表6中,作者表明,對UniNAS空間的搜尋對特定 Agent 的選擇基本保持穩健性,但VKDNW在相同的搜尋預算(12小時)下仍然能夠找到最佳網路架構。

圖片

5. 相關工作

由於NAS搜尋空間的討論已在第II-C節中呈現,在此作者專注於NAS搜尋方法。

a) One-shot NAS:這些方法基於離散架構空間的鬆弛,通常使用一個超網路,該超網路包含了給定搜尋空間中所有可能的節點和邊。該方法首次在DARTS [26]中被提出,其中超網路由所有可能的操作組成,每個操作都被分配一個權重,該權重在超網路訓練期間透過梯度下降進行調整。一旦超網路訓練完成,就保留權重最高的操作,從而建構出最終的網路架構。Robust-DARTS [47]透過在超網路訓練中引入資料增強來改善測試時的泛化能力,而SGAS [21]則旨在透過在每個階段顯式地選擇允許的操作來提高超網路訓練的穩定性。One-shot NAS的主要挑戰是記憶體消耗,因為超網路必須包含整個搜尋空間中所有可能的操作;以及排名紊亂,即在超網路中評估的架構的性能可能與其作為獨立網路時的性能不同。

6. 結論

作者提出了UniNAS,一個通用的神經網路架構搜尋空間,旨在一個統一的框架內系統地探索、分析和比較網路拓撲。與以往的工作相比,作者進一步將計算模組分解為基本操作,從而能夠表達和擴展手工設計與NAS發現的架構,同時支援對網路設計中拓撲可變性的系統性研究。

作者提出了一種高效的架構搜尋演算法,該演算法直接在UniNAS內運行,實現了對FLOPs和參數預算的精確控制,同時支援細粒度修改以遍歷多樣化的架構家族,

參考

[1]. Universal Neural Architecture Space: Covering ConvNets, Transformers and Everything in Between

主標籤:神經網路架構搜尋

次標籤:圖神經網路深度學習Transformer卷積神經網路


上一篇:您說得對,AGI一年內不會出現!學術界對通用人工智慧的最新定義,由27家機構聯合發佈

下一篇:GPT 越來越保守?史丹佛 Manning 團隊提出 Verbalized Sampling,讓模型重新「多想一點」

分享短網址