PA-Diff：物理知识导向的扩散 Transformer，重塑水下视觉新高度

TL;DR

本文介绍了一种名为 PA-Diff 的新型水下图像增强（Underwater Image Enhancement, UIE）框架。它创新性地将水达成像物理模型作为先验，引导**扩散模型（Diffusion Model）**的推理过程。通过 PPG、INR 与 PDT 三个分支的协同，该方法不仅解决了传统模型鲁棒性差的问题，还突破了纯学习方法缺乏物理可解释性的瓶颈，在多个数据集上达成 SOTA 性能。

核心速览：为什么要让扩散模型“学物理”？

水下成像环境极其复杂，吸收、散射和非均匀光照导致了严重的偏色和雾状模糊。

• 现状：基于扩散模型的方法虽然生成能力强，但常因忽略物理规律导致生成结果出现伪影或不自然的颜色。
• PA-Diff 的直觉：如果能让模型在去噪的每一步都知道“传输深度”和“背景光分量”，它就能更精准地定位哪些区域需要补全，哪些颜色需要校正。

痛点深挖：现有方法的“深水区”

传统物理模型（如暗通道先验）在复杂水域下往往失效；而现有的深度学习方法（Physics-free methods）通常是在原始像素空间硬磨，无法充分利用水下图像特有的物理属性。此外，扩散模型在反向去噪过程中缺乏对图像退化难易程度的感知。

物理感知的架构设计 (Methodology)

PA-Diff 的架构分为三个核心分支，形成了一个从“物理感知”到“隐式表征”再到“增强生成”的闭环：

1. 物理先验生成 (Physics Prior Generation, PPG)

利用修改后的 Koschmieder 模型，PPG 分支估计出传输图 (Transmission Map, $T^c$) 和 背景背景光 (Background Light, $B^c$)。这些物理参数直接作为 Diffusion 的条件输入（Conditional Guidance）。

2. 隐式神经网络重构 (Implicit Neural Reconstruction, INR)

为了增强对不同水下场景的鲁棒性，作者引入了 INR。它将图像编码为连续函数，这有助于减少特征空间的噪声，并为后续的扩散分支提供更稳定的表征，降低了还原难度。

3. 物理感知扩散 Transformer (PDT)

这是系统的核心，作者设计了一个 Physics-aware Diffusion Transformer Block，包含以下黑科技：

• PA-SA (Physics-aware Self-attention)：通过 Query 驱动的注意力机制，将生成的物理先验与 Transformer 特征进行深度融合。
• PPU (Physics Perception Unit)：这是一种即插即用的模块，它在特征空间重构了物理成像公式：$ildeJ=ilde{t}_1\odot ildeI+ilde{t}_2\odot ildeB$。

图 1：PA-Diff 整体框架，展示了 PPG, INR 和 PDT 三个分支的交互流程。

实验与结果分析

SOTA 对比

在 UIEB 和 LSUI 大规模数据集上，PA-Diff 在 PSNR、SSIM 以及感知评估指标 LPIPS 和 FID 上均表现出显著优势。

图 2：与前人扩散模型工作（如 DM-water）的视觉对比。PA-Diff 在真实场景下的颜色校正和细节恢复更具一致性。

PPU 模块的即插即用价值

消融实验证明，PPU 模块可以被无缝集成到其他经典模型中（如 Water-Net, U-color）。实验显示，集成 PPU 后，WaterNet 的 PSNR 从 18.15 提升到了 20.54，性能增幅非常可观。

深度洞察：为什么 PPU 如此有效？

PPU 的本质是在神经特征空间中对“物理反演过程”进行微分逼近。它不像传统公式那样严苛（容易受噪声干扰），而是通过深度学习的非线性拟合能力，学习更稳健的物理映射。这种“物理原型 + 神经近似”的思路，既保留了物理模型的约束，又发挥了深度神经网络的灵活性。

总结与局限性

Takeaway: PA-Diff 证明了将物理先验嵌入扩散模型的每一个 Transformer 块（而不仅仅是前处理）是实现高质量图像恢复的关键。

局限性：尽管 PA-Diff 取得了突破，但在物理先验生成分支上仍存在模型理想化的问题（与真实 Ground Truth 仍有 Gap），且在极端弱光、高浑浊度的极端深海环境下，细节丢失的问题依然具有挑战性。

图 3：PPG 分支重构图与原图的差异展示，揭示了未来进一步优化的空间。