扩散模型在医学图像到图像翻译中展现出卓越的保真度,但其扩展到高分辨率 3D 体积时受到高昂计算成本和 GPU 内存需求的严重限制。现有的内存高效策略往往牺牲全局体积一致性或精细解剖细节。在这项工作中,我们提出了像素拼图扩散模型(PPDM),一个简单而有效的框架,用于内存和速度高效的 3D 医学图像翻译。PPDM 引入了一个可逆的像素拼图-解拼图操作,将空间分辨率转换为通道维度,大幅减少激活内存同时保留全局上下文。为了进一步提高效率和稳定性,我们采用直接桥接扩散公式,从条件输入而非纯噪声开始,使模型专注于任务相关残差。此外,引入拼图梯度损失以增强空间一致性并抑制空间重排引入的网格状伪影。我们在多个具有挑战性的 3D 医学图像翻译任务上评估了 PPDM,包括低计数 PET 去噪、联合 PET 去噪与衰减校正以及跨模态 MRI 翻译。在所有任务中,PPDM 始终匹配或超越全 3D 扩散模型,同时将训练 GPU 内存使用量降低多达一个数量级,并显著加速推理,且优于基于潜在压缩或频率分解的现有内存高效扩散方法。这些结果表明,PPDM 为在有限计算资源下进行高保真 3D 扩散医学图像翻译提供了实用且可扩展的解决方案。
Diffusion models have demonstrated superior fidelity for medical image-to-image translation, but their extension to high-resolution 3D volumes is severely constrained by prohibitive computational cost and GPU memory requirements. Existing memory-efficient strategies often compromise global volumetric consistency or fine anatomical detail. In this work, we propose the Pixel Puzzling Diffusion Model (PPDM), a simple and effective framework for memory- and speed-efficient 3D medical image translation. PPDM introduces a reversible pixel puzzle-unpuzzle operator that trades spatial resolution for channel dimensionality, substantially reducing activation memory while preserving global context. To further improve efficiency and stability, we adopt a direct bridge diffusion formulation that starts from the conditional input rather than pure noise, enabling the model to focus on task-relevant residuals. In addition, a puzzle-gradient loss is incorporated to enforce spatial coherence and suppress grid-like artifacts introduced by spatial rearrangement. We evaluate PPDM on multiple challenging 3D medical image translation tasks, including low-count PET denoising, joint PET denoising and attenuation correction, and cross-modal MRI translation. Across all tasks, PPDM consistently matches or outperforms full 3D diffusion models while reducing training GPU memory usage by up to an order of magnitude and significantly accelerating inference, and it outperforms existing memory-efficient diffusion approaches based on latent compression or frequency decomposition. These results demonstrate that PPDM provides a practical and scalable solution for high-fidelity 3D diffusion-based medical image translation under limited computational resources.
核心贡献 · Key contributions
提出像素拼图扩散模型(PPDM),用于内存和速度高效的 3D 医学图像翻译。 Proposes Pixel Puzzling Diffusion Model (PPDM) for memory- and speed-efficient 3D medical image translation.
引入可逆像素拼图-解拼图算子,以空间分辨率换取通道维度,在保持全局上下文的同时减少激活内存。 Introduces a reversible pixel puzzle-unpuzzle operator that trades spatial resolution for channel dimensionality, reducing activation memory while preserving global context.
采用从条件输入开始的直接桥扩散公式,使模型专注于任务相关残差。 Adopts a direct bridge diffusion formulation starting from conditional input, enabling focus on task-relevant residuals.
引入拼图梯度损失以增强空间一致性并抑制网格状伪影。 Incorporates a puzzle-gradient loss to enforce spatial coherence and suppress grid-like artifacts.
在多个 3D 医学图像翻译任务上实现最先进性能,GPU 内存减少高达 7 倍。 Achieves state-of-the-art performance on multiple 3D medical image translation tasks with up to 7x reduction in GPU memory.
为有限计算资源下的高保真 3D 扩散医学图像翻译提供实用且可扩展的解决方案。 Provides a practical and scalable solution for high-fidelity 3D diffusion-based medical image translation under limited computational resources.
局限 · Limitations
对训练超参数敏感,尤其是损失权重λ,可能需要针对不同应用进行调整。 Sensitive to training hyperparameters, especially the loss weight lambda, which may require tuning per application.
随着空间下采样因子增加,性能逐渐下降,在准确性与内存和速度之间权衡。 Performance gradually decreases as the spatial down-scale factor increases, trading off accuracy for memory and speed.
拼图梯度损失的最佳核大小取决于下采样因子,增加了训练复杂性。 Optimal kernel size for puzzle-gradient loss depends on the down-scale factor, adding complexity to training.
评估仅限于 PET 和 MRI 的头部区域;对其他解剖区域或模态的泛化性未充分验证。 Evaluation limited to head regions in PET and MRI; generalizability to other anatomical regions or modalities not fully validated.
对于较大的下采样因子,需要仔细调整核大小以保持性能。 Requires careful adjustment of kernel size for larger down-scale factors to maintain performance.