WaveDiT: Wavelet-Domain MRI Synthesis

• WaveDiT is a conditional flow matching framework that synthesizes full-resolution 3D brain MRIs by operating in the coefficient space of a 3D Haar wavelet transform.
• It employs band-wise heteroscedastic uncertainty modeling using higher-order wavelet statistics to adaptively preserve both low- and high-frequency anatomical details.
• It leverages a factorized spatio-depth transformer attention mechanism to drastically reduce computational cost, enabling near-real-time synthesis on a single modern GPU.

WaveDiT is a conditional flow matching framework for full-resolution 3D brain MRI synthesis that operates in the coefficient space of a 3D Haar Discrete Wavelet Transform rather than in voxel space or a learned latent space. It is presented as a response to a specific constraint in neuroimaging data augmentation: large and demographically balanced datasets are essential for reliable neuroimaging biomarkers, yet existing volumetric generative approaches either incur prohibitive computational cost at 3D scale or rely on lossy latent compression that may compromise anatomical detail. The method combines a wavelet-domain representation, factorized spatio-depth attention, and band-wise heteroscedastic uncertainty modeling derived from higher-order wavelet statistics, with the stated aim of supporting full-resolution 3D synthesis under practical memory and time constraints on a single modern GPU (Danese et al., 7 Jun 2026).

1. Problem formulation and design rationale

The motivating problem is 3D brain MRI synthesis for data augmentation in settings where reliable neuroimaging biomarkers depend on large and demographically balanced datasets. In the formulation associated with WaveDiT, the central limitation of prior approaches is twofold: volumetric methods may incur prohibitive computational cost, while latent-space methods may rely on lossy compression that can compromise anatomical detail (Danese et al., 7 Jun 2026).

WaveDiT addresses this by retaining a full-resolution generative target while changing the representation and the dynamics. The model is described as a conditional flow matching framework operating in the coefficient space of a 3D Haar Discrete Wavelet Transform. This design places the synthesis process in a basis where low- and high-frequency anatomical content are explicitly separated into subbands, and where the high-frequency bands exhibit heavy-tailed and heteroscedastic statistics. The paper further states that predicted log-variance is integrated directly into both the flow objective and conditioning pathway, enabling adaptive precision consistent with the heavy-tailed and input-dependent variance structure of anatomical detail (Danese et al., 7 Jun 2026).

A plausible implication is that the framework is not merely using wavelets for dimensionality reduction; rather, it uses the wavelet basis as a distribution-aware generative domain in which uncertainty can be modeled band-wise. The paper’s summary of its novelties makes this explicit by identifying three coupled components: a wavelet-domain flow-matching model preserving invertibility and multiscale detail without learned compression; Morpheus heteroscedastic conditioning from band-wise higher-order statistics; and a factorized spatio-depth transformer backbone that scales to 1.4 M tokens at practical cost (Danese et al., 7 Jun 2026).

2. Wavelet-domain representation and data flow

The input is a full-resolution T₁ 3D MRI volume

$X \in \mathbb{R}^{1\times D\times H\times W}.$

WaveDiT applies a single-level 3D Haar discrete wavelet transform $W$, producing

$W(X)=\{x_b\}_{b=1\ldots 8},$

where $b=1$ is the low-frequency “LLL” band of size $(D/2\times H/2\times W/2)$, and $b=2\ldots 8$ are the seven high-frequency subbands $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$ of the same downsampled size (Danese et al., 7 Jun 2026).

The representation is characterized in three ways. First, the transform is perfectly invertible. Second, it concentrates $\sim 98\%$ of energy in the LLL band. Third, it exposes the heavy-tailed, heteroscedastic statistics of the HF bands, with kurtosis up to $\sim 270$ at $t=1$ (Danese et al., 7 Jun 2026). These statements define the statistical premise of the method: the low-frequency band carries most of the signal energy, while the high-frequency bands encode difficult, non-Gaussian anatomical detail.

After the transform, the $W$0 wavelet tensor is reshaped into $W$1 2D slices, described as tokens of size $W$2, with patch embeddings of size $W$3. The generative backbone then operates over these slice-wise tokens rather than over the original voxel grid (Danese et al., 7 Jun 2026).

A plausible misconception is that wavelet preprocessing necessarily introduces information loss comparable to latent compression. The formulation here points in the opposite direction: the transform is described as perfectly invertible, and the stated novelty is preservation of invertibility and multiscale detail without learned compression (Danese et al., 7 Jun 2026).

