CVPR 2022 Oral | MLP進軍底層視覺！谷歌提出MAXIM模型刷榜多個圖像處理任務，代碼已開源
作者：假熊貓
鏈接：https://zhuanlan.zhihu.com/p/481256924
你是否厭倦了最新的Transformer/MLP模型"不靈活性"和"高空間復雜度"？
你是否對最新的Transformer/MLP模型無法自適應的應用在"不同分辨率"」感到無奈？
來試試 "MAXIM "模型吧！最新的基于 MLP的UNet型骨干網(wǎng)絡，同時實現(xiàn)了"全局"與"局部”感受野，并且可以在"線性復雜度"下直接應用在"高分辨圖片"上，具有“全卷積”特性，可以"即插即用"！代碼和模型均已開源！***家人們，你還在等什么呢***？！
1太長不看
提出了通用的圖像修復/增強任務骨干網(wǎng)絡 MAXIM，第一次把最近爆火的 「MLP」[1]應用在底層視覺，在五大圖像處理任務（去噪，去模糊，去雨，去霧，增強）超過10個數(shù)據(jù)集達到SOTA性能； 提出一個「即插即用」的多軸門限MLP模塊（Multi-Axis gMLP block），實現(xiàn)了線性復雜度下的全局 / 局部的空間信息交互，解決了MLP/Transformer無法處理不同分辨率圖片的痛點[2]，并且具有全卷積[3]特性，為底層視覺任務量身定做，也可以應用在其他的密集預測任務（留給未來填坑）； 提出另一個「即插即用」的交叉門控模塊（Cross-Gating MLP block），可以無痛替代交叉注意力機制，并且同樣在線性復雜度享有全局 / 局部感受野和全卷積特性。
論文地址：https://arxiv.org/abs/2201.02973
代碼/模型/實驗結(jié)果：https://github.com/google-research/maxim
中文視頻講解：https://youtu.be/gpUrUJwZxRQ（非常詳細，有很多背景介紹，新手友好型）
2背景介紹（Introduction）
都2022年了，你還在執(zhí)著于在「卷積神經(jīng)網(wǎng)絡」中調(diào)參內(nèi)卷嗎？Vision Transformer (ViT)發(fā)表才不過一年多，就已經(jīng)在各大視覺領域「鯊」「瘋」「了」！受到 ViT 優(yōu)雅架構(gòu)的啟發(fā)，各種奇技淫巧也應運而生 —— 谷歌大腦提出的MLP-Mixer把 自注意力機制換成MLP，構(gòu)建一個純MLP架構(gòu)，性能威猛！谷歌另一個大腦提出 gMLP 模型，構(gòu)建了門限MLP模塊，在視覺和語言建模上均無痛吊打 Transformer！有知乎大V不禁發(fā)問：MLP is all you need？
ViT, Mixer, gMLP這些新的視覺骨干網(wǎng)絡帶領了了一波從根本上區(qū)別于傳統(tǒng)卷積神經(jīng)網(wǎng)絡（CNN）架構(gòu)設計的模式轉(zhuǎn)變（paradigm shift），即為「全局模型」（Global Models or Non-Local Networks) —— 我們不再依賴于長期以來人們對二維圖像的先驗知識（prior)：平移不變性和局部依賴；而是無腦使用全局感受野和超大規(guī)模數(shù)據(jù)預訓練的「鈔能力」。當然另一個ViT的特性是從注意力機制的本身定義而來，即為自適應于輸入的動態(tài)加權(quán)平均，但這里我們主要討論這些Transformer-like模型的全局交互屬性。
全局模型允許在輸入的特征圖上進行全局空間交互，即每個輸出像素是由輸入特征的每個點加權(quán)而來，需要次乘法操作，這即為注意力機制/Transformer高計算復雜度的由來。但本質(zhì)上來說，密集感受野的全局模型如 ViT, Mixer, gMLP 都具有平方計算復雜度。這種沒法Scaling up的平方算子是很難作為通用模塊來廣泛使用在各大視覺任務上的，例如需要在高分辨率上訓練/推理的目標檢測，語義分割等，甚至對于幾乎所有的底層視覺任務如去噪、去模糊、超分、去雨、去霧、去雪、去陰影、去摩爾、去反射、去水印、去馬賽克等等等等。。。
雖然但是，不妨也直接用用！華為北大等聯(lián)手打造的IPT模型第一次把ViT模型應用在多個底層視覺任務，刷新了各大榜單并發(fā)表在CVPR 2021。雖然性能很好，但IPT使用的全局注意力機制具有一些明顯的局限性：（1）需要大量數(shù)據(jù)預訓練（如ImageNet），（2）無法直接在高分辨率圖片上進行推理。在實際推理時，往往需要對輸入圖像進行切塊，分別對每個圖像塊進行推理，然后再進行拼接來還原大圖。這種辦法往往會導致輸出圖片中有一些明顯的“塊狀效應”（如下圖），同時推理速度也比較慢，限制了其實際落地和部署能力。
這種能在小圖像塊上訓練，并且直接在大圖上推理的屬性我們稱之為“全卷積”（fully-convolutional）[3]。全卷積屬性對于底層視覺任務來說至關重要，因為底層視覺如圖像修復和增強均需要對圖像進行像素級操作，輸出圖像需和輸入圖像尺寸一樣，無法像圖像分類一樣先進行resize等操作。很明顯，目前的主流全局網(wǎng)絡 ViT，Mixer，gMLP都無法很好的解決這個無法自適應于不同圖像分辨率的痛點。
他來了，他來了，他披著CNN的外皮來了！
Swin Transformer橫空出世，并且一舉摘獲ICCV 2021 Marr Prize。Swin的貢獻可圈可點，譬如引入了層次化的結(jié)構(gòu)，譬如提出了局部注意力機制（Local/Window attention）來解決計算復雜度問題。但是，重點來了！Swin提出的local attention為底層視覺帶來了福音：具有“全卷積 ”屬性！其根本原因是 Self-Attention 作用在一個小的 7 x 7 窗口內(nèi)，而在整個空間內(nèi)的不同窗口都是共享權(quán)重的。如果在更大的尺寸下推理，無非是有更多的窗口罷了。想想看，這玩意兒是不是跟 Strided Convolution 的思想如出一轍？CNN，yyds！
局部注意力作為一個十分自然的受啟發(fā)于CNN的改進，非常適用于底層視覺任務，因此馬不停蹄的就進軍各大底層視覺任務?！笆桩斊錄_”的兩個工作就是 （1）中科大提出的 Uformer （CVPR 22）和（2）Luc Van Gool組提出的SwinIR(ICCVW 21)，均借鑒或改進了Local Attention的思想，并且應用在多個不同的底層視覺任務，取得了驚人的性能。
但是，Local Attention重新引入了Locality的思想，返璞歸真，反而把全局模型中很重要的一個特性"全局交互"給舍棄掉了。我們謙虛地認為Shifted window attention只是對Local Attention做了一個補充，并沒有真正解決怎樣更高效的進行全局交互的痛點（個人觀點，大神請輕噴QAQ。不過這是本文的研究動機之一）
3模型方法（Method）
我們設計了第一個基于MLP的通用底層視覺類UNet骨干網(wǎng)絡稱之為MAXIM。對比前人的一些底層視覺網(wǎng)絡工作，MAXIM具有以下幾個優(yōu)點：
MAXIM 在任意尺寸圖片上都具有全局感受野，并且只需要線性復雜度； MAXIM 可以在任意尺寸圖片上直接推理，具有“全卷積”屬性； MAXIM 平衡了局部和全局算子的使用，使得網(wǎng)絡在不需要超大數(shù)據(jù)集預訓練的情況下達到SOTA性能。
MAXIM骨干網(wǎng)絡架構(gòu)圖如上所示，其具有一個對稱UNet的基本結(jié)構(gòu)，包含降采樣的Encoder模塊，最底層的Bottleneck, 和上采樣的Decoder模塊。其中，每一塊Encoder/Decoder/Bottleneck均采用同樣的設計如Figure 2(b)：多軸門控MLP塊（全局交互）和殘差卷積通道注意力塊（局部交互）。受啟發(fā)于Attention-UNet，我們在UNet的中間層加入了 交叉門控模塊（Cross-gating block)，使用Bottleneck輸出的高階語義特征來調(diào)制編碼器到解碼器之間的跳躍連接特征。值得注意的是，區(qū)別于傳統(tǒng)的各種UNet魔改網(wǎng)絡，MAXIM骨干中的每一個模塊都具有全局/局部感受野，因此具有更大的學習潛力。
魔改一：多軸門控MLP模塊
本模型中的核心貢獻就是提出的多軸門控MLP模塊（Multi-axis gated MLP block），一個即插即用的并行模塊，可進行全局/局部的空間交互，并且具有線性復雜度。我們是受啟發(fā)于 [NeurIPS 21] HiT-GA中提出的多軸自注意力模塊，可在低分辨率特征圖上進行有效的全局/局部信息交互，在多個圖片生成任務達到SOTA。然而，此多軸非彼多軸，我們要能夠使用在高分辨率底層視覺任務上，并且同時需要具有“全卷積”屬性。但同時又不想犧牲全局感受野的重要屬性，于是乎，魔改開始了：
如上圖所示，輸入的特征首先進行通道映射，然后分成兩個頭，分別進行全局和局部交互。其中一半的頭進入局部分支（圖中紅色），我們使用 gMLP 算子在固定的窗口大小內(nèi)進行局部空間交互；另一半頭喂進全局分支（圖中綠色），我們同樣使用 gMLP 算子在固定的網(wǎng)格位置進行全局（膨脹）空間交互。值得注意的是，圖中的 Block 和 Grid 操作均為窗口劃分（和Swin一樣），但Block操作中我們固定【窗口大小】，而在 Grid 操作中我們固定【窗口數(shù)量】（or 網(wǎng)格大小）。在兩個并行分支結(jié)構(gòu)中，我們每次只對固定維度的坐標進行操作，而在其他坐標都共享參數(shù)，從而實現(xiàn)了同時具有“全卷積”屬性和全局/局部感受野。由于我們總是使用固定的窗口大小和網(wǎng)格大小，該模塊也具有線性計算復雜。
魔改二：交叉門控模塊
UNet比較經(jīng)典的魔改網(wǎng)絡是Attention-UNet，其在對稱的跳躍連接中加入了交叉注意力機制來自適應的加權(quán)濾波可以通過的特征圖。受此啟發(fā)，我們進行了第二個魔改，提出了【交叉門控模塊】，如 Figure 2(c)所示。其設計理念嚴格遵守多軸門控MLP模塊的模范，同樣采用多軸全局/局部交互的gMLP模塊。唯一的區(qū)別是在提取了gMLP算子的空間門權(quán)重（gating weights）后，我們采用了交叉相乘的方式來進行信息交互。例如 是兩個不同的輸入特征，交叉門控的概念可以簡單表示為（具體的公式可以參見文章或代碼）：
至此，我們提出了第一個可以進行多特征交互的純純MLP的交叉門控模塊，可以用來進行全局/局部的交叉信息傳遞和互相調(diào)制，功能上等效于交叉注意力機制，可以無腦即插即用。
魔改三：多階段多尺度架構(gòu)
電路圖預警！??！電路圖預警?。。‰娐穲D預警?。?！
為了平衡性能-計算復雜度，MAXIM采用了一個改進的多階段網(wǎng)絡，并且采用了深度監(jiān)督策略來監(jiān)督多階段多尺度的輸出們。本文中針對不同的任務分別使用了2和3階段網(wǎng)絡：MAXIM-2S，MAXIM-3S。雖然MAXIM是多階段網(wǎng)絡，其仍然是可以進行端到端訓練而不需要分步或漸進訓練。在推理階段，只需要把最后階段的最大尺寸輸出保留作為最終的結(jié)果即可。使用的損失函數(shù)是把多個階段、多個尺度所有的輸出和輸入計算Charbonnier損失函數(shù)和頻域變換后的L1損失的加權(quán)和
其中代表網(wǎng)絡在階段s尺度n的輸出圖像，表示尺度n的目標圖片(groundtruth)。該多階段多尺度的設計參考了一些前人工作中的網(wǎng)絡設計經(jīng)驗如MPRNet，HINet，和MIMO-UNet。
不得不說此模塊是作者至今的唯一遺憾！一直覺得提出的網(wǎng)絡結(jié)構(gòu)太復雜了，不夠優(yōu)雅。在作者早期調(diào)參時使用單個MAXIM骨干一直沒法達到接近SOTA的性能，痛苦、徘徊、迷茫了好久，浪費了很多碳排放。后來才明白了一個經(jīng)驗教訓：一定要respect每個領域的domain knowledge，不要輕易懷疑為什么該領域的網(wǎng)絡這樣那樣設計，那都是前輩調(diào)參俠們996的血汗成果啊T.T。多階段網(wǎng)絡在比較難的圖像修復任務如去模糊和去雨中已經(jīng)成為了經(jīng)典架構(gòu)；同樣地，MAXIM在使用了多階段架構(gòu)后，性能也立馬起飛。令人喜出望外的是，使用多階段小網(wǎng)絡比單階段大網(wǎng)絡的性能與計算量均有很大提升（參見消融實驗）！驚不驚喜，意不意外？
4實驗（Experiments）
實驗設置
我們旨在建立一個大一統(tǒng)的骨干網(wǎng)絡可以適用于廣泛的底層視覺/圖像處理任務。因此，我們在五個不同的視覺任務多達 17 個數(shù)據(jù)集上進行了訓練測試。使用的數(shù)據(jù)集總結(jié)如下：
實驗結(jié)果
量化和視覺的實驗結(jié)果如下所示。一圖勝千言，就不多贅述了。更多的實驗結(jié)果參見論文的附錄部分。
1. 去噪（Denoising)?
2. 去模糊（Deblurring)?
3. 去雨（Deraining）
4. 去霧（Dehazing）
5. （光照）增強（Enhancement）
消融實驗
我們做了詳盡的消融實驗來理解MAXIM網(wǎng)絡：
【模塊消融】我們發(fā)現(xiàn)使用每一個新提出的模塊都對最終性能有提升，測試的模塊有階段內(nèi)和階段之間的的交叉門限模塊，SAM模塊，和多尺寸深度監(jiān)督方法； 本文主要提出了可以適用于高分辨率圖片的全局交互MLP網(wǎng)絡，那么問題來了：【使用全局有多好】？消融實驗B表示局部和全局MLP對網(wǎng)絡的提升效果不相上下，合在一起食用效果更佳； 【為什么要使用多階段？】實踐是檢驗真理的唯一標準。我們發(fā)現(xiàn)使用多階段比使用更深、更寬的單個網(wǎng)絡的性能提升更加明顯，并且參數(shù)和計算量也比較平衡； 【通用性】我們提出的多軸并行模塊是一個通用辦法可以把無法處理不同分辨率的算子轉(zhuǎn)化成局部/全局算子，并且具有線性復雜度和分辨率自適應性。我們嘗試了使用自注意力，gMLP，MLP-Mixer，F(xiàn)FT作為不同的空間混合算子，發(fā)現(xiàn)使用自注意力和gMLP可以取得最佳的性能而使用Mixer和FFT具有更快的計算速度。