MRIComp4Flow: Compression of 3D Brain MRI for Training Multi-Modal Generative Models

MRIComp4Flow: Compression of 3D Brain MRI for Training Multi-Modal Generative Models

Abstract Summary: The paper investigates whether standard image codecs can effectively compress semantically rich brain tumor MRI while preserving the fidelity required to train and deploy a 3D MRI generative model. Each 3D volume is compressed with JPEG2000 or a near-lossless JPEG-LS pipeline. A Wavelet Flow Matching model, conditioned on BraTS image sequences (T1n, T1c, T2, T2f), is trained on compressed data and evaluated on the validation set. At a 20:1 compression ratio, synthesis quality is statistically equivalent to a model trained on uncompressed data within a pre-specified margin (ΔPSNR < 1 dB, ΔSSIM < 0.02; paired TOST p = [[p]]): mean PSNR is 27.3 dB vs. 27.0 dB and mean SSIM is 0.95 vs. 0.96 across modalities. The results indicate that JPEG2000 compression is a practical step toward scalable 3D MRI generative modeling without degrading synthesis quality. The code is available at https://github.com/lisafis/MRIComp4Flow.

Introduction Summary: Multi-modal brain MRI is central to tumor segmentation, treatment planning, and longitudinal monitoring. Public benchmarks such as BraTS 2024 provide co-registered non-enhanced T1-weighted (T1n), contrast-enhanced T1-weighted (T1c), T2-weighted (T2w), and FLAIR (T2f) volumes. Training 3D generative models on these data requires repeated random access to full volumes during optimization, making storage bandwidth and disk capacity first-order bottlenecks. For instance, reducing a 1 TB raw NIfTI dataset by a factor of 20 allows the entire corpus to fit on a local NVMe drive, drastically accelerating data loading during training. Prior work by Kurmukov et al. shows that discriminative 3D segmentation networks tolerate JPEG2000 compression up to 20:1 on BraTS MRI with no loss of accuracy, raising the question of whether deep generative models, which must learn the full data distribution rather than a decision boundary, inherit the same robustness. The paper proposes MRIComp4Flow, a pipeline that compresses BraTS volumes using widely supported codecs, reconstructs 3D tensors on the fly during training, and trains a unified Wavelet Flow Matching (WFM) model on the compressed corpus. Models trained on uncompressed NIfTI data are compared against models trained on JPEG2000 (ratios 2:1–400:1) and near-lossless JPEG-LS (NEAR levels 1–255) across all four modalities.

Contributions:

1. Compression pipeline: A slice-wise JPEG2000/JPEG-LS pipeline and compressed data loader for 3D MRI, compatible with WFM training.

2. Generative robustness: The first characterization of how generative synthesis quality varies with input compression, showing that WFM tolerates JPEG2000 up to 20:1 and thereby extending prior discriminative model findings to generative models.

3. Comprehensive evaluation: Quantitative and qualitative synthesis results across nine JPEG2000 ratios and seven JPEG-LS levels, downstream tumor segmentation with a state-of-the-art BraTS2025 algorithm, and ablations on the number of sampling steps and the contribution of the informed prior.

Related Work Summary: JPEG2000 and JPEG-LS are established standards in radiological workflows. JPEG2000 offers embedded rate control via compression ratios and is used in DICOM storage profiles; JPEG-LS provides efficient near-lossless coding through a maximum pixel error (NEAR) parameter. Prior work on MRI compression has focused on diagnostic fidelity rather than downstream machine learning tasks. Conditional generative models, including GANs, diffusion models, and flow-matching methods, synthesize missing MRI contrasts from available sequences. Wavelet Flow Matching (WFM) learns a velocity field in Haar wavelet space from an informed prior (the mean of observed modalities) to the target contrast, enabling accurate synthesis with a single 82M-parameter network for all four BraTS modalities. Training on compressed images is common in natural-image domains but remains underexplored for 3D medical volume synthesis.

Method Summary: MRIComp4Flow comprises three stages: compression, synthesis, and reconstruction. First, raw BraTS NIfTI volumes are compressed offline into per-slice archives. Next, for training a generative model, slices are decoded, stacked into 3D volumes, intensity-normalized, and fed to WFM in Haar wavelet space. Finally, during inference, the model synthesizes missing modalities.

For compression, each volume Vm ∈ RH×W×D is processed independently. Voxel intensities are transformed to [0, ∞), cast to uint16, and segmentation masks are excluded from lossy coding. Each axial slice is encoded either with JPEG2000 (9/7 irreversible transform, glymur) at ratios r ∈ [2, 400], skipping empty background slices, or with near-lossless JPEG-LS (NEAR error; NEAR=0 is visually lossless for 16-bit data). The ratio r is the codec’s nominal per-slice target; the realized whole-dataset reduction is smaller because empty background slices are skipped and container overhead is added. For example, the nominal 20:1 JPEG2000 setting compresses the BraTS training set from 40.37 GB to 3.13 GB, a 12.9× reduction. The compressed loader decodes all.jp2 (or.jls) slices, and inserts each slice at its axial index into a 182 × 218 × 182 buffer. Volumes are clipped to the 0.1st - 99.9th intensity percentiles, min–max normalized to [0, 1], spatially padded to 240 × 240 × 160 and center-cropped, then trimmed to 224 × 224 × 160 to match the WFM input geometry.