3. Conditional flow matching and heteroscedastic uncertainty

WaveDiT uses rectified conditional flow matching (CFM). The interpolation path is given by

$W$4

with target velocity

$W$5

The state $W$6 is decomposed into wavelet bands $W$7 (Danese et al., 7 Jun 2026).

The heteroscedastic component is provided by Morpheus, which predicts a band-wise log-variance

$W$8

from six statistics of $W$9—mean, $W(X)=\{x_b\}_{b=1\ldots 8},$0, max, $W(X)=\{x_b\}_{b=1\ldots 8},$1, skewness, kurtosis—plus a time embedding:

$W(X)=\{x_b\}_{b=1\ldots 8},$2

Under a Gaussian likelihood on the per-band velocity errors, the loss is

$W(X)=\{x_b\}_{b=1\ldots 8},$3

Equivalently, writing $W(X)=\{x_b\}_{b=1\ldots 8},$4,

$W(X)=\{x_b\}_{b=1\ldots 8},$5

The terms are described directly in the paper: $W(X)=\{x_b\}_{b=1\ldots 8},$6 is the model’s predicted instantaneous velocity in band $W(X)=\{x_b\}_{b=1\ldots 8},$7; $W(X)=\{x_b\}_{b=1\ldots 8},$8 is the true velocity along the linear path for band $W(X)=\{x_b\}_{b=1\ldots 8},$9; $b=1$0 re-weights the MSE so that high-variance HF bands are down-weighted; and $b=1$1 prevents trivially blowing up $b=1$2 (Danese et al., 7 Jun 2026).

The interpretation supplied in the source is explicit. Because $b=1$3 depends on $b=1$4 via its higher-order statistics, the model automatically assigns large $b=1$5 (low precision) to bands and spatial states that are inherently unpredictable, such as HHH near tissue boundaries, and small $b=1$6 (high precision) to stable, Gaussian-like content, such as LLL or early $b=1$7 (Danese et al., 7 Jun 2026). This is the sense in which the method is “distribution-aware”: the uncertainty schedule is conditioned on the empirical statistics of the evolving wavelet coefficients rather than being fixed a priori.

4. HDiT backbone and factorized spatio-depth attention

The conditional flow-matching backbone is termed HDiT. Its transformer blocks are modulated via AdaRMSNorm, and the conditioning variables are listed precisely: time, slice index, age metadata, and frequency hints from Morpheus are injected as multiplicative scales without concatenating to the token vectors (Danese et al., 7 Jun 2026).

The architecture has two resolution levels. At Level 1, the model uses “neighborhood attention”: 2D sliding-window self-attention inside each slice with a $b=1$8 window, for example $b=1$9, and AxialRoPE positional encoding. This stage is described as capturing local texture and edge structure in HF bands (Danese et al., 7 Jun 2026).

At Level 2, the model uses “factorized spatio-depth attention,” split into two stages. The spatial stage applies full 2D self-attention independently within each slice to capture long-range in-plane dependencies, with bilateral symmetry given as an example. The depth stage then gathers, for each fixed spatial location $(D/2\times H/2\times W/2)$0, the corresponding tokens across the $(D/2\times H/2\times W/2)$1 slices and applies a small self-attention along the depth axis. This stage is described as restoring volumetric coherence, with continuity of cortical folds given as an example (Danese et al., 7 Jun 2026).

The computational analysis is stated in asymptotic form. Let $(D/2\times H/2\times W/2)$2 be the number of 2D tokens per slice and $(D/2\times H/2\times W/2)$3 the number of slices. A full 3D self-attention would cost

$(D/2\times H/2\times W/2)$4

WaveDiT instead uses:

• Level 1 neighborhood attention of cost $(D/2\times H/2\times W/2)$5.
• Level 2 two-stage attention with spatial cost $(D/2\times H/2\times W/2)$6 and depth cost $(D/2\times H/2\times W/2)$7.

The resulting total cost is

$(D/2\times H/2\times W/2)$8

The description further quantifies the benefit for the depth-aware hourglass transformer: factorized spatio-depth attention restores volumetric coherence at a cost of $(D/2\times H/2\times W/2)$9, about $b=2\ldots 8$0 cheaper than naïve 3D attention (Danese et al., 7 Jun 2026).

