這篇文章是 2022 ECCV 的一篇文章,是曠視科技的一篇文章,針對圖像恢復任務各種網絡結構進行了梳理,最后總結出一種非常簡單卻高效的網絡結構,這個網絡結構甚至不需要非線性激活函數。
文章一開始就提到,雖然在圖像復原領域,可以看到非常明顯的進步和提升,但是這些算法框架的復雜性也隨之增加。這篇文章重新思考了算法框架的設計,提出了一種非常簡單的基線算法框架,這個基線算法框架不僅效果上可以超越當前的一些 SOTA 方法,而且運算效率也非常高。文章為了進一步簡化這個基線模型,甚至將其中的非線性激活函數,比如 Sigmoid,ReLU,GELU,Softmax 等都可以等效替換成簡單的線性相乘,最終文章提出了一個非常簡單高效的基線算法框架,稱為 NAFNet,全稱叫 Nonlinear Activation Free Network,這個算法框架在多個圖像復原任務上都以極低的運算代價取得與 SOTA 相當甚至超越 SOTA 的效果。如圖 1 所示:
- 圖1
圖像恢復任務就是將一張退化的圖像(比如模糊,有噪)恢復成清晰圖像,目前這類圖像恢復任務基本都是基于深度學習的方法在做了,這類方法也取得了越來越好的復原效果,指標也越刷越高,絕大多數的深度學習方法,都是基于一個經典的架構,UNet,將多個 blocks 疊加成一個 U 型,中間還會帶著 skip 的連接,在這個經典架構的基礎上,衍生出很多變體,這些變體帶來了處理效果的提升,也導致了模型的復雜度隨之增加,這篇文章將模型的復雜度分為兩種,一種是 block 之間的復雜度,一種是 block 內的復雜度。
-
Inter-block Complexity,一般都是多階段的網絡結構,后面階段的網絡結構會對前面階段的輸出進行 refine,每個階段都是一個 U 型的結構,這種多階段的設計,是想把一個困難的任務,拆分成幾個簡單的子任務。
-
Intra-block Complexity,有很多工作會聚焦在 block 之間的設計,文章作者也列舉了一些例子。比如有的工作會將 attention 機制進行簡化,用 channel attention 代替常規的 spatial attention。此外,gated liner units 和 depthwise 卷積也比較常見。還有的工作是設計了基于窗口的多頭注意力機制。不過本文作者想說明的是,提升系統復雜度不是唯一的提升性能的方法,有的時候一些簡單的設計也能取得 SOTA 的表現。
Build A Simple Baseline
接下來,文章開始介紹,如何構建一個足夠簡單的基線結構,文章作者構建的原則就是奧卡姆剃刀原則:如無必要,勿增實體,不過文章也說了,必要性的判斷依據是實驗經驗。文章將模型的運算量限制在 16GMACs,對應的輸入圖像大小為 256 x 256。然后主要評估的圖像復原任務是降噪和去模糊。文章中用到的網絡結構,是最簡單的 U 型結構,如圖 2C 所示
- 圖 2
神經網絡一般都是由 block 疊加而成的,文章用最簡單的 residual block 構成基礎模型 PlainNet,如圖 3b 所示;
- 圖 3
文章接下來對 PlainNet 中的主要模塊進行了討論,包括 Normalization, Activation 以及 Attention 機制。
Normalization
現在已經有好幾種不同的 Normalization 的方法,包括 batch Normalization, instance Normalization, layer Normalization 等,不同的任務,不同的網絡結構,這些 Normalization 的方法發揮的作用可能也不一樣。所以文章作者經過比較,最后采用了 layer Normalization 這種方法。
Activation
對于激活函數,relu 是大家經常用的一種激活函數,不過,現在也逐漸發展出來了一些其它的激活函數類型,文章提到了 GELU,實驗證明,用 GELU 比用 gelu 的效果要好:
Attention
隨著 transformer 結構在 CV 領域的逐漸盛行,attention 機制也變得越來越常用,文章中也用到了 attention 機制,不過與原始的 transformer 中的 attention 機制不同,文章中用到的是 channel attention,而且是經過進一步簡化的 channel attention,具體結構如圖 4a 所示:
- 圖 4
文章后面還會對 4a 的結構做進一步的簡化。
Nonlinear Activation Free Network
上面的 PlainNet 運算復雜度相比之前的一些結構,已經降低了很多,但是依然可以取得 SOTA 的效果,不過文章作者并沒有滿足,他們認為這個網絡結構可以進一步的簡化,甚至可以去掉非線性激活函數,首先,文章作者分析了一些方法之后,發現 Gated Linear Units (GLU) 這個激活函數被經常用到, 所以對這個函數的性質做進一步的分析。
Gated Linear Units 的形式如下所示:
G a t e ( X , f , g , σ ) = f ( X ) ⊙ σ ( g ( X ) ) (1) Gate(\mathbf{X}, f, g, \sigma) = f(\mathbf{X}) \odot \sigma(g(\mathbf{X})) \tag{1} Gate(X,f,g,σ)=f(X)⊙σ(g(X))(1)
其中 X X X 表示 feature map, f , g f, g f,g 表示線性變換, σ \sigma σ 表示一個非線性激活函數,比如 Sigmoid, ⊙ \odot ⊙ 表示元素乘積。如果把 GLU 放入網絡結構,效果可以進一步提升,但是運算復雜度也會增加,這不是文章作者想看到的,所以他們對基線網絡結構中的 activation GELU 進行了改進:
G E L U ( x ) = x ? ( x ) (2) GELU(x) = x \phi(x) \tag{2} GELU(x)=x?(x)(2)
其中 ? \phi ? 表示標準正態分布的累積分布函數,一般 GELU 可以估計成如下形式:
gelu ( x ) = 0.5 x ? ( 1 + tanh ? ( 2 / π ? ( x + 0.044715 ? x 3 ) ) ) (3) \text{gelu}(x) = 0.5 x \ast (1 + \tanh(\sqrt{2/\pi} \ast (x + 0.044715 \ast x^{3} ))) \tag{3} gelu(x)=0.5x?(1+tanh(2/π??(x+0.044715?x3)))(3)
從公式 (1) 和公式 (2),可以注意到,GELU 是 GLU 一種特殊情況,如果 f , g f, g f,g 取恒等變換, σ \sigma σ 取 ? \phi ?,那 GLU 就變成 GELU 了,從這種相似性,文章作者猜測 GLU 可以看成是一類廣義的激活函數,或許可以替換其它的激活函數。進一步,文章作者認為 GLU 本身就包含非線性,即使去掉 σ \sigma σ 函數, G a t e ( X ) = f ( X ) ⊙ g ( X ) Gate(\mathbf{X}) = f(\mathbf{X}) \odot g(\mathbf{X}) Gate(X)=f(X)⊙g(X) 也包含了非線性,基于這些推論,文章作者提出了一種 GLU 的簡化版,直接將 feature map 在 channel 層面平分成兩部分,然后再直接相乘,如圖4 c 所示,這個簡化后的 GLU 稱為 SimpleGate,與 GELU 相比,SimpleGate 只需要簡單的乘法就可以實現:
S i m p l e G a t e ( X , Y ) = X ⊙ Y (4) SimpleGate(\mathbf{X}, \mathbf{Y}) = \mathbf{X} \odot \mathbf{Y} \tag{4} SimpleGate(X,Y)=X⊙Y(4)
模塊的具體結構如圖 4c 所示,通過這個替換,網絡的表現進一步有提升,到目前為止,網絡中還剩下的非線性函數就是 channel attention 中的 Sigmoid 以及 ReLU 了。
在前面的 PlainNet 中,采用了 channel attention 的機制,如圖 4a 所示,這種機制同樣可以捕獲到全局信息,同時保持運算的高效,它將空間信息在channel 維度先進行壓縮,然后利用一個 MLP 去計算 channel attention,作為每個 feature map 的權重,再與之前的 feature map 相乘:
C A ( X ) = X ? σ ( W 2 max ? ( 0 , W 1 p o o l ( X ) ) ) (5) CA(\mathbf{X}) = \mathbf{X} \ast \sigma(W_{2} \max (0, W_{1} pool(\mathbf{X}))) \tag{5} CA(X)=X?σ(W2?max(0,W1?pool(X)))(5)
其中, X \mathbf{X} X 表示 feature map,pool 表示全局平均池化操作, σ \sigma σ 表示 Sigmoid 函數, W 1 , W 2 W_1, W_2 W1?,W2? 表示全連接層,如果將 channel attention 的計算看成是一個函數,那么公式 (5) 可以寫成:
C A ( X ) = X ? Φ ( X ) (6) CA(\mathbf{X}) = \mathbf{X} \ast \Phi(\mathbf{X}) \tag{6} CA(X)=X?Φ(X)(6)
可以看到,公式 6 與公式 1 很像,受此啟發,文章作者也推測 channel attention 可能也是 GLU 的一種特殊形式,或許也可以進行類似的簡化,文章將 channel attention 最終簡化成如下所示:
S C A ( X ) = X ? W p o o l ( X ) (7) SCA(\mathbf{X}) = \mathbf{X} \ast W pool(\mathbf{X}) \tag{7} SCA(X)=X?Wpool(X)(7)
模塊的具體結構如圖 4b 所示,相比 4a 的結構,4b 去掉了 Sigmoid 以及 ReLU,全連接也減少了一層
基于 PlainNet, 文章將其中的 GELU 以及 channel attention 做了簡化,最終得到了一個不含非線性激活函數的輕量級網絡,稱為 NAFNet。
文章最后給出了很多實驗對比,說明 NAFNet 在保持低運算代價的同時,還能保持很好的效果指標,真有點大道至簡,大巧若拙的感覺了。