CVPR22 Oral｜MLP進(jìn)軍底層視覺！谷歌提出MAXIM模型刷榜多個(gè)圖像處理任務(wù)，代碼已開源

你是否厭倦了最新的Transformer/MLP模型的「不靈活性」和「高空間復(fù)雜度」？

你是否對最新的Transformer/MLP模型無法自適應(yīng)的應(yīng)用在「不同分辨率」感到無奈？

來試試 「MAXIM 」模型吧！最新的基于 MLP的UNet型骨干網(wǎng)絡(luò)，同時(shí)實(shí)現(xiàn)了「全局」「局部」感受野，并且可以在「線性復(fù)雜度」下直接應(yīng)用在「高分辨圖片」上，具有「“全卷積”」特性，可以「即插即用」！代碼和模型均已開源！_家人們，你還在等什么呢？！_

太長不看（TL；DR）

• 提出了通用的圖像修復(fù)/增強(qiáng)任務(wù)骨干網(wǎng)絡(luò) MAXIM，第一次把最近爆火的 「MLP」[1]應(yīng)用在底層視覺，在五大圖像處理任務(wù)（去噪，去模糊，去雨，去霧，增強(qiáng)）超過10個(gè)數(shù)據(jù)集達(dá)到SOTA性能；
• 提出一個(gè)「即插即用」的多軸門限MLP模塊（Multi-Axis gMLP block），實(shí)現(xiàn)了線性復(fù)雜度下的全局 / 局部的空間信息交互，解決了MLP/Transformer無法處理不同分辨率圖片的痛點(diǎn)[2]，并且具有全卷積[3]特性，為底層視覺任務(wù)量身定做，也可以應(yīng)用在其他的密集預(yù)測任務(wù)（留給未來填坑）；
• 提出另一個(gè)「即插即用」的交叉門控模塊（Cross-Gating MLP block），可以無痛替代交叉注意力機(jī)制，并且同樣在線性復(fù)雜度享有全局 / 局部感受野和全卷積特性。

論文地址：https://arxiv.org/abs/2201.02973

代碼/模型/實(shí)驗(yàn)結(jié)果：https://github.com/google-research/maxim

中文視頻講解：https://youtu.be/gpUrUJwZxRQ（非常詳細(xì)，有很多背景介紹，新手友好型）

背景介紹（Introduction）

都2022年了，你還在執(zhí)著于在「卷積神經(jīng)網(wǎng)絡(luò)」中調(diào)參內(nèi)卷嗎？

Vision Transformer (ViT)[4]發(fā)表才不過一年多，就已經(jīng)在各大視覺領(lǐng)域「鯊」「瘋」「了」！受到 ViT 優(yōu)雅架構(gòu)的啟發(fā)，各種奇技淫巧也應(yīng)運(yùn)而生 —— 谷歌大腦提出的MLP-Mixer[5]把 自注意力機(jī)制換成MLP，構(gòu)建一個(gè)純MLP架構(gòu)，性能威猛！谷歌另一個(gè)大腦提出 gMLP 模型[6]，構(gòu)建了門限MLP模塊，在視覺和語言建模上均無痛吊打 Transformer！有知乎大V不禁發(fā)問：MLP is all you need？[1]

ViT, Mixer, gMLP這些新的視覺骨干網(wǎng)絡(luò)帶領(lǐng)了了一波從根本上區(qū)別于傳統(tǒng)卷積神經(jīng)網(wǎng)絡(luò)（CNN）架構(gòu)設(shè)計(jì)的模式轉(zhuǎn)變（paradigm shift），即為「全局模型」（Global Models or Non-Local Networks[7]) —— 我們不再依賴于長期以來人們對二維圖像的先驗(yàn)知識（prior)：平移不變性和局部依賴；而是無腦使用全局感受野和超大規(guī)模數(shù)據(jù)預(yù)訓(xùn)練的「鈔能力」[8]。當(dāng)然另一個(gè)ViT的特性是從注意力機(jī)制的本身定義而來，即為自適應(yīng)于輸入的動態(tài)加權(quán)平均，但這里我們主要討論這些Transformer-like模型的全局交互屬性。

