這是 2022 CVPR 上的一篇文章，介紹了用 MLP 做 low-level 圖像處理的工作

Abstract

近年來，Transformer 和多層感知機（MLP）模型的發展為計算機視覺任務提供了新的網絡架構設計。盡管這些模型在圖像識別等許多視覺任務中已被證明是有效的，但在將它們應用于底層視覺任務時仍存在挑戰。對高分辨率圖像支持的靈活性不足以及局部注意力機制的局限性，可能是主要的瓶頸所在。在這項工作中，我們提出了一種基于多軸 MLP 的架構，稱為 MAXIM，它可以作為一種高效且靈活的通用視覺主干架構，用于圖像處理任務。MAXIM 采用了 U 型網絡的分層結構，并通過空間門控的 MLP 實現了長距離的交互。具體來說，MAXIM 包含兩個基于 MLP 的構建模塊：一個多軸門控 MLP，它能夠對局部和全局視覺線索進行高效且可擴展的空間混合；以及一個交叉門控模塊，作為交叉注意力機制的替代方案，用于處理跨特征條件。這兩個模塊都完全基于 MLP，但同時也受益于其全局特性和 “全卷積” 特性，這兩個特性在圖像處理中是非常理想的。我們廣泛的實驗結果表明，所提出的 MAXIM 模型在十多個基準測試中，涵蓋一系列圖像處理任務，包括去噪、去模糊、去雨、去霧和增強等，均取得了當前最優的性能表現。而且，與其他具有競爭力的模型相比，它所需的參數數量和浮點運算次數（FLOPs）更少或相當。

Introduction

圖像處理任務，如圖像復原和增強，是重要的計算機視覺問題，其目標是從退化的輸入圖像中生成理想的輸出圖像。不同類型的圖像退化可能需要不同的圖像增強處理方法，例如去噪、去模糊、超分辨率、去霧、低光照增強等等。

鑒于經過精心整理的大規模訓練數據集日益增多，近期基于精心設計的卷積神經網絡（CNN）的高性能方法 [15, 17, 20, 22, 50, 52, 60, 61, 110, 111, 125] 在許多任務上展現出了當前最優（SOTA）的性能。

改進基礎模型的架構設計是提升包括圖像復原在內的大多數計算機視覺任務性能的關鍵因素之一。眾多研究人員發明或借鑒了各種獨立的模塊或構建單元，并將它們應用于底層視覺任務中，這些應用涉及到殘差學習 [43, 95, 120]、密集連接 [95, 121]、分層結構 [37, 41, 42]、多階段框架 [16, 34, 111, 113] 以及注意力機制 [66, 91, 110, 111]。

最近對視覺 Transformer（ViT）[11, 24, 57] 的研究探索已經證明了它們作為常用卷積神經網絡（CNN）模型替代方案的巨大潛力。ViT [24] 的出色表現也激發了類似的模型設計，這些設計采用了更簡單的全局算子，比如 MLP-Mixer [87]、gMLP [54]、GFNet [76] 和 FNet [44] 等等。盡管這些模型在許多高級任務中取得了成功應用 [4, 24, 57, 85, 89, 102, 104]，但它們在底層圖像增強和復原問題上的有效性尚未得到廣泛研究。在將 Transformer 應用于底層視覺領域的開創性工作 [10, 15] 中，直接采用了全自注意力機制，而這種機制只接受相對較小且固定尺寸的圖像塊（例如 48×48）。當使用裁剪的方法將其應用于較大尺寸的圖像時 [15]，這種策略不可避免地會導致圖像塊邊界偽影的出現。基于局部注意力的 Transformer [52, 97] 改善了這一問題，但它們也受到限制，其感受野大小有限，或者會失去非局部性 [24, 93]，而非局部性是 Transformer 和 MLP 模型相對于分層 CNN 而言極具吸引力的一個特性。

為了克服這些問題，我們提出了一種通用的圖像處理網絡，名為 MAXIM，用于底層視覺任務。MAXIM 的一個關鍵設計要素是采用了多軸方法（第 3.2 節），該方法可以并行捕捉局部和全局的交互信息。通過在每個分支的單軸上混合信息，這種基于多層感知機（MLP）的算子變成了 “全卷積” 形式，并且在圖像尺寸方面呈線性擴展，這顯著提高了它在密集型圖像處理任務中的靈活性。

