NeurIPS 2022｜清華提出首個(gè)退化可感知的展開式Transformer

來(lái)源丨h(huán)ttps://zhuanlan.zhihu.com/p/576280023編輯丨極市平臺(tái)

本文介紹我們 NeurIPS 2022 關(guān)于 Spectral Compressive Imaging （SCI）重建的工作：

《Degradation-Aware Unfolding Half-Shuffle Transformer for Spectral Compressive Imaging》

文章：https://arxiv.org/abs/2205.10102

代碼：https://github.com/caiyuanhao1998/MST

這個(gè)github倉(cāng)庫(kù)是一個(gè)針對(duì) Snapshot Compressive Imaging 重建的工具包，集成了超過(guò)12種深度學(xué)習(xí)算法。我們之前的工作如 MST, CST, MST++, HDNet 也都在這個(gè)github倉(cāng)庫(kù)中開源。本文也是我們做的 Transformer in SCI 系列的第三個(gè)工作。

MST：https://arxiv.org/abs/2111.07910

CST：https://arxiv.org/abs/2203.04845

MST++：https://arxiv.org/abs/2204.07908

HDNet：https://arxiv.org/abs/2203.02149

單曝光快照壓縮成像（Snapshot Compressive Imaging，SCI）的任務(wù)是將一個(gè)三維的數(shù)據(jù)立方塊如視頻（H×W×T）或高光譜圖像（H×W×λ）通過(guò)預(yù)先設(shè)計(jì)好的光學(xué)系統(tǒng)壓縮成一個(gè)二維的快照估計(jì)圖（H×W）從而大幅度地降低數(shù)據(jù)存儲(chǔ)和傳輸?shù)拈_銷。常見的單曝光快照壓縮成像系統(tǒng)有 Coded Aperture Snapshot Spectral Compressive Imaging （CASSI），如下圖所示

那么在 SCI 中一個(gè)至關(guān)重要的問(wèn)題就是如何從被壓縮過(guò)后的二維快照估計(jì)圖重建出原始的三維數(shù)據(jù)，當(dāng)前主流的方法大都基于深度學(xué)習(xí)，可以分為兩類：端到端（End-to-end）的方法和深度展開式（Deep Unfolding）的方法。端到端的方法直接采用一個(gè)深度學(xué)習(xí)模型，去擬合一個(gè)從 2D 快照壓縮估計(jì)圖到 3D 高光譜數(shù)據(jù)的映射。這種方法比較暴力，確實(shí)可解釋性。深度展開式方法將神經(jīng)網(wǎng)絡(luò)嵌入到最大后驗(yàn)概率（Maximum A Posteriori，MAP）模型中來(lái)迭代地重建出高光譜圖像，能更好地和光學(xué)硬件系統(tǒng)適配。因此，本文主要研究深度展開式算法。當(dāng)前這些方法主要有兩大問(wèn)題：

• 當(dāng)前的深度展開式框架大都沒有從 CASSI 中估計(jì)出信息參數(shù)用于引導(dǎo)后續(xù)的迭代，而是直接簡(jiǎn)單地將這些所需要的參數(shù)設(shè)置為常數(shù)或者可學(xué)習(xí)參數(shù)。這就導(dǎo)致后續(xù)的迭代學(xué)習(xí)缺乏蘊(yùn)含 CASSI 退化模式和病態(tài)度信息指導(dǎo)。
• 當(dāng)前的 Transformer 中全局的 Transformer 計(jì)算復(fù)雜度與輸入的圖像尺寸的平方成正比，導(dǎo)致其計(jì)算開銷非常大。而局部 Transformer 的感受野又受限于位置固定的小窗口當(dāng)中，一些高度相關(guān)的 token 之間無(wú)法match。

為了解決上述兩個(gè)問(wèn)題，我們提出了首個(gè)深度展開式的Transformer。我們貢獻(xiàn)可以概括為：

• 首先，我們推導(dǎo)出了一個(gè)能夠感知 CASSI 退化模式與病態(tài)度的深度展開框架，它從壓縮估計(jì)圖和編碼掩膜中估計(jì)出信息參數(shù)來(lái)引導(dǎo)后續(xù)的迭代學(xué)習(xí)。
• 接著，我們?cè)O(shè)計(jì)了一個(gè)能夠同時(shí)捕獲局部和全局依賴關(guān)系的 Transformer 并且計(jì)算復(fù)雜度相較于全局的Transformer而言，大幅減低。
• 最終，我們將我們?cè)O(shè)計(jì)的 Transformer 嵌入到我們推導(dǎo)的深度展開框架中來(lái)極大提升光譜圖像重建的效果。我們的算法在使用更低參數(shù)量和更少計(jì)算量的前提之下，性能大幅度地超過(guò)了前人的方法。

我們定義向量化后的壓縮估計(jì)圖為 y , 被偏移后的輸入數(shù)據(jù)為 x， 傳感矩陣為 φ ，則 CASSI 的退化數(shù)學(xué)模型為

其中的 n 表示成像時(shí)產(chǎn)生的隨機(jī)噪聲，同樣地，它也經(jīng)過(guò)了向量化。

我們首先推導(dǎo)出一個(gè) CASSI 退化模式和病態(tài)度可感知的深度展開框架，Degradation-Aware Unfolding Framework （DAUF），如圖 2 所示。它以最大后驗(yàn)概率為理論基礎(chǔ)來(lái)進(jìn)行推導(dǎo)。結(jié)合公式（1），我們可以得到 CASSI 的最大后驗(yàn)概率能量?jī)?yōu)化函數(shù)為：