全局模型允許在輸入的特征圖上進(jìn)行全局空間交互，即每個(gè)輸出像素是由輸入特征的每個(gè)點(diǎn)加權(quán)而來，需要 次乘法（假設(shè) 為空間尺寸）。因此，輸出整張大小為 的特征圖需要 次乘法操作，這即為注意力機(jī)制/Transformer高計(jì)算復(fù)雜度的由來。但本質(zhì)上來說，密集感受野的全局模型如 ViT, Mixer, gMLP 都具有平方計(jì)算復(fù)雜度。 這種沒法Scaling up的平方算子是很難作為通用模塊來廣泛使用在各大視覺任務(wù)上的，例如需要在高分辨率上訓(xùn)練/推理的目標(biāo)檢測，語義分割等，甚至對于幾乎所有的底層視覺任務(wù)如去噪、去模糊、超分、去雨、去霧、去雪、去陰影、去摩爾、去反射、去水印、去馬賽克等等等等。。。

雖然但是，不妨也直接用用！華為北大等聯(lián)手打造的IPT模型第一次把ViT模型應(yīng)用在多個(gè)底層視覺任務(wù)，刷新了各大榜單并發(fā)表在CVPR 2021[9][10]。雖然性能很好，但I(xiàn)PT使用的全局注意力機(jī)制具有一些明顯的局限性：（1）需要大量數(shù)據(jù)預(yù)訓(xùn)練（如ImageNet），（2）無法直接在高分辨率圖片上進(jìn)行推理。在實(shí)際推理時(shí)，往往需要對輸入圖像進(jìn)行切塊，分別對每個(gè)圖像塊進(jìn)行推理，然后再進(jìn)行拼接來還原大圖。這種辦法往往會導(dǎo)致輸出圖片中有一些明顯的“塊狀效應(yīng)”（如下圖），同時(shí)推理速度也比較慢，限制了其實(shí)際落地和部署能力。

這種能在小圖像塊上訓(xùn)練，并且直接在大圖上推理的屬性我們稱之為“全卷積”（fully-convolutional）[3]。全卷積屬性對于底層視覺任務(wù)來說至關(guān)重要，因?yàn)榈讓右曈X如圖像修復(fù)和增強(qiáng)均需要對圖像進(jìn)行像素級操作，輸出圖像需和輸入圖像尺寸一樣，無法像圖像分類一樣先進(jìn)行resize等操作。很明顯，目前的主流全局網(wǎng)絡(luò) ViT，Mixer，gMLP都無法很好的解決這個(gè)無法自適應(yīng)于不同圖像分辨率的痛點(diǎn)。

他來了，他來了，他披著CNN的外皮來了！

Swin Transformer橫空出世，并且一舉摘獲ICCV 2021 Marr Prize[11]。 Swin的貢獻(xiàn)可圈可點(diǎn)，譬如引入了層次化的結(jié)構(gòu)，譬如提出了局部注意力機(jī)制（Local/Window attention）來解決計(jì)算復(fù)雜度問題。但是，重點(diǎn)來了！Swin提出的local attention為底層視覺帶來了福音：具有“全卷積 ”屬性！其根本原因是 Self-Attention 作用在一個(gè)小的 7 x 7 窗口內(nèi)，而在整個(gè)空間內(nèi)的不同窗口都是共享權(quán)重的。如果在更大的尺寸下推理，無非是有更多的窗口罷了。想想看，這玩意兒是不是跟 Strided Convolution 的思想如出一轍？CNN，yyds！

局部注意力作為一個(gè)十分自然的受啟發(fā)于CNN的改進(jìn)，非常適用于底層視覺任務(wù)，因此馬不停蹄的就進(jìn)軍各大底層視覺任務(wù)?！笆桩?dāng)其沖”的兩個(gè)工作就是 （1）中科大提出的 Uformer [12]（CVPR 22）和（2）Luc Van Gool組提出的SwinIR[13] (ICCVW 21)，均借鑒或改進(jìn)了Local Attention的思想，并且應(yīng)用在多個(gè)不同的底層視覺任務(wù)，取得了驚人的性能。

但是，Local Attention重新引入了Locality的思想，返璞歸真，反而把全局模型中很重要的一個(gè)特性「全局交互」給舍棄掉了。我們謙虛地認(rèn)為Shifted window attention只是對Local Attention做了一個(gè)補(bǔ)充，并沒有真正解決怎樣更高效的進(jìn)行全局交互的痛點(diǎn)（個(gè)人觀點(diǎn)，大神請輕噴QAQ。。不過這是本文的研究動機(jī)之一）

模型方法（Method）