我們還定義并構建了一個純粹基于 MLP 的交叉門控模塊，該模塊使用相同的多軸方法，對 MAXIM 中間部分的跳躍連接進行自適應門控，從而進一步提升了性能。受近期圖像復原模型的啟發，我們開發了一種簡單而有效的多階段、多尺度架構，該架構由一系列 MAXIM 主干網絡堆疊而成。

MAXIM 在一系列圖像處理任務中都取得了出色的性能，同時所需的參數數量和浮點運算次數（FLOPs）極少。我們的貢獻在于：

• 一種新穎且通用的圖像處理架構，名為 MAXIM，采用了一系列編碼器 - 解碼器主干網絡，并由多尺度、多階段損失函數進行監督。
• 專為底層視覺任務量身定制的多軸門控 MLP 模塊，該模塊始終具備全局感受野，且相對于圖像大小具有線性復雜度。
• 一個交叉門控模塊，用于對兩個獨立的特征進行交叉條件處理，同樣具備全局特性且為全卷積結構。
• 大量實驗表明，MAXIM 在包括去噪、去模糊、去雨、去霧和圖像增強等在內的 10 多個數據集上取得了當前最優的結果。

MAXIM

我們提出了首個用于底層視覺任務的有效通用多層感知機（MLP）架構，我們將其稱為用于圖像處理的多軸多層感知機（MAXIM）。與之前的底層視覺 Transformer 模型 [10,15,52,97] 不同，MAXIM 具有幾個理想的特性，使其在圖像處理任務中頗具吸引力。首先，MAXIM 能夠以線性復雜度在任意大尺寸的圖像上實現全局感受野；其次，它可以直接支持任意的輸入分辨率，也就是說，它是全卷積的；最后，它提供了一種局部（卷積層）和全局（多層感知機）模塊的平衡設計，在無需大規模預訓練的情況下 [15]，性能超越了當前最優（SOTA）方法。

• 圖 2 MAXIM 架構。我們采用（a）一個編解碼器主干網絡，其中每個（b）編碼器、解碼器以及瓶頸層都包含一個多軸門控多層感知器（MLP）模塊（圖 3）以及一個殘差通道注意力模塊。該模型通過（c）一個交叉門控模塊得到進一步提升，這個交叉門控模塊使得全局上下文特征能夠對跳躍連接進行門控。更詳細的描述見附錄 A.2。

Main Backbone

MAXIM 主干網絡（圖 2a）遵循源自 U-Net [80] 的編碼器 - 解碼器設計原則。我們發現，像 3×3 卷積這樣作用范圍小的算子對類 U-Net 網絡的性能至關重要。因此，我們對每個模塊采用混合模型設計（圖 2b）—— 使用卷積層處理局部信息，使用多層感知機（MLP）實現長距離交互 —— 以便充分發揮它們各自的優勢。

為了在不同尺度上實現長距離的空間信息混合，我們將多軸門控多層感知機模塊（MAB）插入到每個編碼器、解碼器以及瓶頸層中（圖 2b），隨后堆疊一個殘差通道注意力模塊（RCAB）[100, 111]（LayerNormConv-LeakyReLU-Conv-SE [31]）。
受跳躍連接的門控濾波機制 [67, 71] 的啟發，我們對門控多層感知機（gMLP）進行擴展，構建了一個交叉門控模塊（CGB，圖 2c）。它是一種高效的二階模塊，可作為交叉注意力機制（三階相關性）的替代方案，用于實現兩種不同特征之間的交互或條件處理。我們利用來自瓶頸層（圖 2a）的全局特征對跳躍連接進行門控，同時將經過優化的全局特征向上傳播到下一個交叉門控模塊。多尺度特征融合 [20, 84, 110]（紅色和藍色線條）被用于在編碼器到交叉門控模塊（Encoder $\to$ CGB）以及交叉門控模塊到解碼器（CGB $\to$ Decoder）的數據流中聚合多層次的信息。

Multi-Axis Gated MLP

我們的工作受到了文獻 [123] 中提出的多軸分塊自注意力機制的啟發，該機制可在不止一個軸上執行注意力操作。在分塊圖像的兩個軸上執行的注意力操作對應于兩種形式的稀疏自注意力機制，即區域注意力和擴張注意力。盡管該模塊能夠并行捕捉局部和全局信息，但它無法適應圖像復原或增強任務，因為在這些任務中，測試圖像的尺寸往往是任意的。