At inference time, the process is also specified exactly: draw $b=2\ldots 8$1, integrate the learned ODE $b=2\ldots 8$2 from $b=2\ldots 8$3 with a 2nd-order Heun solver, then apply the inverse Haar DWT to recover the voxel-space MRI (Danese et al., 7 Jun 2026).

5. Computational profile and practical deployment

The reported computational setting is unusually constrained for full-resolution 3D synthesis. All results were produced on a single NVIDIA H100 with $b=2\ldots 8$4 GB VRAM and batch size 1 (Danese et al., 7 Jun 2026). The paper states that this supports full-resolution 3D synthesis under practical memory and time constraints on a single modern GPU.

Training is reported to take $b=2\ldots 8$5 hours to converge over 200 epochs. Two baselines are compared on the same hardware: FlowLet took $b=2\ldots 8$6 days and WDM (wavelet diffusion) $b=2\ldots 8$7 days (Danese et al., 7 Jun 2026). Inference with 10 integration steps takes $b=2\ldots 8$8 second per volume; by comparison, FlowLet with 10 steps takes $b=2\ldots 8$9 s, and WDM with 1000 steps takes $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$0 s (Danese et al., 7 Jun 2026).

The contextual claim attached to these numbers is that latent diffusion baselines and pixel-space diffusion typically require hundreds–thousands of steps or an additional encoder/decoder, making them slower or introducing reconstruction blur (Danese et al., 7 Jun 2026). Within the paper’s framing, WaveDiT’s practical significance lies in the combination of full-resolution generation, wavelet-domain invertibility, and near-real-time synthesis on a single affordable GPU.

A plausible implication is that the method’s efficiency is not attributable to a single optimization alone. Rather, it follows from the interaction of the single-level 3D Haar DWT, the slice-wise tokenization strategy, the factorized spatio-depth attention scheme, and the ODE-based flow-matching inference with a 2nd-order Heun solver.

6. Empirical evaluation, ablations, and reported significance

WaveDiT was evaluated on a 5,989-subject multi-site cohort comprising OpenBHB, ADNI, and OASIS-3, with 20% held out for downstream BAP (Danese et al., 7 Jun 2026). The paper reports three complementary evaluation axes.

For global distribution alignment, the metrics are FID, MMD on features of a medical-pretrained ResNet-50, and MS-SSIM for intra-set diversity. WaveDiT-CFM with 10 steps achieves FID $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$1 (best), MMD $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$2 (second best), and MS-SSIM $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$3 (best) (Danese et al., 7 Jun 2026).

For Brain Age Prediction (BAP), a separate 3D DenseNet is trained on synthetic data to predict age, and test MAE is measured on real subjects $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$4 yrs. The reported values are: real-data-only MAE $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$5 yrs; FlowLet $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$6; WaveDiT-CFM $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$7 yrs (best) (Danese et al., 7 Jun 2026).

For region-of-interest anatomical agreement, 95 brain regions are segmented with FastSurfer, and region-wise intensity MAE (iMAE), KL divergence, and Dice coefficient are computed. WaveDiT-CFM reports iMAE $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$8, KLD $(\mathrm{LLH}, \mathrm{LHL}, \ldots, \mathrm{HHH})$9, and Dice $\sim 98\%$0, all best in class; the second-best FlowLet reports iMAE $\sim 98\%$1, KLD $\sim 98\%$2, and Dice $\sim 98\%$3 (Danese et al., 7 Jun 2026).

The ablation findings are summarized directly: operating in the wavelet domain, using conditional (CFM) trajectories, and heteroscedastic scheduling via Morpheus are all critical for stable training and superior anatomical fidelity (Danese et al., 7 Jun 2026). This suggests that the reported gains are not attributed solely to the transformer backbone or solely to the wavelet transform; instead, the empirical case is that performance depends on the specific combination of representation, trajectory formulation, and uncertainty modeling.

The paper’s own summary identifies the main contributions as follows: a wavelet-domain flow-matching model preserving invertibility and multiscale detail without learned compression; Morpheus heteroscedastic conditioning from band-wise higher-order statistics, yielding adaptive precision across subbands; a factorized spatio-depth transformer backbone that scales to 1.4 M tokens at practical cost; and real-time 3D MRI synthesis on a single GPU, with state-of-the-art distributional, clinical (BAP), and ROI-level performance (Danese et al., 7 Jun 2026).