ECCV 2022｜CST: 首個(gè)嵌入光譜稀疏性的Transformer

近年來(lái)，基于學(xué)習(xí)的方法表現(xiàn)出了良好的性能，并主導(dǎo)了主流研究方向。然而，現(xiàn)有的基于CNN的方法在捕獲長(zhǎng)程相關(guān)性和非局部自相似性方面存在局限性。本工作提出了一種能將光譜表征嵌入到Transformer的重建方法：CST。

本文介紹我們 ECCV 2022 關(guān)于 Snapshot Compressive Imaging 重建的工作：

《Coarse-to-Fine Sparse Transformer for Hyperspectral Image Reconstruction》

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

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

這個(gè)github倉(cāng)庫(kù)是一個(gè)針對(duì) Snapshot Compressive Imaging 重建的工具包，集成了11種深度學(xué)習(xí)算法。

快照壓縮成像系統(tǒng)如圖1所示。左邊為待成像的場(chǎng)景，也即三維光譜型號(hào)（空間維度長(zhǎng)和寬，通道維度是不同波段的光譜）。它通過(guò)預(yù)先設(shè)計(jì)好的光路，首先被編碼孔徑掩膜進(jìn)行調(diào)制，然后被三棱鏡進(jìn)行散射，在探測(cè)器上不同的空間位置進(jìn)行成像，這些像疊加在一起之后便得到一個(gè)二維的快照估計(jì)圖，如右下角 Measurement 所示。如此一來(lái)，原先輸入的三維光譜數(shù)據(jù)（x，y，λ）便被壓縮成了二維數(shù)據(jù)（x，y），從而降低數(shù)據(jù)存儲(chǔ)和傳輸?shù)呢?fù)擔(dān)。

然而壓縮估計(jì)圖并不能直接應(yīng)用，我們還需要將其復(fù)原成高光譜圖像才可以進(jìn)一步分析處理。通常來(lái)說(shuō)，進(jìn)行高光譜重建的方法可以分為四類：

（1）基于先驗(yàn)?zāi)Ｐ偷膫鹘y(tǒng)方法。這一類方法會(huì)預(yù)先手工設(shè)計(jì)一些圖像先驗(yàn)，如 total variation，low rank property 等。然而這類方法的泛化性差，并且每次使用前需要調(diào)整參數(shù)，花費(fèi)大量時(shí)間。

（2）端到端的深度學(xué)習(xí)方法。這類方法直接采用一個(gè)深度學(xué)習(xí)模型，去擬合一個(gè)從 2D 快照壓縮估計(jì)圖到 3D 高光譜數(shù)據(jù)的映射。這類方法目前主要基于卷積神經(jīng)網(wǎng)絡(luò)，沒(méi)有適配光譜表征在空間維度呈現(xiàn)出的稀疏性。并且卷積神經(jīng)網(wǎng)絡(luò)在捕獲非局部依賴關(guān)系（non-local / long-range dependences）上有明顯的短板。

（3）迭代式的深度學(xué)習(xí)方法。這一類方法運(yùn)用迭代公式，將卷積神經(jīng)網(wǎng)絡(luò)嵌入到每一個(gè)迭代中作為去噪網(wǎng)絡(luò)。

（4）即插即用的方法。這類方法將預(yù)先訓(xùn)練好的一個(gè)深度學(xué)習(xí)網(wǎng)絡(luò)插入到每一個(gè)迭代始終，無(wú)需微調(diào)直接應(yīng)用。

本文主要研究第（2）類方法，主要貢獻(xiàn)點(diǎn)可以概括為如下：

• 提出了一種能將光譜表征嵌入到Transformer的重建方法。名字是 Coarse-to-Fine Sparse Transformer (CST)
• 提出一種檢測(cè)密集光譜表征區(qū)域的方法，光譜感知篩選機(jī)制 ，Spectrum-Aware Screening Mechanism (SASM)
• 提出一種基于哈希來(lái)聚合相關(guān)光譜表征的多頭自注意機(jī)制，Spectra-Aggregation Hashing Multi-head Self-Attention (SAH-MSA)。
• 在仿真數(shù)據(jù)集上，我們的 CST 系列模型用了更少的參數(shù)量取得了更高的結(jié)果。在真實(shí)數(shù)據(jù)上，效果更逼真。