我們對用于圖像處理任務的 “多軸” 概念進行了改進，構建了一個（分頭式）多軸門控多層感知機模塊（MAB），如圖 3 所示。與文獻 [123] 中在單層應用多軸注意力機制不同，我們首先將頭部分為兩半，每一半獨立進行劃分。在局部分支中，尺寸為 $(H,W,C/2)$ 的特征的半個頭被分塊成形狀為 $(\frac{H}{b}\times \frac{W}{b},b\times b,C/2)$ 的張量，這表示將其劃分為大小為 $(b\times b)$ 的非重疊窗口；在全局分支中，另一半頭部使用固定的 $(d\times d)$ 網格劃分為形狀為 $(d\times d,\frac{H}{d}\times \frac{W}{d},C/2)$ 的形式，每個窗口大小為 $(\frac{H}{d}\times \frac{W}{d})$。為便于可視化，在圖 3 中我們設 $b=2$，$d=2$。為使其成為全卷積結構，我們僅在每個分支的單個軸上應用門控多層感知機（gMLP）模塊 [54] —— 局部分支在第二個軸上應用，全局分支在第一個軸上應用 —— 同時在其他空間軸上共享參數。直觀來講，并行應用多軸 gMLP 分別對應著空間信息的局部和全局（擴張式）混合。最后，處理后的頭部被拼接起來并進行投影以減少通道數，再通過來自輸入的長跳躍連接進一步合并。值得注意的是，與處理固定大小圖像塊的方法 [15] 相比，這種方法為我們的模型帶來了優勢，避免了圖像塊邊界偽影的問題。

復雜度分析。我們所提出的多軸門控多層感知機模塊（MAB）的計算復雜度為。

$$\Omega (\text{MAB})=\underset{\text{Global?gMLP}}{\underbrace{d^{2}HWC}}+\underset{\text{Local?gMLP}}{\underbrace{b^{2}HWC}}+\underset{\text{Dense?layers}}{\underbrace{10HW{C}^2}}\text{\hspace{1em}(1)}$$

這相對于圖像尺寸 HW 呈線性關系，而像視覺 Transformer（ViT）、Mixer 和門控多層感知機（gMLP）等其他全局模型則呈二次方關系。

多軸方法的通用性。我們提出的并行多軸模塊（圖 3）提供了一種以可擴展方式在二維圖像上應用一維算子的原則性方法。它還具有顯著的靈活性和通用性。例如，直接將門控多層感知機（gMLP）替換為空間多層感知機 [87]、自注意力機制 [24]，甚至傅里葉變換 [44, 76]，就能得到一系列 MAXIM 變體，這些變體都具備全局性和全卷積性。它也很容易擴展到未來可能在語言模型等中定義的任何一維算子。

• 圖 3 多軸門控多層感知器（MLP）模塊。輸入首先被投影為一個 [6; 4; C] 的特征，然后拆分為兩個頭。
  在局部分支中，半個頭被劃分成 3×2 個不重疊的 [2; 2; C/2] 的補丁塊，而在全局分支中，我們使用 2×2 的網格劃分另一半。我們僅在每個分支的單一軸上應用門控多層感知器（gMLP）模塊 [54]（如右側的 gMLP 模塊所示）—— 局部分支應用于第二軸，全局分支應用于第一軸，同時在其他空間維度上共享。并行運行的 gMLP 操作符對應于局部和全局（擴張的）關注區域，如不同顏色所示（即，使用 gMLP 操作符對相同顏色的區域進行空間混合）。我們提出的模塊在任意輸入分辨率下都能表達全局和局部感受野。

Cross Gating MLP Block

對 U-Net 的一種常見改進方法是利用上下文特征，在跳躍連接中選擇性地控制特征傳播 [67,71]，這通常通過交叉注意力機制來實現 [13, 90]。在這里，我們構建了一種有效的替代方案，即交叉門控模塊（CGB，圖 2c），它是多軸門控多層感知機模塊（MAB，3.2 節）的擴展，因為 MAB 只能處理單個特征。CGB 可以看作是一種更通用的條件層，能夠與多種特征進行交互 [13, 70, 90]。我們采用了與 MAB 類似的設計模式。

