fastWDM3D: Efficient 3D Inpainting for Brain MRIs

• fastWDM3D is a 3D diffusion inpainting model that reconstructs healthy brain tissue in masked regions of T1-weighted MRIs using a wavelet-domain approach.
• It employs a variance-preserving noise schedule and image-space reconstruction losses to achieve high SSIM, low MSE, and rapid inference in just two diffusion steps.
• The method significantly reduces computation time to 1.81 seconds per image while maintaining 3D anatomical consistency compared to prior DDPM-based inpainting techniques.

Searching arXiv for the primary paper and closely related inpainting/diffusion context. fastWDM3D is a wavelet-domain 3D diffusion inpainting model for healthy brain tissue synthesis in T1-weighted MRI, introduced for the task of reconstructing plausible healthy anatomy inside a masked pathological region while preserving the surrounding observed tissue. It was developed in the context of the BraTS Local Synthesis of Healthy Brain Tissue via Inpainting Challenge, where the input is a T1-weighted brain MRI together with a mask, and the output is a 3D inpainted healthy tissue volume (Durrer et al., 17 Jul 2025). The method is motivated by applications including pseudo-healthy baselines for tumor growth modeling and improved image registration and related MRI analysis methods that benefit from healthy-looking anatomy. Its defining result is that a modified WDM3D formulation can perform 3D healthy tissue inpainting using only two diffusion steps, reaching a SSIM of 0.8571, a MSE of 0.0079, a PSNR of 22.26, and 1.81 s per image on the BraTS inpainting test set (Durrer et al., 17 Jul 2025).

1. Problem setting and motivation

Healthy tissue inpainting denotes the reconstruction of healthy-appearing anatomy within a masked region of a pathological scan, yielding a pseudo-healthy image while preserving the observed tissue outside the mask (Durrer et al., 17 Jul 2025). In the formulation used for fastWDM3D, the model receives a T1-weighted brain MRI with a masked region to fill, together with the corresponding mask, and outputs a 3D inpainted healthy tissue volume. The work is explicitly situated within the BraTS Local Synthesis of Healthy Brain Tissue via Inpainting Challenge (Durrer et al., 17 Jul 2025).

The practical motivation is twofold. The first stated use is the generation of pseudo-healthy baselines for tumor growth modeling. The second is the facilitation of downstream tasks such as image registration and other MRI analysis methods that perform better on healthy-looking anatomy (Durrer et al., 17 Jul 2025). The paper positions the task as one where diffusion models are attractive because they can preserve 3D anatomical realism, but where conventional denoising diffusion probabilistic models had been hindered by long sampling chains and correspondingly slow inference (Durrer et al., 17 Jul 2025).

A central contrast is therefore between qualitative realism and computational tractability. Earlier DDPM-based healthy tissue inpainting approaches had shown convincing results, including in BraTS challenge settings, but often required hundreds or thousands of denoising steps and inference times measured in minutes per volume (Durrer et al., 17 Jul 2025). fastWDM3D was designed to retain the diffusion-based advantages of 3D consistency and realism while reducing inference to only two diffusion steps.

2. Lineage from WDM3D and the rejection of adversarial training

fastWDM3D is a modified version of WDM3D, a 3D wavelet diffusion model (Durrer et al., 17 Jul 2025). The original WDM3D used 1000 time steps, a linear variance schedule, and MSE on wavelet coefficients (Durrer et al., 17 Jul 2025). The fastWDM3D study first examined whether a 2D fast-generation strategy based on combining DDPMs with GANs and using a variance-preserving schedule could transfer to 3D inpainting. This led to two intermediate variants: WDDGAN3D, which retained adversarial training, and GO3D, which removed the discriminator (Durrer et al., 17 Jul 2025).

