技术摘要：像素级残差扩散 Transformer (PRDiT)

问题陈述

生成高分辨率 3D 医学体积数据（特别是 CT 扫描）面临着显著挑战，这是由于体素数据随分辨率呈立方级增长所致。现有的生成模型在局部细节保真度与全局结构连贯性之间面临权衡。

• 基于 GAN 的方法（如 HA-GAN、3D-StyleGAN）通常受到模式崩溃（mode collapse）、训练不稳定以及处理高分辨率 3D 体积时极高显存需求的困扰。
• 基于扩散的方法（如 3D-DDPM、潜在扩散模型 LDM、小波扩散模型 WDM）提供了更稳定的训练，但通常依赖于卷积 U-Net 架构。这些架构难以捕捉实现连贯 3D 解剖学合成所需的长程依赖关系和全局上下文。
• 潜在空间方法（使用 VAE 或 VQ-VAE）由于训练样本有限，往往难以学习鲁棒的 3D 医学特征，导致重建质量不佳并造成关键解剖细节的丢失。
• 直接应用 3D Transformer 在计算上是极其昂贵的；分辨率翻倍会导致 Token 数量增加 8 倍，注意力机制成本增加约 64 倍，使得在高分辨率数据上进行端到端训练变得低效且不稳定。

方法论：PRDiT

作者提出了 像素级残差扩散 Transformer (PRDiT)，这是一个可扩展的框架，能够直接在体素层面合成高质量的 3D 医学体积，从而绕过了自动编码器瓶颈。其核心创新在于一种两阶段残差学习策略，将生成过程分解为局部和全局组件。

1. 两阶段架构

• 局部去噪器（第一阶段）： 一个轻量级的、基于 MLP 的盲估计器，在重叠的 3D 补丁（patches）上运行。它独立处理每个补丁，以从补丁内部信息中估计粗略的低频结构和噪声。该阶段利用受时间嵌入调制的自适应层归一化（AdaLN），以根据扩散时间步长调整处理过程。
• 全局残差扩散 Transformer（第二阶段）： 一个具有内存高效注意力机制的 Transformer 模块，用于对残差（高频细节）进行建模。它接收来自局部去噪器的预测结果，并通过利用全局上下文来修正补丁边界处的误差并捕捉长程解剖依赖关系，从而对预测进行精细化处理。
• 训练策略： 首先训练局部去噪器以最小化单补丁损失。随后，训练全局残差 DiT 来预测真实值与局部去噪器输出之间的差异，从而有效地将其容量集中在处理高频精细化而非重新学习全局解剖结构上。

2. 预测器-校正器采样

为了提高样本多样性和稳定性，作者引入了一种结合了确定性步骤和随机性步骤的改进型逆向扩散采样方法：

• 冷预测器（Cold Predictor）： 一个确定性的梯度更新步骤，向前推进 $k$ 个时间步。
• 热校正器（Hot Corrector）： 一个重新引入受控随机性（噪声）以精细化样本的步骤。
  这种“预测器-校正器”方案平衡了确定性引导与自适应噪声注入，在保持全局连贯性的同时增强了对精细结构的保护。

3. 高效的高分辨率缩放

为了解决扩展到更高分辨率（如 $256^3$）时的计算成本问题，PRDiT 采用了迁移学习策略：

• 训练并冻结一个低分辨率模型（如 $128^3$）。
• 对于高分辨率生成，输入通过下采样（使用最近邻下采样以保留噪声能量）进行处理，由冻结的低分辨率模型处理，然后通过上采样（使用三线性插值处理信号，最近邻处理噪声）进行恢复。
• 一个专门的高分辨率残差精细化模块（基于局部去噪器的设计）被训练用于恢复在下采样/上采样过程中丢失的缺失高频细节。这避免了从头开始训练整个架构的需求。

核心贡献

1. 无瓶颈的体素级合成： 一种两阶段扩散 Transformer 框架，直接在体素层面生成 3D 医学体积，消除了自动编码器瓶颈并保留了关键的解剖细节。
2. 残差学习策略： 将生成任务创新性地分解为用于粗略结构的基于 MLP 的局部去噪器和用于高频残差精细化的全局 Transformer，简化了优化过程并增强了训练稳定性。
3. 预测器-校正器采样： 引入了生成过程中的“热”扩散采样（受控随机性），以平衡确定性引导与噪声注入，提高了样本多样性和保真度。
4. 可扩展的高分辨率训练： 一种高效的缩放策略，可以复用预训练的低分辨率组件，显著降低了高分辨率合成的训练成本和计算开销。
5. 达到最先进水平（SOTA）的性能： 通过包括 3D Fréchet Inception Distance (FID)、最大均值差异 (MMD) 和 Wasserstein 距离在内的多种指标进行了广泛验证，证明了其卓越性能。

实验结果

模型在 LIDC-IDRI 和 RAD-ChestCT 数据集上进行了评估，使用了包括 3D FID、MMD 和 Wasserstein 距离在内的指标。

• 性能对比 SOTA： PRDiT 在包括 HA-GAN、3D-LDM 和 WDM-3D 在内的基准模型中表现出持续优越的性能。
  • 在 LIDC-IDRI ($128^3$) 上，PRDiT-12L 变体实现了 1.41 的 3D FID 和 0.1501 的 MMD，显著低于竞争对手（例如 HA-GAN FID 为 3.26，3D-LDM FID 为 7.62）。
  • 定性分析显示，与 GAN 和其他扩散模型观察到的模糊和伪影相比，PRDiT 生成了更锐利的骨骼边缘、更平滑的器官边界以及更清晰的解剖细节（如支气管壁）。
• 可扩展性： 将 Transformer 深度从 4 层增加到 12 层带来了生成质量的持续提升。
• 效率：
  • 训练： 高分辨率 PRDiT-4L$\uparrow$256 模型在 36 GPU 小时内达到了 2.28 的 FID，而 WDM-3D 需要 120 GPU 小时，且 3D-LDM 出现了显存溢出错误。
  • 推理： PRDiT 各变体（11.5s 至 21.9s）比 3D-LDM (26.1s) 和 WDM-3D (34.6s) 更快，但慢于单次通过的 HA-GAN (0.01s)。
• 消融研究： 去除补丁间的重叠或移除全局 DiT 模块会导致性能显著下降，证实了局部-全局分解以及残差精细化策略的必要性。

重要性与主张

论文声称，PRDiT 通过有效平衡计算效率与高保真生成，代表了 3D 医学图像合成领域的重大进展。通过将粗略局部细节的建模与高频全局残差的建模解耦，该框架克服了以往方法中固有的优化困难和显存限制。作者断言，PRDiT 是一种强大且高效的生成高分辨率 3D 医学体积的方法，对于解剖准确性和结构连贯性至关重要的临床和诊断场景具有广阔的应用前景。该工作强调了残差扩散 Transformer 在无需高昂成本从头训练的情况下，有效地扩展到更高分辨率的潜力。