更具體地說，設 $\mathbf{X},\mathbf{Y}$ 為兩個輸入特征，且 ${\mathbf{X}}_1,{\mathbf{Y}}_1\in {\mathbb{R}}^{H\times W\times C}$是圖 2c 中經過第一個全連接層（Dense 層）投影后的特征。然后進行輸入投影操作:

$${\mathbf{X}}_2=\sigma ({\mathbf{W}}_1\text{LN}({\mathbf{X}}_1)),{\mathbf{Y}}_2=\sigma ({\mathbf{W}}_2\text{LN}({\mathbf{Y}}_1))\text{\hspace{1em}(2)}$$

其中 $\sigma$ 是高斯誤差線性單元（GELU）激活函數 [30] ，LN 是層歸一化 [5] ，${\mathbf{W}}_1$、${\mathbf{W}}_2$是多層感知機（MLP）投影矩陣。多軸分塊門控權重分別從 ${\mathbf{X}}_2$, ${\mathbf{Y}}_2$ 計算得出，但應用方式相反：

$$\hat{\mathbf{X}}={\mathbf{X}}_2\odot G({\mathbf{Y}}_2),\hat{\mathbf{Y}}={\mathbf{Y}}_2\odot G({\mathbf{X}}_2)\text{\hspace{1em}(3)}$$

其中 $\odot$ 表示逐元素相乘，函數 $G(?)$ 使用我們提出的多軸方法從輸入中提取多軸交叉門控權重。

$$G(\mathbf{x})={\mathbf{W}}_5([{\mathbf{W}}_3{\text{Block}}_b({\mathbf{z}}_1),{\mathbf{W}}_4{\text{Grid}}_d({\mathbf{z}}_2)])\text{\hspace{1em}(4)}$$

其中$[?,?]$表示拼接操作。這里$({\mathbf{z}}_1,{\mathbf{z}}_2)$是沿著通道維度從$\mathbf{z}$拆分出的兩個獨立頭部，其中$\mathbf{z}$表示經過激活后的投影特征$\mathbf{x}$。

$$[{\mathbf{z}}_1,{\mathbf{z}}_2]=\mathbf{z}=\sigma ({\mathbf{W}}_6\text{LN}(\mathbf{x}))\text{\hspace{1em}(5)}$$

${\mathbf{W}}_3$、${\mathbf{W}}_4$是空間投影矩陣，分別應用于具有固定窗口大小 $b\times b$（${\text{Block}}_b$）的分塊特征的第二個軸，以及具有固定網格大小$d\times d$（${\text{Grid}}_d$）的網格化特征的第一個軸。最后，我們采用輸入的殘差連接，在輸出通道投影之后，使用投影矩陣${\mathbf{W}}_7$、${\mathbf{W}}_8$保持與輸入（${\mathbf{X}}_1$、${\mathbf{Y}}_1$）相同的通道維度，記為:

$${\mathbf{X}}_3={\mathbf{X}}_1+{\mathbf{W}}_7\hat{\mathbf{X}},{\mathbf{Y}}_3={\mathbf{Y}}_1+{\mathbf{W}}_8\hat{\mathbf{Y}}\text{\hspace{1em}(6)}$$

Multi-Stage Multi-Scale Framework

我們進一步采用了多階段框架，因為與增加模型寬度或高度相比（見 4.3A 節消融實驗），我們發現這種方法更有效。我們認為全分辨率處理 [16, 69, 77] 比多圖像塊分層結構 [83, 111, 113] 更好，因為后者可能會在圖像塊之間產生邊界效應。

為了實施更強的監督，我們在每個階段都采用多尺度方法 [18, 20, 48] 來輔助網絡學習。我們利用監督注意力模塊 [111] 沿各個階段逐步傳播注意力特征。我們使用交叉門控模塊進行跨階段特征融合。具體細節請讀者參考圖 9。

• 圖 9 我們采用了一種通用的多階段框架，以提升 MAXIM 在具有挑戰性的圖像修復任務中的性能。受文獻 [16, 111] 的啟發，我們使用了監督注意力模塊（SAM）以及跨階段特征融合，來輔助后續階段的學習。與先前的方法不同，由于所提出的多軸多層感知器（MLP）方法，我們的 MAXIM 主干網絡在每個階段的每一層都能實現全局感知，這使得它在學習低級和高級特征中的全局交互方面更加強大。