The experiments identified the transferable components precisely. The variance-preserving noise schedule and the selected reconstruction losses were found to be effective for high-quality 3D inpainting in a few time steps, whereas adversarial training was not (Durrer et al., 17 Jul 2025). The paper states that the adversarial loss did not decrease during training and that the discriminator brought no performance benefit. Training cost also differed sharply: WDDGAN3D required about 112 h and 31.69 GiB, whereas GO3D required about 11 h and 24.24 GiB (Durrer et al., 17 Jul 2025).

This empirical outcome determined the final method design. fastWDM3D is not a GAN hybrid. It is a pure diffusion model derived from WDM3D, modified to use the variance-preserving schedule and a reconstruction loss in image space focused both globally and on the masked region (Durrer et al., 17 Jul 2025). The resulting system simplifies training and lowers cost while preserving or improving performance. A plausible implication is that, for this specific 3D medical inpainting regime, diffusion trajectory design and supervisory signal were more consequential than adversarial discrimination.

3. Diffusion formulation and variance-preserving scheduling

The paper follows standard DDPM notation for the forward and reverse processes (Durrer et al., 17 Jul 2025). The forward process gradually perturbs an input image $x_0$ over $T$ steps:

$q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$

with $\vec{I}$ the identity matrix and $\beta_t$ the variance at time step $t$ (Durrer et al., 17 Jul 2025). The intended meaning is stated to be the standard Gaussian transition with signal scaling by $\sqrt{1-\beta_t}$ and isotropic variance $\beta_t\vec I$.

The reverse process is defined as

$$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$

where $\mu_\theta$ and $T$0 are parameterized by a time-conditioned model $T$1 (Durrer et al., 17 Jul 2025). The reverse process aims to match the true denoising distribution $T$2, and the paper notes that $T$3 can be trained to predict $T$4, after which $T$5 can be sampled using the posterior $T$6 (Durrer et al., 17 Jul 2025).

The critical modification enabling few-step operation is the variance-preserving schedule. Following Song et al., the forward process corresponds to the SDE

$T$7

and the paper adopts the VP schedule from Xiao et al.:

$T$8

again preserving the paper’s typesetting as given (Durrer et al., 17 Jul 2025). The practical point is explicit: this schedule makes the overall perturbation independent of the number of diffusion steps, so the image can still be fully perturbed even when $T$9 is very small. The definitions

$q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$0

are used to denote stepwise and cumulative signal preservation (Durrer et al., 17 Jul 2025).

The paper compares three schedules: linear, adapted linear, and VP (Durrer et al., 17 Jul 2025). Its analysis states that the usual linear schedule does not sufficiently perturb the image when $q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$1 is small, while the adapted linear schedule can destroy image information too early for large $q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$2. The VP schedule works for both small and large $q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$3 and is therefore preferred (Durrer et al., 17 Jul 2025). This suggests that schedule robustness across step counts was a prerequisite for compressing the denoising chain to two steps.

4. Wavelet-domain representation, conditioning, and architecture

The model operates in the wavelet domain rather than directly in full-resolution voxel space (Durrer et al., 17 Jul 2025). It uses a 3D discrete wavelet transform with low- and high-pass filters of stride 2,

$q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$4

applied along all three spatial dimensions (Durrer et al., 17 Jul 2025). A 3D volume $q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$5 is decomposed into eight subbands,

$q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$6

each at half the spatial resolution of $q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$7, and the inverse discrete wavelet transform reconstructs the original image (Durrer et al., 17 Jul 2025).

The inpainting setup is implemented using Palette-conditioning (Durrer et al., 17 Jul 2025). During training, the ground-truth image $q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$8, the mask $q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta _{t}x_{t-1},\beta _{t}\vec{I}),$9, and the corresponding voided image $\vec{I}$0 are used. Operationally, the wavelet coefficients of $\vec{I}$1, $\vec{I}$2, and the noisy $\vec{I}$3 are concatenated into a 24-channel input, and the model predicts an 8-channel output (Durrer et al., 17 Jul 2025). The output is passed through the IDWT to reconstruct the inpainted image $\vec{I}$4. Because each of $\vec{I}$5, $\vec{I}$6, and noisy $\vec{I}$7 contributes eight wavelet channels, the input dimensionality is 24 channels and the prediction target is the 8 wavelet channels corresponding to the reconstructed target (Durrer et al., 17 Jul 2025).