引入輔助變量 z 之后，我們可以得到

為了得到展開式的推導(dǎo)，同時(shí)使迭代過(guò)程更加簡(jiǎn)單，能夠更快地收斂，我們對(duì)公式（3）采用 Half-Quadratic Splitting （HQS）算法進(jìn)行展開，得到：

我們對(duì)公式（4）中的 x 和 z 進(jìn)行解耦，從而得到兩個(gè)迭代的子問(wèn)題如下：

其中的 x 項(xiàng)有一個(gè)閉式解：

其中 I 是恒等矩陣，上述閉式解涉及到矩陣求逆，對(duì)計(jì)算機(jī)不友好。為簡(jiǎn)化矩陣求逆運(yùn)算，我們做了以下推導(dǎo)：

將公式（7）插入到公式（6），我們可以得到：

請(qǐng)注意, 在 CASSI 系統(tǒng)中  是一個(gè)對(duì)角矩陣, 定義 , 由此可得:

定義  且  表示  的第  個(gè)元素, 將公式 (9) 代入公式 (8), 可得：

返回到迭代公式（5）中，從貝葉斯概率的角度來(lái)看，z 項(xiàng)的求解我們可以視為一個(gè)去噪問(wèn)題：

我們可以用一個(gè)嵌入的神經(jīng)網(wǎng)絡(luò)來(lái)隱式地求解公式（11）。至此，我們便完成了 x 項(xiàng)與 z 項(xiàng)的迭代求解過(guò)程。

總的來(lái)看，我們的 DAUF 可以用下面這個(gè)范式來(lái)概括：

從左到右依次為參數(shù)估計(jì)，線性映射，先驗(yàn)去噪。

我們半交互式 Transformer （Half-Shuffle Transformer，HST）的整體結(jié)構(gòu)如圖3 (a) 所示，采用一個(gè) U 形網(wǎng)絡(luò)，包含 一個(gè) Encoder，Bottleneck，Decoder。其中基本單元是 Half-Shuffle Attention Block （HSAB）。HSAB中最重要的模塊是 Half-Shuffle Multi-head Self-Attention （HS-MSA）。

HS-MSA 如圖3 (d) 所示，將輸入的 token  線性映射為：

然后按通道進(jìn)行二等分：

等分后分別輸入到兩個(gè)分枝，local 分枝和 non-local 分枝，其中的 local 分枝計(jì)算 self-attention 如下：

non-local 分枝首先對(duì)  進(jìn)行網(wǎng)格劃分, 再reshape, 從  到 , 然后再計(jì)算 self-attention 如下:

對(duì)兩個(gè)branch的計(jì)算結(jié)果進(jìn)行融合：

最后將整個(gè) HST 作為一個(gè) stage 的去噪先驗(yàn)網(wǎng)絡(luò)（不同 stage 的網(wǎng)絡(luò)權(quán)重不共享）插入到 DAUF 當(dāng)中便得到我們的 DAUHST。

定量實(shí)驗(yàn)對(duì)比如表 1 所示，我們的 DAUHST 以更低的計(jì)算量和參數(shù)量顯著超越了之前 16 種 state-of-the-art 方法。我們的方法比先前最好的方法 End-to-end 方法 CST-L 和 Deep Unfolding 方法 BIRNAT 要分別高出 2.24 和 0.78 dB。

相較于先前的 Deep Unfolding 方法，我們繪制了 PSNR - FLOPS 坐標(biāo)圖比較 DAUHST 和其他 Deep Unfolding 方法的 性價(jià)比。如圖4所示。我們的方法在消耗相同計(jì)算量的情況下比先前方法要高出 4 dB。

在仿真數(shù)據(jù)集上的定性結(jié)果對(duì)比如圖5所示。左上角是RGB圖像和快照估計(jì)圖（Measurement）。下方四行圖像是不同方法重建的四個(gè)波長(zhǎng)下的高光譜圖像。右上角的圖像是下方圖像中黃色框框內(nèi)的放大圖。從重建的高光譜圖像來(lái)看，我們的方法能更好地恢復(fù)出細(xì)節(jié)內(nèi)容和紋理結(jié)構(gòu)，請(qǐng)注意對(duì)比小立方塊區(qū)域。a 和 b 曲線對(duì)應(yīng)著 RGB 圖像的兩個(gè)綠色框的區(qū)域的光譜強(qiáng)度曲線，可以看出，我們的 DAUHST 與 Ground Truth 的曲線最為接近。

圖6 展示的是各類方法在真實(shí)數(shù)據(jù)集上的對(duì)比?？梢钥闯鲋挥形覀兊姆椒軌蛟诟鞣N波長(zhǎng)的光譜上穩(wěn)定地重建出小花并同時(shí)抑制噪聲的生成。

本文是我們 SCI 系列代表作的第五個(gè)，也是 NeurIPS 上邊首次有 SCI 重建的工作。SCI 重建作為新興的 low-level 方法這兩年迅猛發(fā)展，希望能夠看到有更多的人能夠加入的這個(gè) topic 的研究，畢竟新的領(lǐng)域有更多出成果的機(jī)會(huì)。另附上我們先前在 CVPR 2022 和 ECCV 2022 上的兩個(gè)工作 MST 和 CST 的知乎解讀鏈接：

https://zhuanlan.zhihu.com/p/501101943

https://zhuanlan.zhihu.com/p/544979161

本文僅做學(xué)術(shù)分享，如有侵權(quán)，請(qǐng)聯(lián)系刪文。