For the WFM backbone, the model operates in 3D Haar wavelet space. Let DWT(·) denote the 3D discrete wavelet transform producing eight subbands concatenated along the channel dimension. For target modality y conditioned on the modality set C, the informed prior is defined as the mean of the available modalities in wavelet space: xsource = (1/C) Σc∈C DWT(xc). Following conditional flow matching, noisy states interpolate between source and target with a σ-schedule peaking at t = 0.5: xt = (1 − t) xsource + t xtarget + σ(t) ϵ, σ(t) = σmax t(1 − t). A U-Net fθ predicts the constant velocity v = xtarget − xsource by minimizing L = Et,ϵ [∥fθ(x̃t, c, t, y) − (xtarget − xsource)∥2], where c stacks conditioning wavelet coefficients and y is a class label over the four modalities. At inference, one Euler step from t = 0 to t = 1 followed by IDWT(·) yields the synthesized volume.

Experiments Summary: The BraTS2024 Glioma (GLI) cohort is used with four co-registered MRI modalities per case (1,621/188 train/val), comprising four co-registered 182×218×182 modalities (T1n, T1c, T2w, T2f) at 1 mm isotropic resolution. Evaluation uses a validation set; for each case and each target modality, the model receives three remaining modalities as conditioning and synthesizes the held-out contrast.

For training, all models share an identical configuration: an ∼82M parameter WavUNet (unified 4-class conditioning, 3D Haar wavelets) trained with Adam (lr 10−5), batch size 4, and σmax = 0.5. 18 unified WFM models are trained with the same hyperparameters, including a Baseline (uncompressed) model using the original BraTS NIfTI training data. Compressed models include JPEG2000 at 20:1, JPEG-LS with near=20, and additional experiments at ratios 1 (JPEG-LS only), 2, 5, 10, 50, 100 (JPEG2000 only), 200, 255 (JPEG-LS only), 400 (JPEG2000 only).

Synthesized volumes are saved as NIfTI and compared against uncompressed ground-truth references. PSNR and SSIM (3D, skimage) are reported after per-volume min–max normalization to [0, 1]. Timing is measured for each subject across all four synthesized modalities. Per-subject means across the 188 validation cases with 95% confidence intervals are reported; baseline-vs-Comp-20 differences are assessed with a paired Wilcoxon signed-rank test and a two-one-sided-tests (TOST) equivalence test at the margins above.

Results Summary:

Reconstruction fidelity of compressed inputs: JPEG2000 is effectively lossless at low ratios (∞ dB PSNR, SSIM = 1.0000 at r ∈ 2, 5) and degrades gracefully, still retaining 38.05 dB / 0.9741 SSIM at 20:1. JPEG-LS loses fidelity faster as NEAR grows, dropping from 54.45 dB at NEAR=1 to 24.70 dB at NEAR=50; notably, JPEG2000 at the aggressive 100:1 setting (25.70 dB) is comparable to JPEG-LS at NEAR=50.

Synthesis Quality: Input compression does not degrade synthesis at practical ratios: quality is maintained at 20:1 and, by PSNR/SSIM, remains close to baseline up to 100:1 (T1n 27.66 dB, T2w 25.95 dB), degrading markedly only at the extreme 200:1 and 400:1 settings (T1n 24.65/20.94 dB, T2w 22.63/17.18 dB). Near-lossless JPEG-LS preserves synthesis quality across most NEAR levels. Significant degradation only happens at very severe compression levels. This confirms that JPEG-LS retains the fidelity required for multi-modal generative training even though its raw-slice reconstruction error grows faster than JPEG2000’s.

Qualitative evaluation: Near-lossless JPEG-LS preserves fine anatomical structure and contrast boundaries across all four modalities, while JPEG2000 maintains structural integrity up to 20:1 and only exhibits visible smoothing and edge-ringing at 50:1–100:1. The 20:1 model preserves macroscopic tumor anatomy and contrast boundaries, with residual differences confined to subtle high-frequency texture.

Inference speed: Compression leaves inference throughput unchanged: both baseline and compressed models synthesize all four modalities in a single step at ∼0.35–0.37 s per volume (4 × 188 modality outputs in ∼265–278 s total), since the codec acts only on stored data and not on the sampling path.

Downstream task: Whole-tumor segmentation using the winning solution from the BraTS2025 segmentation challenge shows that whole-tumor delineation is well preserved (Dice 0.814, 95th percentile Hausdorff distance 5.85 mm), whereas the smaller Tumor Core (0.238) region shows lower agreement and high variance. This gradient-robust macrostructure, fragile fine substructure, is consistent with the high-frequency smoothing observed qualitatively and marks the main limitation of compressed-data synthesis. The low Tumor-Core agreement (0.238) indicates that small, high-frequency substructures are the primary casualty of aggressive compression, which is flagged as the main clinical limitation.