Architecturally, fastWDM3D inherits the WDM3D wavelet-domain 3D diffusion design rather than the adversarial WDDGAN route (Durrer et al., 17 Jul 2025). The generator-style models explored in the paper mention 4 residual blocks and, for WDDGAN, an NCSN++-based generator with a StyleGAN-like mapping network for latent conditioning, but the final fastWDM3D is the DDPM-only wavelet model rather than the adversarial NCSN++ setup (Durrer et al., 17 Jul 2025). The processing pipeline is described as follows: preprocess a cropped 3D T1 volume, generate a voided version using the healthy mask, compute DWTs of voided image, mask, and noisy target, concatenate them, run a 3D diffusion backbone in wavelet space, predict 8 wavelet output channels, and reconstruct the inpainted volume via IDWT (Durrer et al., 17 Jul 2025).

The model’s full 3D character is emphasized by comparison with 2D alternatives. The paper notes that a 2D DDPM can look realistic in axial slices but introduces stripe artifacts in coronal view due to lack of true 3D consistency (Durrer et al., 17 Jul 2025). In this context, fastWDM3D is positioned as preserving anatomically coherent volumetric structure.

5. Objective functions and the few-step regime

A major methodological departure from the original WDM3D is the supervisory signal (Durrer et al., 17 Jul 2025). For WDDGAN3D, the total loss is given as

$\vec{I}$8

Here, $\vec{I}$9 is the adversarial loss, $\beta_t$0 is a reconstruction loss between ground truth $\beta_t$1 and prediction $\beta_t$2, and $\beta_t$3 is a region-specific reconstruction loss over the masked area $\beta_t$4 (Durrer et al., 17 Jul 2025).

After removing the discriminator, GO3D uses

$\beta_t$5

and fastWDM3D explicitly replaces the original WDM3D wavelet-coefficient MSE loss with the same reconstruction objective

$\beta_t$6

The paper does not specify scalar weighting factors between these two terms and presents them as an unweighted sum (Durrer et al., 17 Jul 2025).

The significance of this change is stated directly. Rather than supervising only wavelet coefficients with MSE, fastWDM3D supervises the reconstructed image both globally and specifically within the inpainted region (Durrer et al., 17 Jul 2025). The paper attributes its speed and quality to several interacting factors: only two denoising steps instead of the 1000 steps of the original WDM3D; the VP schedule, which preserves meaningful perturbation and denoising behavior even at very low $\beta_t$7; image-space reconstruction losses $\beta_t$8, which help preserve quality in the few-step regime; wavelet-domain processing, which reduces spatial resolution per subband; and the removal of adversarial training, which lowers training complexity and memory use (Durrer et al., 17 Jul 2025).

The paper’s summary of its main experimental insights is correspondingly specific. Adversarial training is unnecessary for this task; the VP schedule is crucial; replacing wavelet-coefficient MSE with global plus masked-region reconstruction losses is important; and, in this 3D inpainting setting, very small diffusion step counts are not merely acceptable but optimal, with the best fastWDM3D model using $\beta_t$9 (Durrer et al., 17 Jul 2025).

6. Empirical results, benchmarks, and computational profile

The ablation results establish the model’s behavior in several stages. For WDDGAN3D(VP), the reported metrics were SSIM $t$0, MSE $t$1, and PSNR $t$2 at $t$3, and SSIM $t$4, MSE $t$5, and PSNR $t$6 at $t$7 (Durrer et al., 17 Jul 2025). GO3D(VP), which removes the discriminator, was reported as better or comparable while being much cheaper: at $t$8, SSIM $t$9, MSE $\sqrt{1-\beta_t}$0, and PSNR $\sqrt{1-\beta_t}$1; at $\sqrt{1-\beta_t}$2, SSIM $\sqrt{1-\beta_t}$3, MSE $\sqrt{1-\beta_t}$4, and PSNR $\sqrt{1-\beta_t}$5 (Durrer et al., 17 Jul 2025).

