引言

想象一下,在炎炎夏日俯视一条长长的柏油路。空气在闪烁,导致景色摇晃、模糊和扭曲。这种现象被称为大气湍流,是由于温度变化影响空气折射率而产生的模糊和几何变形的混沌组合。虽然这种“热浪”在肉眼看来可能具有某种艺术感,但对于用于监控、遥感和天文学的远距离成像系统来说,它简直是一场噩梦。

修复这些图像极其困难。与简单的运动模糊不同,湍流在空间上 (图像的不同部分有不同的扭曲) 和时间上 (每毫秒都在变化) 都是变化的。使用卷积神经网络 (CNN) 的传统方法很难奏效,因为它们的感受野太小,无法捕捉大规模的扭曲。较新的基于 Transformer 的方法虽然可以捕捉全局上下文,但计算复杂度呈二次方增长,这使得它们对于高清视频来说太慢且太耗内存。

在这篇文章中,我们将深入探讨 MambaTM , 这是普渡大学研究人员提出的一个新颖框架。该论文提出了一个双管齐下的解决方案: 一种基于选择性状态空间模型 (Mamba) 的新型深度学习架构,以实现线性计算复杂度;以及一种潜在相位失真 (LPD) 表示法,它能有效地“学习”湍流物理特性以指导恢复过程。

挑战: 为何湍流难以修复

要理解 MambaTM 的必要性,我们首先必须了解现有技术的局限性。大气湍流引入了两种不同类型的退化:

  1. 像素位移 (倾斜) : 图像发生翘曲和摆动。
  2. 模糊: 高频细节丢失。

在数学上,这通常使用泽尼克 (Zernike) 多项式来建模——这是一系列用于描述波前像差的数学函数。模拟器利用一张清晰图像和一组泽尼克系数 (描述扭曲的随机数) 来生成湍流图像。

\[ I \stackrel { \mathrm { d e f } } { = } g ( J ; \mathbf { a } ) = \sum _ { k = 1 } ^ { 1 0 0 } \psi _ { k } \circledast ( \boldsymbol { \beta } _ { k } \cdot \mathcal { W } ( J ; \mathcal { T } ) ) + n , \]

方程 1: 湍流退化的数学模型。

如上式所示,退化图像 \(I\) 是清晰图像 \(J\) 和湍流参数 \(\mathbf{a}\) 的函数。挑战在于逆问题: 即已知 \(I\),求 \(J\)。

现有的深度学习方法试图逆转这一过程,但它们面临着速度、内存和性能的“三难困境”。循环神经网络 (RNN) 速度快但不稳定。Transformer 提供高质量结果,但随着视频长度增加,速度呈指数级下降。这为 MambaTM 的登场奠定了基础,它的目标是将 Transformer 的高质量与简单网络的速度结合起来。

创新点 1: 学习潜在相位失真 (LPD)

这篇论文最重要的贡献之一不仅仅是恢复网络,还在于研究人员如何对问题的物理特性进行建模。

泽尼克系数的问题

标准的基于物理的模拟器依赖泽尼克系数来表示相位失真。一种简单的方法是训练一个网络从模糊图像中估计这些系数,然后使用它们来逆转破坏。

然而,研究人员发现了一个致命缺陷: 病态性 (Ill-posedness) 。 不同的泽尼克系数组合可能导致几乎相同的视觉退化效果。如果神经网络试图预测确切的系数,它很难收敛,因为对于给定的模糊斑块没有唯一的“正确”答案。此外,使用这些系数模拟湍流涉及大核卷积,计算成本高昂,会拖慢训练速度。

解决方案: 潜在相位失真 (LPD)

作者建议学习一种压缩的、“潜在的"失真表示,而不是直接预测泽尼克系数。他们使用变分自编码器 (VAE) 来实现这一点。

图 2: LPD 和 ReBlurNet 的学习方案。

如图 2 所示,该过程涉及两个步骤:

  1. 泽尼克编码器 (VAE) : 网络接收基于物理的参数 (泽尼克系数和核大小) 并将它们压缩成潜在图 (LPD) ,由均值 \(\mu\) 和方差 \(\sigma^2\) 定义。这迫使网络学习失真的分布,而不是生硬的数值。
  2. ReBlurNet: 一个解码器网络,它利用清晰图像和这个新的 LPD 图来重建湍流图像。