我們的 CST 的整體結(jié)構(gòu)如圖2（a）所示。首先將二維快照估計(jì)圖 Y 滑動(dòng)截取為初始的光譜圖像 H 并與偏移的掩膜 M 進(jìn)行 concate。網(wǎng)絡(luò)的第一部分是一個(gè)稀疏度估計(jì)器，它估計(jì)出光譜表征比較密集的區(qū)域，然后輸出一個(gè)

稀疏度圖，第二階段根據(jù)這個(gè)稀疏度圖篩選出光譜密集區(qū)域，集中對(duì)這些區(qū)域進(jìn)行計(jì)算。兩個(gè)階段均采用U-Net。

為了使稀疏度估計(jì)器能夠?qū)崿F(xiàn)我們想要的“篩選出光譜密集區(qū)域”的功能，我們對(duì)它輸出的稀疏度圖進(jìn)行監(jiān)督。我們?cè)O(shè)立的目標(biāo)是重建光譜圖與真值光譜圖之間的差異沿通道的均值，如下式所示：

這很好理解，我們?cè)O(shè)想在高光譜成像的場(chǎng)景中，背景往往很黑暗，這些背景區(qū)域幾乎沒(méi)有信息，很好重建。那么難以重建的部分便是光譜密集的區(qū)域?；诖耍覀?cè)O(shè)立的損失函數(shù)如下：

其中 MsM_sM_s 為估計(jì)的稀疏度圖。那么整體的損失函數(shù)為：

其中 X' 為網(wǎng)絡(luò)重建的光譜圖，X* 為真值。λ 為權(quán)重系數(shù)，平衡兩個(gè)損失函數(shù)。

在進(jìn)行光譜密集區(qū)域篩選時(shí)，我們對(duì)稀疏度圖設(shè)置一個(gè)閾值，高于這一閾值的區(qū)域就被選出。

CST 的第二階段基于第一階段輸出的稀疏度圖來(lái)進(jìn)行光譜圖像重建，其基本組成單元如圖1（b）所示。有兩條支路，頂端支路是一個(gè)恒等連接，底端支路首先基于稀疏度圖的篩選，只對(duì)選出的區(qū)域進(jìn)行計(jì)算，其組件中包含一個(gè)基于哈希聚合相關(guān)光譜表征的多頭自注意機(jī)制（SAH-MSA），如圖1（c）所示。以前的多頭注意力機(jī)制（MSA）總是對(duì)所有的 token 進(jìn)行計(jì)算，有些 token 毫無(wú)關(guān)聯(lián)，這種計(jì)算方式效率低下。我們的 SAH-MSA 正是解決這一問(wèn)題。首先，我們通過(guò)一個(gè)哈希映射函數(shù)，給每一個(gè) token 算出一個(gè)分?jǐn)?shù)，公式如下：

然后，我們根據(jù) token 的分?jǐn)?shù)從高到低排序，劃分成不同的 buckets，如下：

則 SAH-MSA 的輸出為：

其中，每一個(gè) head 的自注意力由下式計(jì)算得到：

同時(shí)每一個(gè) head 中的每一個(gè) bucket 都由 multi-round 機(jī)制計(jì)算得到

其中每一個(gè)round的全總?cè)缦滤?/p>3. 實(shí)驗(yàn)3.1 定量實(shí)驗(yàn)對(duì)比

在仿真數(shù)據(jù)上的定量實(shí)驗(yàn)的結(jié)果如下標(biāo)所示

我們的 CST 系列采用更少的計(jì)算量達(dá)到了更高的性能。特別地，與我們 CVPR 2022的工作 MST 的比較如下圖3所示。

CST 與其他方法在仿真數(shù)據(jù)集和真實(shí)數(shù)據(jù)集上的定性結(jié)果對(duì)比分別如圖4 和圖5 所示。

我們對(duì)稀疏度估計(jì)器輸出的稀疏度圖進(jìn)行可視化，結(jié)果如圖6所示。

ECCV 2022 的這次工作是 Transformer 應(yīng)用于 Snapshot Compressive Imaging 的第二個(gè)工作。它采用與我們CVPR 2022 的工作 MST 不同的技術(shù)方案，實(shí)現(xiàn)了更低計(jì)算量，更高精度的高光譜重建效果。期待未來(lái)能看到更多的人從事 Snapshot Compressive Imaging 的研究。

另附上 MST 與 MST++ 的知乎鏈接：

phantom：[CVPR 2022 & NTIRE 冠軍] 首個(gè)高光譜圖像重建Transformer

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