我們設(shè)計(jì)了第一個(gè)基于MLP的通用底層視覺類UNet骨干網(wǎng)絡(luò)稱之為MAXIM。對比前人的一些底層視覺網(wǎng)絡(luò)工作，MAXIM具有以下幾個(gè)優(yōu)點(diǎn)：

• MAXIM 在任意尺寸圖片上都具有全局感受野，并且只需要線性復(fù)雜度；
• MAXIM 可以在任意尺寸圖片上直接推理，具有“全卷積”屬性；
• MAXIM 平衡了局部和全局算子的使用，使得網(wǎng)絡(luò)在不需要超大數(shù)據(jù)集預(yù)訓(xùn)練的情況下達(dá)到SOTA性能。

MAXIM骨干網(wǎng)絡(luò)架構(gòu)圖如上所示，其具有一個(gè)對稱UNet的基本結(jié)構(gòu)，包含降采樣的Encoder模塊，最底層的Bottleneck, 和上采樣的Decoder模塊。其中，每一塊Encoder/Decoder/Bottleneck均采用同樣的設(shè)計(jì)如Figure 2(b)：多軸門控MLP塊（全局交互）和殘差卷積通道注意力塊（局部交互）。受啟發(fā)于Attention-UNet[14]，我們在UNet的中間層加入了 交叉門控模塊（Cross-gating block)，使用Bottleneck輸出的高階語義特征來調(diào)制編碼器到解碼器之間的跳躍連接特征。值得注意的是，區(qū)別于傳統(tǒng)的各種UNet魔改網(wǎng)絡(luò)，MAXIM骨干中的每一個(gè)模塊都具有全局/局部感受野，因此具有更大的學(xué)習(xí)潛力。

魔改一：多軸門控MLP模塊

本模型中的核心貢獻(xiàn)就是提出的多軸門控MLP模塊（Multi-axis gated MLP block），一個(gè)即插即用的并行模塊，可進(jìn)行全局/局部的空間交互，并且具有線性復(fù)雜度。我們是受啟發(fā)于 [NeurIPS 21] HiT-GAN[15]中提出的多軸自注意力模塊，可在低分辨率特征圖上進(jìn)行有效的全局/局部信息交互，在多個(gè)圖片生成任務(wù)達(dá)到SOTA。然而，此多軸非彼多軸，我們要能夠使用在高分辨率底層視覺任務(wù)上，并且同時(shí)需要具有“全卷積”屬性。但同時(shí)又不想犧牲全局感受野的重要屬性，于是乎，魔改開始了：

如上圖所示，輸入的特征首先進(jìn)行通道映射，然后分成兩個(gè)頭，分別進(jìn)行全局和局部交互。 其中一半的頭進(jìn)入局部分支（圖中紅色），我們使用 gMLP 算子在固定的窗口大小內(nèi)進(jìn)行局部空間交互；另一半頭喂進(jìn)全局分支（圖中綠色），我們同樣使用 gMLP 算子在固定的網(wǎng)格位置進(jìn)行全局（膨脹）空間交互。 值得注意的是，圖中的 Block 和 Grid 操作均為窗口劃分（和Swin一樣）， 但Block操作中我們固定【窗口大小】，而在 Grid 操作中我們固定【窗口數(shù)量】（or 網(wǎng)格大小）。在兩個(gè)并行分支結(jié)構(gòu)中，我們每次只對固定維度的坐標(biāo)進(jìn)行操作，而在其他坐標(biāo)都共享參數(shù)，從而實(shí)現(xiàn)了同時(shí)具有“全卷積”屬性和全局/局部感受野。由于我們總是使用固定的窗口大小和網(wǎng)格大小，該模塊也具有線性計(jì)算復(fù)雜度 。

魔改二：交叉門控模塊

UNet比較經(jīng)典的魔改網(wǎng)絡(luò)是Attention-UNet[14]，其在對稱的跳躍連接中加入了交叉注意力機(jī)制來自適應(yīng)的加權(quán)濾波可以通過的特征圖。受此啟發(fā)，我們進(jìn)行了第二個(gè)魔改，提出了【交叉門控模塊】，如 Figure 2(c)所示。其設(shè)計(jì)理念嚴(yán)格遵守多軸門控MLP模塊的模范，同樣采用多軸全局/局部交互的gMLP模塊。唯一的區(qū)別是在提取了gMLP算子的空間門權(quán)重（gating weights） 后，我們采用了交叉相乘的方式來進(jìn)行信息交互。例如 是兩個(gè)不同的輸入特征，交叉門控的概念可以簡單表示為（具體的公式可以參見文章或代碼）：

