清华提出首个退化可感知的展开式Transformer｜NeurIPS 2022

公众号新闻

2022-10-27 05:10

新智元报道

作者：phantom

编辑：QQ

【新智元导读】NeurIPS 2022关于Specral Compressive Imaging (SCI）重建的工作。

本文介绍我们 NeurIPS 2022 关于 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

这个github仓库是一个针对 Snapshot Compressive Imaging 重建的工具包，集成了超过12种深度学习算法。我们之前的工作如 MST, CST, MST++, HDNet 也都在这个github仓库中开源。本文也是我们做的 Transformer in SCI 系列的第三个工作。

简介

单曝光快照压缩成像（Snapshot Compressive Imaging，SCI）的任务是将一个三维的数据立方块如视频（H×W×T）或高光谱图像（H×W×λ）通过预先设计好的光学系统压缩成一个二维的快照估计图（H×W）从而大幅度地降低数据存储和传输的开销。

常见的单曝光快照压缩成像系统有 Coded Aperture Snapshot Spectral Compressive Imaging （CASSI），如下图所示：

图1 单曝光快照压缩成像光学系统

那么在 SCI 中一个至关重要的问题就是如何从被压缩过后的二维快照估计图重建出原始的三维数据，当前主流的方法大都基于深度学习，可以分为两类：端到端（End-to-end）的方法和深度展开式（Deep Unfolding）的方法。

端到端的方法直接采用一个深度学习模型，去拟合一个从 2D 快照压缩估计图到 3D 高光谱数据的映射。这种方法比较暴力，确实可解释性。

深度展开式方法将神经网络嵌入到最大后验概率（Maximum A Posteriori，MAP）模型中来迭代地重建出高光谱图像，能更好地和光学硬件系统适配。因此，本文主要研究深度展开式算法。当前这些方法主要有两大问题：

当前的深度展开式框架大都没有从 CASSI 中估计出信息参数用于引导后续的迭代，而是直接简单地将这些所需要的参数设置为常数或者可学习参数。这就导致后续的迭代学习缺乏蕴含 CASSI 退化模式和病态度信息指导。
当前的Transformer 中全局的 Transformer 计算复杂度与输入的图像尺寸的平方成正比，导致其计算开销非常大。而局部 Transformer 的感受野又受限于位置固定的小窗口当中，一些高度相关的 token 之间无法match。

为了解决上述两个问题，我们提出了首个深度展开式的Transformer。我们贡献可以概括为：

首先，我们推导出了一个能够感知 CASSI 退化模式与病态度的深度展开框架，它从压缩估计图和编码掩膜中估计出信息参数来引导后续的迭代学习。

接着，我们设计了一个能够同时捕获局部和全局依赖关系的 Transformer 并且计算复杂度相较于全局的Transformer而言，大幅减低。

最终，我们将我们设计的 Transformer 嵌入到我们推导的深度展开框架中来极大提升光谱图像重建的效果。我们的算法在使用更低参数量和更少计算量的前提之下，性能大幅度地超过了前人的方法。

方法

CASSI 压缩退化的数学模型

我们定义向量化后的压缩估计图为 y , 被偏移后的输入数据为 x，传感矩阵为 φ ，则 CASSI 的退化数学模型为

其中的 n 表示成像时产生的随机噪声，同样地，它也经过了向量化。

退化可感知的深度展开框架

图2 退化可感知的深度展开式数学框架

我们首先推导出一个 CASSI 退化模式和病态度可感知的深度展开框架，Degradation-Aware Unfolding Framework （DAUF），如图 2 所示。它以最大后验概率为理论基础来进行推导。结合公式（1），我们可以得到 CASSI 的最大后验概率能量优化函数为：

引入辅助变量 z 之后，我们可以得到

为了得到展开式的推导，同时使迭代过程更加简单，能够更快地收敛，我们对公式（3）采用 Half-Quadratic Splitting （HQS）算法进行展开，得到：

我们对公式（4）中的 x 和 z 进行解耦，从而得到两个迭代的子问题如下：

其中的 x 项有一个闭式解：

其中 I 是恒等矩阵，上述闭式解涉及到矩阵求逆，对计算机不友好。为简化矩阵求逆运算，我们做了以下推导：

将公式（7）插入到公式（6），我们可以得到：

请注意，在 CASSI 系统中

是一个对角矩阵，定义

，由此可得：

定义

且

表示

的第 i 个元素，将公式（9）代入公式（8），可得：

返回到迭代公式（5）中，从贝叶斯概率的角度来看，z 项的求解我们可以视为一个去噪问题：

我们可以用一个嵌入的神经网络来隐式地求解公式（11）。至此，我们便完成了 x 项与 z 项的迭代求解过程。

总的来看，我们的 DAUF 可以用下面这个范式来概括：

从左到右依次为参数估计，线性映射，先验去噪。

半交互式 Transformer

网络整体结构

我们半交互式 Transformer （Half-Shuffle Transformer，HST）的整体结构如图3 (a) 所示，采用一个 U 形网络，包含一个 Encoder，Bottleneck，Decoder。其中基本单元是 Half-Shuffle Attention Block （HSAB）。HSAB中最重要的模块是 Half-Shuffle Multi-head Self-Attention （HS-MSA）。

Half-Shuffle Multi-head Self-Attention

HS-MSA 如图3 (d) 所示，将输入的 token

线性映射为：

然后按通道进行二等分：

等分后分别输入到两个分枝，local 分枝和 non-local 分枝，其中的 local 分枝计算 self-attention 如下：

non-local 分枝首先对

进行网格划分，再reshape，从

，然后再计算 self-attention 如下：

对两个branch的计算结果进行融合：

最后将整个 HST 作为一个 stage 的去噪先验网络（不同 stage 的网络权重不共享）插入到 DAUF 当中便得到我们的 DAUHST。

实验

定量实验对比

定量实验对比如表 1 所示，我们的 DAUHST 以更低的计算量和参数量显著超越了之前 16 种 state-of-the-art 方法。我们的方法比先前最好的方法 End-to-end 方法 CST-L 和 Deep Unfolding 方法 BIRNAT 要分别高出 2.24 和 0.78 dB。