这种方法将病态问题转化为适定问题。LPD 图捕捉的是湍流的效果 , 而不是模棱两可的系数。至关重要的是,这个 LPD 模拟器比标准物理模拟器快 50 倍 , 并且是完全可微的,允许它直接集成到恢复网络的训练循环中。

为了确保学习到的表示遵循易于处理的分布 (高斯分布) ,训练中包含了 KL 散度损失:

\[ \mathcal { L } _ { K L } = - \frac { 0 . 5 } { H \times W } \sum _ { i , j } ( \log ( \pmb { \sigma } _ { i , j } ^ { 2 } ) + 1 - \pmb { \mu } _ { i , j } - \pmb { \sigma } _ { i , j } ) \]

方程 2: KL 散度损失确保潜在空间表现良好。

创新点 2: MambaTM 架构

在通过 LPD 有效建模退化物理特性后,研究人员着手解决恢复架构的问题。视频恢复需要分析长序列的帧,以区分移动物体和湍流的抖动运动。

为什么要用 Mamba?

该论文采用了由 Mamba 架构推广的选择性状态空间模型 (SSM) 。 与计算每对像素之间注意力的 Transformer (二次复杂度 \(O(N^2)\)) 不同,SSM 使用循环状态顺序处理数据 (线性复杂度 \(O(N)\)) 。

SSM 的核心机制由以下离散化演化方程描述:

\[ \pmb { h } _ { t } = \bar { \pmb { A } } \pmb { h } _ { t - 1 } + \bar { \pmb { B } } \pmb { x } _ { t } , \quad \ b { y } _ { t } = \pmb { C } \pmb { h } _ { t } + \pmb { D } \pmb { x } _ { t } \]

方程 7: 离散化状态空间模型方程。

在这里,隐藏状态 \(h_t\) 随时间根据输入 \(x_t\) 演化。Mamba 的“选择性”部分意味着参数 \(A, B,\) 和 \(C\) 不是静态的;它们根据输入发生变化,允许模型有选择地记住或忘记信息。这对于湍流处理至关重要,因为模型需要“记住”静态背景,同时“忘记”随机抖动。

网络结构

MambaTM 网络是一种多尺度架构,旨在高效处理视频帧。

图 1: MambaTM 网络架构。

图 1 详细展示了完整的流程:

  1. 多尺度编码器: 视频在不同分辨率下被处理,以捕捉精细细节和宏观结构。
  2. Mamba 组: 核心处理单元。
  3. LPD 引导: 这是连接两项创新的桥梁。网络从输入视频中估计 LPD (失真图) 。然后将估计的 LPD 注入回 Mamba 块中以“引导”恢复过程。

用 1D 扫描解决 3D 视频问题

Mamba 本质上是一个 1D 序列模型 (类似于处理文本) 。而视频是 3D 的 (时间 \(\times\) 高度 \(\times\) 宽度) 。为了弥补这一差距,作者利用三种不同的扫描机制来展平视频数据,同时不丢失时空关系:

  • 空间优先扫描 (SFMB) : 主要沿空间维度扫描,保留局部图像特征。
  • 时间优先扫描 (TFMB) : 沿时间轴扫描,对于分析湍流随时间的变化至关重要。
  • 局部希尔伯特扫描 (LHMB) : 使用希尔伯特曲线 (一种空间填充曲线) 扫描像素。这是一个巧妙的补充,因为标准光栅扫描 (逐行) 会分离垂直相邻的像素。希尔伯特扫描在展平为 1D 时能更好地保留局部邻域位置。

论文中的表 6 (如下所示) 突出了消融研究的结果,证明了结合这些扫描顺序比单独使用任何一种都能产生最佳性能。

表 6: 消融研究显示了结合不同扫描顺序和 LPD 引导的有效性。

联合训练与损失函数

训练策略是一个联合优化问题。网络试图同时做两件事:

  1. 恢复: 生成清晰图像 \(\hat{J}\)。
  2. 重退化: 估计 LPD 并利用它从清晰的估计中重新创建湍流图像。

这种“重退化”循环充当了自监督的一致性检查。如果网络正确理解了物理特性,它应该能够模拟它刚刚消除的湍流。

总损失函数结合了恢复损失 (像素级 + 感知) 和重退化损失:

\[ \mathcal { L } = \mathcal { L } _ { r e s t o r e } + \alpha \mathcal { L } _ { r e t u r b } \]

方程 6: 结合恢复和重退化目标的总损失函数。

恢复损失确保输出看起来像真实图像 (Ground Truth) :