至此，我們提出了第一個(gè)可以進(jìn)行多特征交互的純純MLP的交叉門控模塊，可以用來進(jìn)行全局/局部的交叉信息傳遞和互相調(diào)制， 功能上等效于交叉注意力機(jī)制，可以無腦即插即用。

魔改三：多階段多尺度架構(gòu)

電路圖預(yù)警?。。‰娐穲D預(yù)警?。?！電路圖預(yù)警！?。?/p>

為了平衡性能-計(jì)算復(fù)雜度，MAXIM采用了一個(gè)改進(jìn)的多階段網(wǎng)絡(luò)，并且采用了深度監(jiān)督策略來監(jiān)督多階段多尺度的輸出們。本文中針對不同的任務(wù)分別使用了2和3階段網(wǎng)絡(luò)：MAXIM-2S，MAXIM-3S。雖然MAXIM是多階段網(wǎng)絡(luò)，其仍然是可以進(jìn)行端到端訓(xùn)練而不需要分步或漸進(jìn)訓(xùn)練。在推理階段，只需要把最后階段的最大尺寸輸出保留作為最終的結(jié)果即可。使用的損失函數(shù)是把多個(gè)階段、多個(gè)尺度所有的輸出和輸入計(jì)算Charbonnier損失函數(shù)和頻域變換后的L1損失的加權(quán)和，可以表示為：

其中 代表網(wǎng)絡(luò)在階段s尺度n的輸出圖像， 表示尺度n的目標(biāo)圖片(groundtruth)。該多階段多尺度的設(shè)計(jì)參考了一些前人工作中的網(wǎng)絡(luò)設(shè)計(jì)經(jīng)驗(yàn)如MPRNet[16]，HINet[17]，和MIMO-UNet[18]。

[PS] 不得不說此模塊是作者至今的唯一遺憾！一直覺得提出的網(wǎng)絡(luò)結(jié)構(gòu)太復(fù)雜了，不夠優(yōu)雅。在作者早期調(diào)參時(shí)使用單個(gè)MAXIM骨干一直沒法達(dá)到接近SOTA的性能，痛苦、徘徊、迷茫了好久，浪費(fèi)了很多碳排放。后來才明白了一個(gè)經(jīng)驗(yàn)教訓(xùn)：一定要respect每個(gè)領(lǐng)域的domain knowledge，不要輕易懷疑為什么該領(lǐng)域的網(wǎng)絡(luò)這樣那樣設(shè)計(jì)，那都是前輩調(diào)參俠們996的血汗成果啊T.T。多階段網(wǎng)絡(luò)在比較難的圖像修復(fù)任務(wù)如去模糊和去雨中已經(jīng)成為了經(jīng)典架構(gòu)；同樣地，MAXIM在使用了多階段架構(gòu)后，性能也立馬起飛。令人喜出望外的是，使用多階段小網(wǎng)絡(luò)比單階段大網(wǎng)絡(luò)的性能與計(jì)算量均有很大提升（參見消融實(shí)驗(yàn)）！驚不驚喜，意不意外？

實(shí)驗(yàn)（Experiments）

實(shí)驗(yàn)設(shè)置

我們旨在建立一個(gè)大一統(tǒng)的骨干網(wǎng)絡(luò)可以適用于廣泛的底層視覺/圖像處理任務(wù)。因此，我們在五個(gè)不同的視覺任務(wù)多達(dá) 17 個(gè)數(shù)據(jù)集上進(jìn)行了訓(xùn)練測試。使用的數(shù)據(jù)集總結(jié)如下：

實(shí)驗(yàn)結(jié)果

量化和視覺的實(shí)驗(yàn)結(jié)果如下所示。一圖勝千言，就不多贅述了。更多的實(shí)驗(yàn)結(jié)果參見論文的附錄部分。

1. 去噪（Denoising) on SIDD, DND 數(shù)據(jù)集

2. 去模糊（Deblurring) on GoPro, HIDE，RealBlur，REDS數(shù)據(jù)集

3. 去雨（Deraining）on Rain100L, Rain100H, Test100, Test1200, Test2800, RainDrop數(shù)據(jù)集

4. 去霧（Dehazing）on RESIDE indoor，outdoor數(shù)據(jù)集

