清华Ð提出首个二值化光谱重建算法

本文介绍我们在 NeurIPS 2023 上的新工作 《Binarized Spectral Compressive Imaging》

  论文：https://arxiv.org/abs/2305.10299 代码1：github.com/caiyuanhao1998/BiSCI 代码2：github.com/caiyuanhao1998/MST

图1 BiSRNet 与当前最先进的二值网络的性能对比

目前所有的代码，预训练模型和测试结果均已开源在我们开发的一个二值化光谱压缩重建工具包 BiSCI 内，该工具包支持八类最主要的二值网络，欢迎大家来使用。同时，我们还将 BiSRNet 嵌入到了我们之前开发的光谱重建工具箱 MST 当中。目前 MST 工具包已支持超过 12 类深度学习算法，并包含各种配套的可视化函数，欢迎大家来对比。

1. 光谱压缩重建任务介绍

相比于常规的三通道 RGB 图像，高光谱图像包含几十上百个波段，从而捕获了关于成像场景更丰富的信息。也正因为这一重要特性，高光谱图像被广泛地应用于医疗，地形勘探，农业等领域。如图2所示，在医院进行检查时，如果只看常规的RGB图像可能很难诊断病因，但是如果采用高光谱图像捕获并在特定波长下渲染的话，就可以看清楚各类血管，骨骼结构等，从而辅助医生诊断。同样的原理也可应用在遥感地形勘探和农业病虫害检测。

图2 高光谱图像的应用

然而高光谱图像并容易获取，传统的成像设备采用光谱仪对成像场景进行逐波段的扫描，费时费力，难以捕捉运动场景。近些年，科学家们专门设计了单曝光压缩成像（Snapshot Compressive Imaging，SCI）系统来解决这一问题。其光路结构如图3所示。该系统首先通过一个编码孔径掩膜对成像场景的各光谱通道进行调制，然后通过一个三棱镜进行色散后在相机上捕获到一个二维的快照估计图（compressive measurement）。通过这个光路系统，我们便可将三维的光谱立方块压缩成一个二维的图像。而光谱压缩重建的任务便是从这个二维的压缩估计图上恢复出三维的高光谱数据。

图3 单曝光压缩成像系统

当前的主流方法是采用的是全精度模型如 CNN 或者 Transformer 来学一个从压缩估计图到三维光谱立方块的映射。几个比较经典算法有我们之前的工作 MST，MST++，CST，DAUHST，HDNet 等（这些方法全部都开源在我们的工具箱 MST 当中）。这类方法虽然取得了很好的重建效果，但却难以部署到移动端设备（如智能手机、相机、无人机等）上，因为移动端设备的内存空间，计算资源和电力均有限，无法运转全精度模型。另一方面，全精度模型的一些计算单元如深度展开算子和多头自注意力机制等相对复杂，移动端设备无法支持。为了推动光谱重建算法的实际应用，本文做了如下贡献：

（1）提出了光谱压缩重建领域内首个基于二值神经网络（Binarized Neural Network，BNN）的算法 — 二值化光谱重分布网络（Binarized Spectral-Redistribution Network，BiSRNet）。

（2）设计了一个新的二值化卷积单元 — 二值化光谱重分布卷积（Binarized Spectral-Redistribution Convolution，BiSR-Conv）。该卷积单元可以调整光谱表征的强度和分布，同时在反向传播中更好地逼近二值化符号（Sign）函数，从而让求得的梯度更加准确。

（3）制作了四个二值化卷积模块来解决特征图形变过程中的维度不匹配问题，从而让全精度信息能流通整个模型的每一层卷积单元以弥补二值卷积造成的信息损失。

（4）我们的 BiSRNet 显著地超越了当前最先进的二值化算法，甚至取得了能与全精度CNN比肩的效果，然而我们的BiSRNet算法却只需要极低的存储空间（约 0.06 %）和计算代价（约 1 %）。

2. 本文方法

2.1 基础模型

一般来说，用于二值化的全精度模型应该是比较轻量的且它的计算单元可以再移动端设备上运行。然而，现存的CNN 或 Transformer 模型均不满足这一要求。为此，我们重新设计了一个简单，轻量，易于部署的基础模型（Base Model）。

受到之前工作 MST，MST++，CST，DAUHST 的启发，我们设计的基础模型也采用一个 U 形结构，如图4所示。

图4 基础模型的网络结构

在这个基础模型中，采用的所有计算单元都可以被移动端设备支持，同时也不涉及计算复杂度高的操作。

2.2 二值化光谱重分布卷积单元

图 5 符号函数与各类逼近函数的对比图

Clip(x) 与 Quad(x) 的具体表达式如下：

然而分段线性函数知识一个粗略的估计，它与符号函数之间依旧有着很大的误差。图 5 中的阴影部分面积就表示这个误差的大小，Clip(x) 的误差是1。此外，一旦激活或者权重的值落在了 [-1, 1] 之外，他们就不会再被更新。尽管分段二次函数是一个更精确的逼近（误差为2/3），上述的两个问题依旧存在。为此，我们设计了一个可缩放的双曲正切函数来在反向传播中拟合符号函数：

通过我们的设计，全精度信息流就不会被二值卷积给阻断，从而能够顺畅地流过我们设计的 BiSR-Conv 单元，如图 4 (c) 中的红色箭头所示。

图6 本文卷积二值卷积块与普通二值卷积块在处理特征图维度变化时的对比

全精度信息流在二值化算法中非常重要，因为它在一定程度上弥补了量化导致的信息损失。然而在特征图的上下采样过程中，由于维度发生了变化，难以直接引入残差链接以补充全精度信息，如图 6 中每个子图的左半部分所示。蓝色箭头表示二值信号，红色箭头表示全精度信息流。普通的二值卷积模块会在特征度维度改变的同时阻断全精度信息流的传递。然而我们的二值卷积模块通过采用分割合并的技巧，让全精度信息流不被阻断，从而在整个二值模型中流通。

3. 实验结果

3.1 量化指标

表1 BiSRNet 与 SOTA BNN，传统方法，全精度CNN方法的量化指标对比

3.2 视觉对比

图7 BiSRNet 与其他 BNN 方法在仿真数据及上的视觉对比

图8 BiSRNet 与其他 BNN 方法在真实数据集上的视觉对比

4. 总结

本文首次探索了压缩量化在光谱压缩重建领域的应用，提出了该领域首个二值化卷积神经网络 BiSRNet，在量化指标和视觉结果上都显著地超越了当前最先进的二值化模型。代码，预训练权重，重建结果均开源在：

github.com/caiyuanhao1998/BiSCI