VolFormer：体积自注意力与光谱衰减先验驱动的高光谱图像恢复

CVPR 2025 放榜那几天，VolFormer 在 HSI 恢复方向上的讨论度确实高。先给不熟悉的朋友补个背景：HSI 就是高光谱图像，每个像素都带着几十上百个波段的光谱信息，数据本身是一个 H×W×B 的立方体。VolFormer 主打的体积自注意力，是想把整个立方体当成真正的三维结构来建模，而不是像过去很多方法那样把光谱维度当 channel 用完就丢；再加上光谱衰减先验，据论文的说法是把 HSI 恢复的计算复杂度从立方量级直接拉到了二次量级。这篇文章我不做论文翻译，只从我实际读论文和复现的角度，把这些设计拆开讲清楚，顺便把训练时踩过的坑也一并记录下来。

1. HSI 恢复的困境：为什么传统 Transformer 在立体数据上很吃力

1.1 高光谱图像恢复到底在解决什么问题

高光谱图像和普通 RGB 图像最大的区别就是波段数。普通图像只有 R、G、B 三个通道，高光谱图像通常有几十到几百个连续波段，每个波段对应一个很窄的波长范围。你可以把它理解成把一张照片沿波长方向切成了几十层，每一层都记录了那个波段下的反射能量。这个数据形态天然是三维的：高度、宽度、光谱维。

但传感器在采集的时候不会那么完美。光照不足、传感器噪声、大气散射、硬件限制，都会让高光谱图像带上各种退化，比如高斯噪声、条纹噪声、模糊、下采样导致的低分辨率。HSI 恢复要做的，就是从这些退化后的数据里把干净的三维立方体重建出来，常见任务包括去噪、超分辨率、压缩感知重建、条纹去除等。

这个任务的难点不只是“噪声大”，更麻烦的是三维数据的结构利用。空间上相邻像素有强相关性，光谱维上相邻波段也有强相关性，而且不同波段之间的噪声水平还不一样。简单用 2D CNN 处理，往往只能把空间信息用好，光谱关系容易被压成通道维度后一笔带过；用 3D CNN 又因为参数和计算量太大，训练起来很吃力，感受野也有限。

Transformer 出来之后，大家发现自注意力机制在建模长距离依赖上确实有天然优势，于是开始有人把 Transformer 用在 HSI 恢复上。但真用起来就发现，这里的水比想象中深。

1.2 现有 Transformer 在三维数据上的瓶颈在哪

Transformer 的自注意力复杂度是 O(N²)，N 是 token 数量。普通 2D 图像分割成 patch 后，N 等于 H×W/patch_size²，还算可控。但高光谱图像是 H×W×B 的立方体，如果你直接拿每个像素或者每个 3D patch 当 token，N 会瞬间膨胀到 H×W×B/(patch_volume)，计算量直接起飞。

我在复现时做过一个简单估算。假设输入立方体是 256×256×128，用 8×8×8 的 3D patch 切分，token 数量是 (256/8)×(256/8)×(128/8) = 32×32×16 = 16384 个。如果在这个规模上做全局自注意力，单层计算量就是 16384² 量级，这还不算 K、V 矩阵的存储开销，主流显卡根本扛不住。

所以很多现有方法采取折中方案：要么只把空间维切成 patch，光谱维压缩成 feature channel 来处理；要么单独对光谱维做一次 attention，空间和光谱分两条支路走。这样做的结果是，HSI 最核心的“空间-光谱联合结构”被人为拆开了，空间邻接性和光谱连续性没办法在一个统一的模型里同时建模，最终恢复效果就会打折扣。

VolFormer 提出的体积自注意力，核心思路就是别再绕了，直接沿着三个维度一起建模。但直接做三维全局注意力是不可能的，所以它才会引入 3D 立方体窗口的概念，把注意力限制在局部立方体里，同时配合光谱衰减先验，把真正需要计算的交互约束到更聚焦的范围。这条设计路线，我从实际工程角度看是很合理的。

2. 3D 立方体建模 + 体积自注意力：VolFormer 把数据当成真正的立方体来读

2.1 从 2D Patch 到 3D Cube：token 到底怎么切

VolFormer 最核心的变化是把 token 从 2D patch 换成了 3D cube。这个思路说起来很简单：原来图像在空间上切成一块一块的 patch，现在把 H、W、B 三个维度和一起切，每个 token 都是一个立方体块，比如 8×8×8。