5. （光照）增強(qiáng)（Enhancement）on MIT-Adobe FiveK，LOL數(shù)據(jù)集

消融實(shí)驗(yàn)

我們做了詳盡的消融實(shí)驗(yàn)來理解MAXIM網(wǎng)絡(luò)：

• 【模塊消融】我們發(fā)現(xiàn)使用每一個(gè)新提出的模塊都對最終性能有提升，測試的模塊有階段內(nèi)和階段之間的的交叉門限模塊，SAM模塊，和多尺寸深度監(jiān)督方法；
• 本文主要提出了可以適用于高分辨率圖片的全局交互MLP網(wǎng)絡(luò)，那么問題來了：【使用全局有多好】？消融實(shí)驗(yàn)B表示局部和全局MLP對網(wǎng)絡(luò)的提升效果不相上下，合在一起食用效果更佳；
• 【為什么要使用多階段？】_實(shí)踐是檢驗(yàn)真理的唯一標(biāo)準(zhǔn)_。我們發(fā)現(xiàn)使用多階段比使用更深、更寬的單個(gè)網(wǎng)絡(luò)的性能提升更加明顯，并且參數(shù)和計(jì)算量也比較平衡；
• 【通用性】我們提出的多軸并行模塊是一個(gè)通用辦法可以把無法處理不同分辨率的算子轉(zhuǎn)化成局部/全局算子，并且具有線性復(fù)雜度和分辨率自適應(yīng)性。我們嘗試了使用自注意力，gMLP，MLP-Mixer，F(xiàn)FT[19]作為不同的空間混合算子，發(fā)現(xiàn)使用自注意力和gMLP可以取得最佳的性能而使用Mixer和FFT具有更快的計(jì)算速度。

參考

1. ^ab[吐槽] MLP is all you need??https://zhuanlan.zhihu.com/p/370780575
2. ^CVPR2021提出的Transformer無法直接處理不同分辨率?https://github.com/huawei-noah/Pretrained-IPT/issues/18
3. ^ab“全卷積”(fully-convolutional)：卷積網(wǎng)絡(luò)天然具有的特性，即可以應(yīng)用在不同分辨率的圖片?https://arxiv.org/abs/1411.4038
4. ^An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale?https://arxiv.org/abs/2010.11929
5. ^MLP-Mixer: An all-MLP Architecture for Vision?https://arxiv.org/abs/2105.01601
6. ^Pay attention to MLPs?https://proceedings.neurips.cc/paper/2021/hash/4cc05b35c2f937c5bd9e7d41d3686fff-Abstract.html
7. ^Non-Local Neural Networks?https://openaccess.thecvf.com/content_cvpr_2018/html/Wang_Non-Local_Neural_Networks_CVPR_2018_paper.html
8. ^GPT-3: Money is All You Need?https://twitter.com/arankomatsuzaki/status/1270981237805592576?s=20&t=jEDzZJ2KrCIRUYVKL0vW7A
9. ^華為北大等聯(lián)手打造的Transformer竟在CV領(lǐng)域超過了CNN：多項(xiàng)底層視覺任務(wù)達(dá)到SOTA https://zhuanlan.zhihu.com/p/328534225
10. ^IPT CVPR 2021 | 底層視覺預(yù)訓(xùn)練Transformer | 華為開源代碼解讀?https://zhuanlan.zhihu.com/p/384972190
11. ^如何看待swin transformer成為ICCV2021的 best paper？?https://www.zhihu.com/question/492057377
12. ^Uformer: A General U-Shaped Transformer for Image Restoration?https://arxiv.org/abs/2106.03106
13. ^SwinIR: Image Restoration Using Swin Transformer?https://arxiv.org/abs/2108.10257
14. ^abAttention U-Net: Learning Where to Look for the Pancreas?https://arxiv.org/abs/1804.03999
15. ^Improved Transformer for High-Resolution GANs?https://arxiv.org/abs/2106.07631
16. ^Multi-Stage Progressive Image Restoration?https://arxiv.org/abs/2102.02808
17. ^HINet: Half Instance Normalization Network for Image Restoration?https://arxiv.org/abs/2105.06086
18. ^Rethinking Coarse-to-Fine Approach in Single Image Deblurring?https://arxiv.org/abs/2108.05054
19. ^Global Filter Networks for Image Classification?https://arxiv.org/abs/2107.00645

公眾號后臺回復(fù)“畫圖模板”獲取90+深度學(xué)習(xí)畫圖模板～