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

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

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

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

（1）基于先驗模型的傳統(tǒng)方法。這一類方法會預先手工設計一些圖像先驗，如 total variation，low rank property 等。然而這類方法的泛化性差，并且每次使用前需要調整參數(shù)，花費大量時間。

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

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

（4）即插即用的方法。這類方法將預先訓練好的一個深度學習網絡插入到每一個迭代始終，無需微調直接應用。

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

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

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

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

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

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

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

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

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

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

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

則 SAH-MSA 的輸出為：

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

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

其中每一個round的全總如下所示

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

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

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

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

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

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

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

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