具體來說，給定一個輸入圖像 $I\in {\mathbb{R}}^{H\times W\times 3}$，我們首先通過下采樣提取其多尺度變體：$I_n$，$n=1,?,N$ 。MAXIM 在S個階段中的每個階段s預測多尺度的復原輸出，總共產生$S\times N$個輸出: $R_{s,n}$ 。盡管 MAXIM 是多階段的，但它是端到端訓練的，損失會在各個階段和尺度上累積:

$$\mathcal{L}=\sum _{s=1}^S\sum _{n=1}^N[{\mathcal{L}}_{char}({\mathbf{R}}_{s,n},{\mathbf{T}}_n)+\lambda {\mathcal{L}}_{freq}({\mathbf{R}}_{s,n},{\mathbf{T}}_n)]\text{\hspace{1em}(7)}$$

其中${\mathbf{T}}_n$表示（經過雙線性縮放的）多尺度目標圖像，${\mathcal{L}}_{char}$是 Charbonnier 損失函數 [111]：

$${\mathcal{L}}_{char}(\mathbf{R},\mathbf{T})=\sqrt{\Vert \mathbf{R}?\mathbf{T}{\Vert }^2+{?}^2}\text{\hspace{1em}(8)}$$

這里我們將$?$設為$10^{?3}$ 。${\mathcal{L}}_{freq}$是頻率重建損失函數，用于強化高頻細節 [20, 35]：
$${\mathcal{L}}_{freq}(\mathbf{R},\mathbf{T})=\Vert \mathcal{F}(\mathbf{R})?\mathcal{F}(\mathbf{T}){\Vert }_1\text{\hspace{1em}(9)}$$

其中$\mathcal{F}(?)$表示二維快速傅里葉變換。在所有實驗中，我們使用$\lambda =0.1$作為加權因子。

Experiments

我們旨在為廣泛的圖像處理任務構建一個通用的骨干網絡。因此，我們在 17 個不同的數據集上針對五種不同的任務對 MAXIM 進行了評估：（1）去噪、（2）去模糊、（3）去雨、（4）去霧，以及（5）增強（潤飾）（匯總于表 8 ）。

• 表 8

Experimental Setup

• 數據集與評估指標
  我們通過測量真實圖像與預測圖像之間的峰值信噪比（PSNR）和結構相似性指數（SSIM）[96] 來進行定量比較。去噪任務使用 SIDD [2] 和 DND [72] 數據集；去模糊任務使用 GoPro [62]、HIDE [81] 和 RealBlur [79] 數據集；去雨任務使用 [111] 中用到的 Rain13k 組合數據集；去霧任務使用 RESIDE [46] 數據集；圖像增強（潤飾）任務則在 Five - K [8] 和 LOL [98] 數據集上進行評估。
• 訓練細節
  我們提出的 MAXIM 模型可進行端到端訓練，既無需大規模預訓練，也無需漸進式訓練。網絡在 256×256 的隨機裁剪圖像塊上進行訓練。針對不同任務，訓練的迭代次數有所不同。數據增強方面，我們采用隨機水平和垂直翻轉、90° 旋轉，以及概率為 0.5 的 MixUp [112] 方法。優化器使用 Adam [39]，初始學習率設為 2×10?? ，并通過余弦退火衰減 [59] 逐步降至 10?? 。測試時，我們使用對稱填充將輸入圖像填充為 64×64 的倍數，推理完成后，再將填充后的圖像裁剪回原始尺寸。關于各任務的更多訓練細節見附錄 A.1。
• 架構配置
  我們設計了兩種 MAXIM 變體：用于不同任務的兩階段模型 MAXIM - 2S 和三階段模型 MAXIM - 3S 。特征提取初始通道數設為 32，包含 3 個下采樣層，特征從 2562×32 經 1282×64 、642×128 ，收縮至 322×256 ，并由兩個瓶頸層處理（圖 2a），然后對稱擴展回全分辨率。當應用于 256×256 圖像時，MAXIM - 2S 和 MAXIM - 3S 的參數量和所需浮點運算次數（FLOPs）列于表 7A 的最后兩行。