fastWDM3D: Fast and Accurate 3D Healthy Tissue Inpainting

2025-07-17

AI Summary

To address the low sampling efficiency and limited reconstruction accuracy in 3D healthy brain tissue inpainting, this paper proposes fastWDM3D—a novel 3D wavelet diffusion model (WDM3D) enhanced with variance-preserving noise scheduling and a dedicated reconstruction loss, eliminating adversarial training. The method achieves high-fidelity generation in merely two sampling steps. It accelerates inference by up to 800× over state-of-the-art denoising diffusion probabilistic models, completing single-volume inpainting in just 1.81 seconds. On the BraTS test set, it attains SSIM = 0.8571, MSE = 0.0079, and PSNR = 22.26—demonstrating superior quantitative performance. The synthesized pseudo-healthy baseline volumes effectively support tumor growth modeling and multi-modal image registration, providing a computationally efficient and clinically reliable 3D generative tool for neuroimaging analysis.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Cognitive Modeling & Cognitive Systems: Computational Creativity

Application Category

• Web Mining and Content Analysis: Web data generation and simulation
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Healthy tissue inpainting has significant applications, including the generation of pseudo-healthy baselines for tumor growth models and the facilitation of image registration. In previous editions of the BraTS Local Synthesis of Healthy Brain Tissue via Inpainting Challenge, denoising diffusion probabilistic models (DDPMs) demonstrated qualitatively convincing results but suffered from low sampling speed. To mitigate this limitation, we adapted a 2D image generation approach, combining DDPMs with generative adversarial networks (GANs) and employing a variance-preserving noise schedule, for the task of 3D inpainting. Our experiments showed that the variance-preserving noise schedule and the selected reconstruction losses can be effectively utilized for high-quality 3D inpainting in a few time steps without requiring adversarial training. We applied our findings to a different architecture, a 3D wavelet diffusion model (WDM3D) that does not include a GAN component. The resulting model, denoted as fastWDM3D, obtained a SSIM of 0.8571, a MSE of 0.0079, and a PSNR of 22.26 on the BraTS inpainting test set. Remarkably, it achieved these scores using only two time steps, completing the 3D inpainting process in 1.81 s per image. When compared to other DDPMs used for healthy brain tissue inpainting, our model is up to 800 x faster while still achieving superior performance metrics. Our proposed method, fastWDM3D, represents a promising approach for fast and accurate healthy tissue inpainting. Our code is available at https://github.com/AliciaDurrer/fastWDM3D.

Problem

• Improving slow sampling speed in 3D healthy tissue inpainting
• Combining DDPMs and GANs for efficient 3D image generation
• Achieving high-quality inpainting with minimal time steps

Innovation

• Combines DDPMs with GANs for 3D inpainting
• Uses variance-preserving noise schedule effectively
• Achieves fast 3D inpainting in two steps