Discussion & Conclusion Summary: Across eight JPEG2000 ratios and six JPEG-LS levels, multi-modal synthesis quality remains comparable to training on lossless NIfTI up to 20:1 (a 12.9× storage reduction) with unchanged sub-second inference; by PSNR/SSIM it holds up to 100:1 and collapses only at 200:1–400:1, while visible wavelet artifacts emerge earlier, at 50:1–100:1. Notably, at several intermediate ratios, synthesis and downstream metrics exceed the uncompressed baseline, suggesting that mild lossy compression can act as a beneficial denoising regularizer, an effect recently reported for frozen foundation-VAE reconstructions of 3D body CT, and a promising avenue for deliberate exploitation. A natural next step is to train generative models directly on the compressed latent representations rather than on decoded volumes, coupling storage efficiency with representation reuse.

Improvements to AI systems based on this paper:

1. Compression-aware training pipeline for 3D generative models: Implement a data-loading system that decodes JPEG2000/JPEG-LS compressed volumetric slices on-the-fly during training, reducing storage requirements by 13× (from 40.37 GB to 3.13 GB for BraTS) without sacrificing synthesis quality (ΔPSNR < 1 dB, ΔSSIM < 0.02 at 20:1). The improved system can train multi-modal generative models on datasets that previously exceeded local NVMe capacity, enabling faster iteration and larger-scale experiments.

2. Robustness thresholding for generative models: Use the identified 20:1 JPEG2000 compression threshold (and up to 100:1 for PSNR/SSIM) to automatically select optimal compression ratios for new medical imaging datasets, balancing storage efficiency against downstream synthesis fidelity. The system can predict when compression artifacts will degrade outputs (e.g., at 200:1–400:1) and flag them for human review.

3. Denoising regularization via mild compression: Exploit the observed effect where intermediate compression ratios (e.g., 20:1–100:1) sometimes improve synthesis metrics over uncompressed baselines (e.g., T1n PSNR 27.66 dB vs. 27.0 dB baseline). The improved system can deliberately apply mild lossy compression as a regularizer during training, potentially reducing overfitting to high-frequency noise in raw MRI data.

4. Compressed-domain generative modeling: Extend the pipeline to train generative models directly on compressed latent representations (e.g., wavelet coefficients or JPEG2000 bitstreams) rather than decoded volumes, coupling storage efficiency with representation reuse. This could enable training on even larger cohorts or higher-resolution volumes without additional hardware.

5. Adaptive compression selection per modality and task: Use the per-modality quality metrics (e.g., T1n, T1c, T2w, T2f) to dynamically choose compression settings—for instance, applying lower compression to modalities critical for tumor-core delineation (which showed fragile substructure at Dice 0.238) while using higher compression for less critical contrasts. The system can optimize the storage-quality tradeoff based on the target downstream task (e.g., segmentation vs. synthesis).

6. Quality-aware inference with compression-aware uncertainty: Integrate the observed artifact patterns (high-frequency smoothing, edge-ringing at 50:1–100:1) into uncertainty estimation for synthesized volumes. The improved system can flag regions where compression-induced artifacts are likely (e.g., tumor core boundaries) and provide confidence maps, aiding clinical decision-making.

7. Benchmarking framework for compression-generative co-design: Provide a standardized evaluation protocol (paired TOST equivalence tests, PSNR/SSIM margins, downstream segmentation Dice) that other researchers can use to validate new codecs or generative architectures, accelerating progress in storage-efficient medical AI.

Large-scale multi-modal MRI datasets impose substantial storage and I/O costs, limiting the training of 3D generative models on commodity infrastructure. While lossy compression is known to preserve accuracy for discriminative segmentation networks, its effect on generative models, which must learn the full data distribution rather than a decision boundary, is unexplored. We study whether standard image codecs can effectively compress semantically rich brain tumor MRI while preserving the fidelity required to train and deploy a 3D MRI generative model. Each 3D volume is compressed with JPEG2000 or a near-lossless JPEG-LS pipeline. Next, a Wavelet Flow Matching model, conditioned on BraTS image sequences (T1n, T1c, T2, T2f), is trained on compressed data, and the resulting models are evaluated on the validation set. At a 20:1 compression ratio, synthesis quality is statistically equivalent to a model trained on uncompressed data within a pre-specified margin (PSNR <1,dB, SSIM <0.02; paired TOST p=[[p]]): mean PSNR is 27.3,dB vs. 27.0,dB and mean SSIM is 0.95 vs. 0.96 across modalities. Our results indicate that JPEG2000 compression is a practical step toward scalable 3D MRI generative modeling without degrading synthesis quality. The codebase is available at https://github.com/lisafis/MRIComp4Flow.