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

近年來，基于學習的方法表現(xiàn)出了良好的性能，并主導了主流研究方向。然而，現(xiàn)有的基于CNN的方法在捕獲長程相關(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

這個github倉庫是一個針對 Snapshot Compressive Imaging 重建的工具包，集成了11種深度學習算法。

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

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

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

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

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

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

本文主要研究第（2）類方法，主要貢獻點可以概括為如下：

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

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

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

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

這很好理解，我們設(shè)想在高光譜成像的場景中，背景往往很黑暗，這些背景區(qū)域幾乎沒有信息，很好重建。那么難以重建的部分便是光譜密集的區(qū)域。基于此，我們設(shè)立的損失函數(shù)如下：

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

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

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

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

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

則 SAH-MSA 的輸出為：

其中，每一個 head 的自注意力由下式計算得到：

同時每一個 head 中的每一個 bucket 都由 multi-round 機制計算得到

其中每一個round的全總?cè)缦滤?/p>3. 實驗3.1 定量實驗對比

在仿真數(shù)據(jù)上的定量實驗的結(jié)果如下標所示

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

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

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

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

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

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

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