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

2026-06-13

AI Summary

This work addresses the high computational and memory demands of high-resolution 3D medical image diffusion models, which often limit their practicality, while existing efficient approaches typically compromise global consistency or anatomical fidelity. The authors propose the Pixel Puzzle Diffusion Model (PPDM), which leverages invertible pixel puzzle–depuzzle operations to compress spatial dimensions into channel dimensions, substantially reducing memory consumption without sacrificing global context. PPDM further incorporates a direct bridging diffusion mechanism and a puzzle gradient loss to effectively suppress grid-like artifacts and enhance spatial coherence. Evaluated across multiple 3D medical image translation tasks, PPDM matches or exceeds the performance of full 3D diffusion models while reducing training memory usage by nearly an order of magnitude and significantly accelerating inference, outperforming current methods based on latent-space compression or frequency-domain decomposition.

Abstract

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.

Problem

• 3D medical image translation
• diffusion models
• GPU memory efficiency
• computational cost
• volumetric consistency

Innovation

• Pixel Puzzling
• Memory-Efficient Diffusion
• 3D Medical Image Translation
• Direct Bridge Diffusion
• Puzzle-Gradient Loss