The schedule ablation further supports the VP choice. For GO3D(LA), the paper reports SSIM $\sqrt{1-\beta_t}$6, MSE $\sqrt{1-\beta_t}$7, and PSNR $\sqrt{1-\beta_t}$8 at $\sqrt{1-\beta_t}$9, and SSIM $\beta_t\vec I$0, MSE $\beta_t\vec I$1, and PSNR $\beta_t\vec I$2 at $\beta_t\vec I$3 (Durrer et al., 17 Jul 2025). Although the adapted linear schedule gave lower MSE numerically, it was worse in SSIM and PSNR than GO3D(VP) and WDDGAN3D, and VP was preferred because it worked robustly across different $\beta_t\vec I$4 without modification (Durrer et al., 17 Jul 2025).

The final fastWDM3D comparison against GO3D is the paper’s central benchmark. At 120k iterations, fastWDM3D achieved the following results (Durrer et al., 17 Jul 2025):

Configuration	Metrics	AST
fastWDM3D, $\beta_t\vec I$5	SSIM $\beta_t\vec I$6; MSE $\beta_t\vec I$7; PSNR $\beta_t\vec I$8	$\beta_t\vec I$9 s
fastWDM3D, $$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$0	SSIM $$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$1; MSE $$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$2; PSNR $$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$3	$$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$4 s
fastWDM3D, $$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$5	SSIM $$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$6; MSE $$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$7; PSNR $$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$8	$$p_{\theta}(x_{t-1}|x_t) := \mathcal{N}\left(x_{t-1}; \mu_{\theta}(x_t, t), \sigma_t^2\vec{I}\right),$$9 s

These results are reported on the BraTS inpainting test set, with the $\mu_\theta$0 model trained for 120k iterations, corresponding in that setup to 300 epochs with batch size 3 (Durrer et al., 17 Jul 2025). The paper notes that GO3D with $\mu_\theta$1 and batch size 2 had slightly higher SSIM in another experimental setting, but fastWDM3D($\mu_\theta$2) was selected as the best overall configuration because it had the best PSNR, the same rounded MSE, the lowest sampling time, and much lower memory use (Durrer et al., 17 Jul 2025).

The runtime comparison with prior diffusion inpainting methods is especially stark. Table 4 reports 20 min for DDPM 2D slice-wise, 25 min for DDPM 2D seq-pos, 25 min for DDPM Pseudo3D, 20 min for DDPM 3D mem-eff, 1 min for LDM3D, 5 min for WDM3D, and 1.81 s for fastWDM3D($\mu_\theta$3) (Durrer et al., 17 Jul 2025). The paper therefore states that fastWDM3D is up to $\mu_\theta$4 faster than other DDPMs used for healthy brain tissue inpainting (Durrer et al., 17 Jul 2025). Relative to the prior WDM3D specifically, it uses 500× fewer time steps, 2 instead of 1000, and the paper reports that WDM3D had SSIM $\mu_\theta$5, MSE $\mu_\theta$6, PSNR $\mu_\theta$7, and 5 min runtime, versus fastWDM3D’s SSIM $\mu_\theta$8, MSE $\mu_\theta$9, PSNR $T$00, and 1.81 s (Durrer et al., 17 Jul 2025). The authors explicitly attribute this “vast improvement” to changing the schedule and the loss.