\[ \mathcal { L } _ { r e s t o r e } ( \hat { \pmb { J } } , \pmb { J } ) = \mathcal { L } _ { c } ( \hat { \pmb { J } } , \pmb { J } ) + \alpha _ { p } \mathcal { L } _ { p } ( \hat { \pmb { J } } , \pmb { J } ) \]

方程 4: 利用 Charbonnier 和感知损失的恢复损失。

而重退化损失确保估计的物理参数是准确的:

\[ \mathcal { L } _ { r e t u r b } = \mathcal { L } _ { c } ( \hat { I } _ { t i l t } , I _ { t i l t } ) + \mathcal { L } _ { c } ( \hat { I } _ { t u r b } , I ) + \alpha _ { k } \mathcal { L } _ { K L } \]

方程 5: 重退化损失。

实验与结果

研究人员在合成数据集和真实世界数据集上将 MambaTM 与最先进的方法进行了比较。结果表明其在三个关键领域具有优势: 重建质量、速度和泛化能力。

定量性能

在湍流抑制中,通常使用 PSNR (峰值信噪比) 和 SSIM (结构相似性) 等标准指标。

表 3: 动态场景数据集上的性能显示了 MambaTM 的速度优势。

如表 3 所示,MambaTM 取得了比 TMT 和 DATUM 等竞争方法更高的 PSNR 分数。然而,最引人注目的列是 FPS (每秒帧数) 。 MambaTM 的运行速度达到 55.4 FPS , 几乎是 DATUM (32.7 FPS) 的两倍,并且比像 TMT 这样的 Transformer 方法 (1.50 FPS) 快了几个数量级。这使得 MambaTM 成为第一个可用于实时高分辨率湍流抑制的可行候选者。

我们还可以更仔细地看看计算成本:

表 9: 计算成本比较。

表 9 证实 MambaTM 的延迟 (0.030s) 明显低于其他方法,验证了线性复杂度状态空间模型的效率。

定性结果: 文本与物体

数字很有用,但视觉结果更能说明问题。在文本识别任务 (湍流去除的常见基准) 中,MambaTM 在其他模型失败的地方恢复了可读性。

图 3: 文本恢复的定性比较。

在图 3 中,请注意底行。输入几乎无法阅读。虽然 TMT 和 DATUM 等方法对其进行了改进,但 MambaTM 生成了最清晰、最独特的字符,且伪影 (如颜色噪声或振铃效应) 最少。

该方法在涉及移动物体的动态场景中也表现出色,这通常是众所周知的难题,因为网络必须区分物体运动和湍流“运动”。

图 9: 真实世界动态场景的比较。

在图 9 中,对比蓝色汽车的裁剪图 (H) ,MambaTM (d) 恢复的车辆和行人的边缘明显比 DATUM (c) 更好,后者留下了残留的模糊。

理解学习到的 LPD

为了验证 LPD 是否真的学到了有意义的物理特性,作者对潜在图进行了可视化。

图 7: LPD 图与泽尼克系数的可视化对比。

在图 7 中,我们在顶部看到基于泽尼克的模拟,在底部看到基于 LPD 的模拟。视觉相似性令人震惊。热图 (d, e, f) 显示 LPD 有效地捕捉了湍流的空间强度 (模糊最强烈的地方) ,验证了使用 VAE 压缩退化物理特性的有效性。

结论与启示

MambaTM 论文代表了计算成像领域向前迈出的重要一步。它成功地将深度学习与物理原理以一种既准确又高效的方式结合在一起。

主要收获:

  1. 效率至关重要: 通过使用选择性状态空间模型 (Mamba) ,作者实现了线性复杂度,解锁了视频恢复的实时性能。
  2. 物理感知学习: 潜在相位失真 (LPD) 表示证明,我们不需要强迫神经网络预测精确的物理参数 (如泽尼克系数) 。学习一种行为类似于物理特性的潜在表示往往更稳定、更有效。
  3. 全局上下文: 通过新颖的扫描机制 (空间、时间、希尔伯特) ,该模型有效地将 3D 视频数据作为 1D 序列处理,而不会丢失至关重要的时空联系。

这项工作为远距离监控、炎热环境中的自动导航以及其他物理退化模型复杂或病态的逆问题领域的实时应用打开了大门。通过摆脱笨重的 Transformer 并转向高效的状态空间模型,MambaTM 为高性能视频恢复设定了新的基准。