PPDM: Pixel Puzzling Diffusion Model for Speed and Memory Efficient Volumetric Medical Image Translation

技术摘要：像素拼图扩散模型 (Pixel Puzzling Diffusion Model, PPDM)

问题陈述
与传统的基于 CNN 的方法相比，扩散模型在医学图像到图像（I2I）转换方面建立了更高的保真度，而后者往往存在过度平滑和感受野受限的问题。然而，将这些模型扩展到高分辨率 3D 体积医学数据面临着严峻挑战，其计算成本和 GPU 显存需求极高。3D 扩散模型的激活显存和计算量随空间分辨率呈近似立方级（$O(N^3)$）增长，很快就会超过普通 GPU 的容量。

现有的内存优化策略试图缓解这一问题，但引入了显著的权衡：

• 分块与聚合方法（例如切片式或分块处理）通常无法捕捉长程解剖上下文，从而引入边界伪影并降低全局结构一致性。
• 小波方法将图像分解为不同的频率带，这可能导致纹理过于平滑、细节模糊以及解剖边界减弱，尤其是在高频结构中。
• 潜在扩散模型将输入压缩到低维潜在空间。虽然这种方式节省内存，但这种激进的压缩可能会掩盖具有临床意义的细节，且高度依赖于难以在有限医学数据集上训练的鲁棒编码器/解码器。

方法论
作者提出了像素拼图扩散模型 (PPDM)，这是一个旨在降低内存和速度成本，同时不牺牲全局体积一致性的框架。该流水线由三个核心组件组成：

1. 可逆像素拼图–解拼（Unpuzzle）算子：
   该算子通过空间分辨率与通道维度之间的交换进行操作。给定输入体积 $X \in \mathbb{R}^{C \times D \times H \times W}$，“解拼”算子将体素重新排列为低分辨率矩阵 $X_L \in \mathbb{R}^{(C \cdot r^3) \times \frac{D}{r} \times \frac{H}{r} \times \frac{W}{r}}$，其中 $r$ 是空间降采样因子。

   • 机制： 空间坐标 $(x, y, z)$ 被映射到通道索引，从而有效地扩展了通道维度并减少了空间维度。
   • 优势： 这使得扩散网络可以在降低的空间维度上运行，在大幅降低激活显存的同时，通过扩展后的通道维度保留了全局体积上下文。原始结构通过重建过程中的逆“拼图”算子得以恢复。

2. 直接桥接扩散公式：
   不同于从纯高斯噪声开始的标准 DDPM，PPDM 采用了直接桥接公式。该过程从时间 $T$ 时的条件输入图像 $y$ 开始，构建一个向目标 $x_0$ 演化的随机桥接。

   • 公式： $x_t = (1 - m_t)x_0 + m_t y + \sqrt{\delta_t}\epsilon$，其中 $m_t$ 是插值权重。
   • 优势： 通过从输入而非纯噪声开始，缩小了初始状态与目标之间的分布差距。模型仅专注于学习任务相关的残差（修正），从而减少了不确定性和所需的去噪步数。

3. 拼图梯度损失（Puzzle-Gradient Loss）：
   为了解决空间重排可能引入的通道间不一致性和网格状伪影，模型引入了拼图梯度损失。

   • 机制： 在重建的体积上，使用多个方向的 3D Sobel 核计算空间梯度。该损失函数最小化预测输出与地面真值之间梯度幅值的差异。
   • 优势： 这强制执行了拼图体素块之间的空间相干性，抑制了伪影并增强了精细解剖边界的恢复。

核心贡献

• 内存与速度效率： PPDM 引入了一种新机制，可在不牺牲精度的情况下将 GPU 显存使用量降低高达一个数量级（例如，在特定任务中从 ~50 GB 降至 ~7 GB），并加速了推理过程，使高分辨率 3D 扩散在商用硬件上变得可行。
• 保持全局上下文： 与基于补丁（patch-based）或潜在压缩的方法不同，PPDM 通过在内存中保留完整体积（尽管经过了维度重排）而非处理不连续的补丁或高度压缩的潜在代码，维持了全局体积的一致性。
• 伪影抑制： 拼图梯度损失的集成专门针对并减轻了空间重排策略固有的网格状伪影，确保了高保真的解剖重建。
• 可扩展性： 该框架允许通过调整下采样因子 $r$ 在内存使用量和准确性之间进行灵活的权衡。

实验结果
作者在三个具有挑战性的 3D 医学图像转换任务上评估了 PPDM：

1. 低计数 PET 去噪：（从 1% 低计数到标准剂量）。
2. 联合 PET 去噪与衰减校正：（5% 低计数、非衰减校正到全计数、衰减校正）。
3. 跨模态 MRI 转换：（T2 加权到质子密度加权）。

性能表现：

• 准确性： 在所有任务中，PPDM 一致地达到或超过了全 3D 扩散模型（特别是 3D BBDM）的表现。例如，在 PET 去噪任务中，与 3D BBDM 基准（NMSE 0.0229, PSNR 30.762 dB）相比，PPDM 实现了更低的 NMSE (0.0223) 和更高的 PSNR (30.868 dB)。
• 与高效基准模型的对比： PPDM 显著优于现有的内存高效型方法，如 PatchDDM、WDM 和 ALDM。虽然这些方法降低了内存占用，但其准确性出现了大幅下降（例如，在 PET AC 任务中，PatchDDM 的 NMSE 为 0.2518，而 PPDM 仅为 0.0130）。
• 资源利用率： 与 3D BBDM 基准相比，PPDM 将训练 GPU 显存消耗降低了约 7 倍（例如，从 49.07 GB 降至 6.94 GB），并显著减少了推理 TFLOPs。

意义与主张
论文声称 PPDM 为有限计算资源下的高保真 3D 扩散医学图像转换提供了一个实用且可扩展的解决方案。它证明了在不产生传统 3D 扩散模型所带来的高昂内存成本的前提下，实现最先进的重建质量是可能的。

作者强调，PPDM 避免了现有高效方法的缺陷：它不依赖于补丁拼接（这会破坏全局上下文）或激进的潜在压缩（这有丢失临床细节的风险）。相反，它利用可逆的空间-通道变换，在保持完整 3D 体积完整性的同时，实现了高效的训练和推理。这项工作表明，PPDM 有助于扩散技术在医学成像领域的广泛部署，特别是对于需要高分辨率体积数据的应用，如 PET 重建和跨模态 MRI 合成。

作者承认的局限性

• 超参数敏感性： 模型性能对拼图梯度损失权重 ($\lambda$) 和空间下采样因子 ($r$) 较为敏感。最优值可能因具体应用而异。
• 与分辨率的权衡： 增加下采样因子 $r$ 会进一步降低内存使用，但会导致性能出现适度下降（例如，将 $r$ 从 2 增加到 3 会导致 NMSE 轻微上升）。
• 核大小复杂度： 更大的下采样因子需要仔细调整 Sobel 核的大小，以平衡边缘保留与平滑效果，这增加了训练过程的复杂性。