Relative to the 2023 BraTS challenge podium, fastWDM3D is competitive but not dominant (Durrer et al., 17 Jul 2025). Zhang et al. are reported at SSIM $T$01, MSE $T$02, and PSNR $T$03; Durrer et al. at SSIM $T$04, MSE $T$05, and PSNR $T$06; Huo et al. at SSIM $T$07, MSE $T$08, and PSNR $T$09; and fastWDM3D($T$10) at SSIM $T$11, MSE $T$12, and PSNR $T$13 (Durrer et al., 17 Jul 2025). Thus, the first-place challenge method remains quantitatively stronger, but fastWDM3D outperforms the second- and third-place methods in MSE and PSNR while remaining similar in SSIM. The paper also emphasizes a qualitative advantage: fastWDM3D reconstructs 3D realistic structures, whereas some challenge methods are blurrier and 2D methods exhibit through-plane artifacts (Durrer et al., 17 Jul 2025).

7. Data regime, training protocol, limitations, and reproducibility

The training data are the BraTS 2023 Local Synthesis of Healthy Brain Tissue via Inpainting Challenge training set, containing 1251 patient scans with T1 images, tumor masks, and healthy masks (Durrer et al., 17 Jul 2025). The split used was 1200 for training and 51 for validation, while the hidden test set contains 568 patients (Durrer et al., 17 Jul 2025). Original image size was $T$14. Intensities were clipped by removing the top and bottom 0.5 percentile, then normalized to $T$15 (Durrer et al., 17 Jul 2025).

During training, the region defined by the healthy mask $T$16 was cropped from the T1 image $T$17 to produce the voided image $T$18, and then $T$19, $T$20, and $T$21 were all cropped to $T$22 with the inpainting region centered (Durrer et al., 17 Jul 2025). For evaluation, each prediction $T$23 was normalized back to the individual input range of $T$24 (Durrer et al., 17 Jul 2025). All models were trained on an NVIDIA A100 40 GB GPU (Durrer et al., 17 Jul 2025). In the Table 1 experiments, the settings were batch size 2, 4 residual blocks in the generator, learning rate $T$25, and 100 epochs = $T$26 iterations. In the Table 2 experiments, the settings were batch size 3, 4 residual blocks, learning rate $T$27, and results shown for 40k and 120k iterations (Durrer et al., 17 Jul 2025). The paper states that convergence followed by overfitting was observed in the first comparison, that GO3D did not improve after 100 epochs, and that fastWDM3D continued improving up to 300 epochs (Durrer et al., 17 Jul 2025). During training, GO3D required 38.8 GiB, while fastWDM3D required 18.3 GiB (Durrer et al., 17 Jul 2025).

The method’s strengths are positioned as a combination of high quality, 3D consistency, fast sampling, lower memory use, and simpler training (Durrer et al., 17 Jul 2025). The authors also mention lower environmental footprint because of reduced compute. However, the limitations are stated with comparable clarity. A more thorough ablation disentangling the effects of variance schedule, loss, and architecture would be beneficial; future work should investigate why the VP schedule works so well here, test additional schedules, and evaluate robustness on other datasets, modalities, and generative tasks (Durrer et al., 17 Jul 2025). The paper therefore does not claim broad generalization beyond healthy brain tissue inpainting in a BraTS-like T1 MRI setting.

For reproducibility, the paper states that code is publicly available at https://github.com/AliciaDurrer/fastWDM3D, and it identifies the practical ingredients that matter for reproduction: 3D DWT/IDWT, Palette-style conditioning, the VP noise schedule, image-space reconstruction losses $T$28, training on $T$29 crops centered on the inpainting region, intensity clipping and normalization to $T$30, and selecting very small $T$31, especially $T$32 for the final model (Durrer et al., 17 Jul 2025).

In summary, fastWDM3D is defined less by adversarial augmentation than by a particular convergence of design choices within 3D diffusion inpainting: a wavelet-domain representation, Palette-style conditioning, a variance-preserving schedule, and direct image-space supervision over both the full image and the masked region (Durrer et al., 17 Jul 2025). The paper’s central contribution is the empirical finding that, in this setting, these ingredients make two-step diffusion not merely feasible but effective, yielding a 3D healthy tissue inpainting model that is dramatically faster than earlier DDPM-based approaches while preserving strong quantitative performance and volumetric anatomical consistency (Durrer et al., 17 Jul 2025).