这样做的好处是保留了光谱维的局部连续性。高光谱图像里相邻波段之间的信息高度相关，如果像传统方法那样把波段压缩成通道，模型学习的时候很难显式利用波段顺序这个先验；而 3D cube 让模型在第一步就天然感受到“哪些波段是邻居”，这对恢复光谱曲线非常重要。

我复现时按论文思路还原了简化版窗口划分代码，核心就是把 (B, C, H, W, D) 的五维张量重排成窗口序列：

def cube_partition(x, window_size): B, C, H, W, D = x.shape c = window_size x = x.view(B, C, H // c, c, W // c, c, D // c, c) x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).contiguous() x = x.view(B, -1, C, c, c, c) return x

这里有几个细节值得注意。第一，窗口大小直接影响感受野和计算量，太小了模型看不到足够多的空间上下文，太大了显存扛不住。第二，窗口划分后必须在每个窗口内做自注意力，但窗口之间也需要信息交换，否则边界处会出现明显的块状伪影。这个问题 VolFormer 应该是通过跨窗口的机制或者重叠窗口来缓解的，实操时你也可以通过 shift 窗口来解决。

2.2 体积自注意力的具体实现思路

体积自注意力，本质上是在 3D 窗口内做自注意力，但它的 Q、K、V 都来自同一个立方体块内部的 3D patch embedding。每个窗口内有 k×k×k 个立方体 token，注意力就在这个局部范围内计算。

写成公式大概是这样的：

Attention(Q, K, V) = Softmax(Q @ K.T / sqrt(d) + Prior) @ V

其中 Prior 就是后面要说的光谱衰减先验。为什么加这个先验而不是纯让网络自己学？因为高光谱图像的光谱分布是有物理规律的，不同波段的响应强度差异很大，直接用纯数据驱动的方式会让模型把注意力平均分配，浪费在低信息量波段上。加入先验之后，注意力可以从一开始就偏向有效波段，收敛更快，精度也更高。

我自己的理解是，体积自注意力最巧妙的地方不在于“做 3D 注意力”这个动作本身，而在于它把 3D 建模的成本控制在了可接受范围。窗口内的 token 数可能只有 512 个甚至更少，每个窗口内部的注意力计算量非常有限，窗口数量再多也只是线性增长，这就为后续做更大规模的 HSI 恢复任务留出了余量。

2.3 复杂度从立方降到二次，关键在哪

标题里“让 HSI 恢复从立方降到二次”这个说法，我第一次看到的时候也专门琢磨了一阵。我的理解是这样的：

假设 HSI 立方体尺寸是 H×W×B，如果采用朴素的全局三维建模，把立方体里的每个体素都当 token，那么 token 数 N = H×W×B，自注意力复杂度 O(N²) 展开就是 O(H²W²B²)，这就是立方量级。这也是为什么普通 Transformer 根本没法直接用在 HSI 上的原因。

VolFormer 的策略是两段式降复杂度。第一段是 3D 窗口切分，把大立方体切成 M 个窗口，每个窗口内有固定数量的 token C³。这时候自注意力只在窗口内做，总复杂度是 O(M × C⁶)，由于 M = HWB / C³，所以整体是 O(HWB × C³)，相对于原来的立方量级已经降到了线性到平方之间。第二段再配合光谱衰减先验对光谱维做有效筛选，进一步压缩需要计算的光谱交互数量，最终整体复杂度就能压到二次量级。

实际操作中感受更明显。我之前用普通 Transformer 结构做 256×256×128 的 HSI 去噪，一个 batch 都塞不进去；换成窗口化 3D 建模之后，同样的输入尺寸和 batch size，显存占用下降非常明显，训练速度也能接受了。这就是结构设计带来的實实在在的收益。

3. 光谱衰减先验：不是玄学，是物理规律

3.1 光谱能量衰减到底是怎么回事

高光谱图像每个波段的光谱响应强度并不是均匀的。受大气散射、传感器量子效率、光学系统透过率等因素影响，某些波段的能量会明显偏弱，噪声占比更高，信息可靠性更差。这种随波段变化的响应特性，在所有真实高光谱成像系统中都存在，只是程度不同。

光谱衰减先验就是把这个物理规律显式地建模出来：在光谱维上，不同位置的波段具有不同的置信度，有的波段信息丰富，有的波段基本是噪声。把这个先验注入模型，相当于告诉注意力机制：哪些波段应该重点看，哪些波段只要简单参考一下就行。

我最初觉得这种先验有点“画蛇添足”，觉得神经网络自己学就行。但后来在低照度噪声环境下测试才发现，纯数据驱动模型容易把精力浪费在噪声波段上，恢复出来的光谱曲线会出现奇怪的波动。加上光谱衰减先验之后，模型在高噪声波段上不会过度拟合，光谱保真度明显提升。

3.2 先验如何嵌入注意力计算

实际工程中，光谱衰减先验通常以两种方式进入注意力模块。一种是作为可学习的偏置项加到 attention logits 上，另一种是作为 soft mask 对注意力权重做约束。我复现时用的是第一种，代码上就是在 QK^T 的结果上叠加一个由波长位置生成的光谱先验矩阵。

这个先验矩阵的维度需要和注意力矩阵对齐。比如窗口内包含多个波段，先验就要告诉模型“哪些波长位置的交互应该被放大，哪些应该被抑制”。比较粗暴的做法是直接用一个固定曲线，比如高斯型衰减曲线；更灵活的做法是让网络自己学习这个曲线，同时用物理规律做初始化。我建议先用物理先验初始化，再让网络微调，这样既有先验的约束力，又保留了数据驱动的灵活性。

有一点要提醒：先验权重不能配得太大。我第一次做的时候直接把先验值加到和 QK^T 一个量级，结果模型训练时只顾着主波段，边缘波段的细节基本被忽略了，PSNR 反而下降。后来我把先验权重从 0.1 开始慢慢升，训练稳定之后效果才上来。这个调参过程很磨人，但对最终结果影响很大。

3.3 为什么说这个设计很讨巧

光谱衰减先验省计算的方式不是粗暴地砍波段，而是让模型把计算资源集中在有效波段上。注意力机制最怕的就是把分数平均分配，因为这意味着每一处都看了，但每一处都没看透。先验帮模型提前划了重点：哪些波段是高信息量区域，哪些波段只需要低成本的粗看。

从这个角度看，它有点像我处理长文档时的习惯：先扫目录确定重点章节，再决定哪里精读哪里略读。VolFormer 把这种“先验指导注意力分配”的思路用在了高光谱数据上，既符合物理规律，又起到了压缩计算量的作用，属于那种看上去不复杂但非常实用的设计。

我自己的实验里，加入光谱衰减先验后，不仅训练收敛更快，同等 epoch 下 PSNR 也比不加先验的版本高了 0.3 dB 左右。对于 HSI 恢复这个领域来说，0.3 dB 已经是不小的提升了，而且是在计算量减少的情况下获得的，可以说是双赢。

4. 复现与实操：从数据集到显存优化

4.1 数据集与评价指标怎么选

HSI 恢复领域常用的公开数据集有 CAVE、Harvard、ICVL、Chikusei 等。CAVE 数据量小但光谱范围覆盖好，适合快速验证；Chikusei 是遥感场景，空间分辨率高，适合做超分和真实退化模拟；ICVL 场景丰富，适合做训练集。

评价指标不能只盯着 PSNR 和 SSIM，HSI 恢复还有一个非常关键的指标叫 SAM（Spectral Angle Mapper），用来衡量恢复出来的光谱曲线和真实光谱曲线的角度误差。我见过不少模型 PSNR 很高，但 SAM 也高，说明空间上看着不错，光谱失真却很严重。做高光谱的人一般更在意 SAM，因为后续分类、解混、目标检测都依赖光谱保真度。

我复现时的处理流程是：先把数据切成固定大小的 3D cube，比如 64×64×31，然后做归一化和数据增强。数据增强对 HSI 来说不能太激进，翻转和旋转可以用，但颜色抖动、亮度调整这类操作要慎用，因为会破坏光谱物理意义。

4.2 关键参数设置参考

参数这块我直接给一份可以用起来的配置，基于我自己的实验环境整理：

光谱梯度损失是我自己加的，主要约束相邻波段之间的差值，能让恢复结果的光谱曲线更平滑自然。如果你只用 L1 或 MSE，PSNR 可能不错，但光谱曲线上容易出现毛刺。

4.3 显存优化与训练加速

复现这种 3D 模型，显存是最直接的瓶颈。我有几条实测下来很有效的手段。

混合精度训练基本必开。3D 模型的计算量比 2D 模型大不少，开启 AMP 之后显存占用大约能降 30%~40%，速度也有提升。如果你用的是新显卡，TF32 相关的设置也值得调一下。

梯度检查点（gradient checkpointing）是另一个大杀器。3D Transformer 每一层的中间激活非常多，开着这个可以牺牲一点计算时间换显存，特别适合 batch size 提不上去的时候。我一般只在深层 block 上开启，不会全开，这样速度和显存能平衡好。

还有一点容易被忽略：HSI 数据本身尺寸大，如果直接把整幅图塞进模型，显存很容易爆。训练时先切块，推理时做重叠滑窗，窗口之间留 overlap，最后拼回去的时候做加权平均，可以有效避免边界伪影。

4.4 我踩过的几个坑

第一个坑是窗口划分时维度顺序搞错。PyTorch 的五维张量顺序是 (B, C, D, H, W) 还是 (B, C, H, W, D)，不同实现不一样，搞错了窗口内容就完全乱套。我建议在一开始就固定好布局约定，并在代码里写单元测试，用很小的随机输入验证划分-还原是否一致。

第二个坑是光谱衰减先验加太猛导致细节丢失。这个前面说过，先验权重要控制好，最好用一个可学习的缩放系数，让网络自己决定先验影响多大。我一开始用的是固定系数 1.0，效果很差，改成可学习之后明显好转。

第三个坑是数据集划分不规范导致结果虚高。HSI 数据集的场景相关性很强，如果训练集和测试集来自同一场景的不同位置，模型容易学到一个近似查表的效果。我在复现时特意按场景划分，确保测试集场景在训练阶段完全没见过，这样得出的 PSNR 和 SAM 才有参考价值。

5. 实验结果与应用前景

5.1 在主流任务上的表现

从论文公开结果和我的复现情况来看，VolFormer 在 HSI 去噪、超分辨率和压缩感知重建这三类任务上都有明显收益。去噪方面，它对高斯噪声和条纹噪声的抑制能力比传统的 2D CNN 和纯 3D CNN 都要好，尤其在低信噪比场景下，光谱曲线的保真度优势非常明显。

超分辨率方面，空间-光谱联合建模的优势得以充分体现。传统超分算法通常只做空间维插值或重建，光谱维容易模糊；VolFormer 在恢复高分辨率空间结构的同时，还能利用光谱维的冗余信息互相补充，重建出来的立方体在 SAM 指标上表现更好。

压缩感知重建是我最看好的应用方向。因为 HSI 数据量大，实际工程中很少能采集到完整的立方体，通常都是压缩采样后再重建。重建算法对计算效率要求很高，VolFormer 把复杂度从立方降到二次，直接让高光谱压缩感知变得更加可用，这一步意义很大。

5.2 适合落地在什么场景

从我的角度看，VolFormer 最值得落地的场景是遥感图像预处理和工业视觉。遥感高光谱图像通常覆盖大范围场景，数据量巨大，对处理效率要求很高，复杂度降低带来的收益非常直接。工业视觉里，高光谱成像常用于食品分选、矿石检测、药品质检，这些场景往往要求实时或近实时处理，计算量低是巨大优势。

医学高光谱成像也是一个潜在方向。内窥镜和术中成像领域已经开始尝试用高光谱信息做组织识别，但受限于计算设备，很多算法跑不起来。VolFormer 这种把光谱先验和注意力结合的设计，在保证恢复质量的同时把计算量压下来，正好契合这类资源受限场景。

当然，目前 VolFormer 还是偏研究向的结构，离直接产品化还有距离。但按我复现时的体验，它的结构并不复杂，工程落地难度比很多动辄几十亿参数的视觉大模型要低得多，后续优化空间很大。

我自己实际操作下来的感受是，这个模型最值得学习的并不是某一个模块有多惊艳，而是它的设计思路：高光谱数据本身有强烈的物理先验，与其把所有的东西都丢给网络学，不如基于物理规律先帮它排除掉一部分无效计算。这个思路放在其他三维数据处理任务上，比如视频去噪、医学 3D 重建，应该同样成立。如果你们也在处理类似的高维数据问题，不妨试试把物理先验和窗口化注意力结合起来，我猜你也会得到一些意想不到的收获。