Source-linked AI summary
Domain-Specific Latent Representations Improve the Fidelity of Diffusion-Based Medical Image Super-Resolution
Sebastian Cajas, Ashaba Judith, Rahul Gorijavolu, Sahil Kapadia, Hillary Clinton Kasimbazi, Leo Kinyera, Emmanuel Paul Kwesiga, Sri Sri Jaithra Varma Manthena, Luis Filipe Nakayama, Ninsiima Doreen, Leo Anthony Celi
TL;DR
Medical image super-resolution systems commonly rely on autoencoders designed for natural photographs, leaving the value of domain-specific latent representations insufficiently tested. This paper isolates the VAE by replacing Stable Diffusion’s VAE with MedVAE, finding a +2.91 to +3.29 dB PSNR improvement across three medical imaging modalities and a practical basis for prioritizing VAE selection.
Problem
The study addresses limited evidence on whether generic natural-image VAEs constrain latent diffusion medical image super-resolution relative to domain-specific autoencoders.
Method
A controlled experiment replaces Stable Diffusion’s VAE with MedVAE while holding the remaining super-resolution pipeline components constant.
Results
+2.91 to +3.29 dB PSNR improvement was achieved across knee MRI, brain MRI, and chest X-ray after swapping the VAE latent space.
Takeaways & Limitations
Autoencoder reconstruction fidelity can serve as a training-free screening criterion, supporting domain-specific VAE selection before diffusion architecture search.
Takeaways & Limitations
The evaluation uses a single 4× upsampling factor and a 2D MedVAE, so performance at other scales and in volumetric 3D SR remains uncertain.
Abstract
from arXiv · showhide
Latent diffusion models for medical image super-resolution universally inherit variational autoencoders designed for natural photographs. We show that this default choice, not the diffusion architecture, is the dominant constraint on reconstruction quality. In a controlled experiment holding all other pipeline components fixed, replacing the generic Stable Diffusion VAE with MedVAE, a domain-specific autoencoder pretrained on more than 1.6 million medical images, yields +2.91 to +3.29 dB PSNR improvement across knee MRI, brain MRI, and chest X-ray (n = 1,820; Cohen's d = 1.37 to 1.86, all p < 10^{-20}, Wilcoxon signed-rank). Wavelet decomposition localises the advantage to the finest spatial frequency bands encoding anatomically relevant fine structure. Ablations across inference schedules, prediction targets, and generative architectures confirm the gap is stable within plus or minus 0.15 dB, while hallucination rates remain comparable between methods (Cohen's h < 0.02 across all datasets), establishing that reconstruction fidelity and generative hallucination are governed by independent pipeline components. These results provide a practical screening criterion: autoencoder reconstruction quality, measurable without diffusion training, predicts downstream SR performance (R^2 = 0.67), suggesting that domain-specific VAE selection should precede diffusion architecture search. Code and trained model weights are publicly available at https://github.com/sebasmos/latent-sr.
1 Introduction
Medical image super-resolution addresses globally uneven access to high-resolution imaging, while latent diffusion models may impose a bottleneck through generic natural-image autoencoders. This study isolates the VAE’s contribution and evaluates whether domain-specific latent representations improve SR fidelity without changing the diffusion pipeline.
- Motivation: Low- and middle-income countries disproportionately rely on low-field MRI systems that produce inferior spatial resolution, motivating medical image super-resolution.An estimated two-thirds of the world’s population lacks MRI access, and most installed devices in sub-Saharan Africa operate below 0.3 T.
- Problem: Generic natural-image VAEs may discard modality-specific noise textures, intensity distributions, and anatomical fine structure before diffusion processing begins.The resulting distorted latent manifold cannot recover information already discarded by the encoder through additional diffusion training.
- Method: The controlled experiment replaces SD-VAE with MedVAE while fixing the UNet, training objective, noise schedule, and evaluation protocol across knee MRI, brain MRI, and chest X-ray.MedVAE was pretrained on > 1.6M medical images, and evaluation used n = 1,820 validation images with direct PSNR and LPIPS comparisons.
- Contributions: +3.93 to +6.48 dB: MedVAE’s autoencoder reconstruction ceiling exceeds SD-VAE’s across all datasets, identifying latent-space fidelity as the SR quality constraint.The result supports a modality-agnostic, training-free VAE screening principle based on encode-then-decode reconstruction quality.
- Contributions: < 0.15 dB: the PSNR gap varies across inference schedules, prediction targets, and generative architectures, while hallucination rates remain comparable with Cohen’s h < 0.02.These findings separate reconstruction fidelity, governed by the VAE, from generative hallucination, governed by the diffusion model.
2 Results
MedVAE substantially improves diffusion-based medical image super-resolution fidelity over SD-VAE, with gains concentrated in fine spatial frequencies and stable across tested settings. These fidelity gains occur without reducing hallucination rates, indicating distinct VAE and diffusion error components.
- +2.91 dB and +3.29 dB PSNR: MedVAE SR outperforms SD-VAE SR on BraTS and CXR, respectively, in 4× super-resolution.MedVAE SR reaches 26.42 ± 2.10 dB on BraTS and 28.87 ± 2.62 dB on CXR, versus 23.51 ± 1.99 dB and 25.58 ± 2.17 dB for SD-VAE SR.
- 0.088–0.135 vs. 0.097–0.173 LPIPS: MedVAE SR has lower perceptual error than SD-VAE SR and lies on the quality-perception Pareto frontier.Across all three datasets, MedVAE SR also has better PSNR than SD-VAE SR and lower LPIPS than bicubic interpolation.
- +1.18 dB, +1.41 dB, and +0.70 dB HH1 improvements: MedVAE’s advantage concentrates in the finest wavelet detail bands on MRNet, BraTS, and CXR.Coarse LL3 gains are only +0.05 dB on MRNet and +0.02 dB on BraTS, while CXR shows slightly better SD-VAE performance at mid-to-low frequencies.
- 12.9% vs. 13.3% on BraTS and 3.3% vs. 3.4% on CXR: MedVAE and SD-VAE have comparable hallucination rates.Cohen’s h < 0.02 across all datasets, while MedVAE still achieves lower mean absolute error on MRNet (0.152 vs. 0.154) and BraTS (0.060 vs. 0.061).
3 Discussion
The discussion identifies the VAE latent space as the dominant, architecture-independent bottleneck in medical image super-resolution, with a controlled swap producing a large, stable PSNR gain. It also separates reconstruction fidelity from hallucination and outlines frequency-aware design, resource-allocation implications, and study limitations.
- Core finding: +2.91 to +3.29 dB PSNR improvement resulted from swapping only the VAE latent space across three independent medical imaging modalities.The effect sizes were d = 1.37–1.86, while all other pipeline components remained fixed.
- Core finding: 31 vs. 38 on BraTS and 48 vs. 59 on CXR show lower FID for SD-VAE SR than MedVAE SR, consistent with ImageNet-feature preference for natural-image-like textures.The discussion interprets this counterpoint through the perception-distortion tradeoff and possible preference for hallucinated structures.
- Core finding: The frozen encoder imposes an absolute reconstruction ceiling, and ablations found that prediction targets, inference steps, loss weighting, and flow matching could not close the VAE-swap gap.The discussion presents this VAE-as-bottleneck principle as a general consequence of latent diffusion construction.
- Frequency-band analysis: ≈−0.81/−0.61 dB at LH2/HL2 on chest X-ray favored SD-VAE-AE, but +0.70 dB at HH1 for MedVAE and fine-band dominance yielded +3.29 dB overall.The same mid-low-frequency anomaly appeared in encode-then-decode images, implicating the SD-VAE encoder rather than diffusion.
- Practical implications: Improved VAE domain specificity raises reconstruction fidelity but leaves hallucination rates unchanged, so hallucination reduction requires diffusion-stage interventions.The discussion attributes fabricated content proximately to the stochastic diffusion process and recommends score guidance, conformal prediction intervals, or voxel-level uncertainty estimation.
- Practical implications: ±0.1 dB PSNR-gap stability from T = 50 to T = 1,000 and across prediction targets and generative architectures argues against prioritizing longer sampling or architecture search.The tested generative architectures included DDPM and flow matching.
4 Methods
The study evaluates latent-diffusion super-resolution across knee MRI, brain MRI, and chest radiography using matched datasets and a controlled MedVAE versus SD-VAE comparison. It standardizes latent diffusion training and assesses fidelity, perceptual quality, frequency content, hallucinations, latent alignment, and statistical significance.
- Datasets: The evaluation uses 1,370 knee MRI exams, 700 validation and 720 test brain-MRI slices, and 1,000 validation chest radiographs, with images resized to 256 × 256.MRNet uses mildly degraded HR images for refinement, while BraTS and MIMIC-CXR use bicubic 4× downsampling from 256 × 256 to 64 × 64.
- Model comparison: MedVAE and SD-VAE use the identical UNet architecture, isolating the VAE contribution in the primary diffusion-based comparison.MedVAE encodes latents with shape 3 × 64 × 64 and was pretrained on 1.6M medical images, whereas SD-VAE uses 4 × 32 × 32 latents.
- Diffusion training: The LDM predicts clean HR latents x0 from noisy latents xt conditioned on LR latents, using unweighted L1 loss, a cosine schedule, and 100 inference steps.Training uses T = 1,000 diffusion steps; latents are cached before UNet training.
- Evaluation metrics: PSNR and MS-SSIM are the primary fidelity metrics, while LPIPS and FID provide complementary perceptual indicators with ImageNet-feature domain-mismatch caveats.MRNet FID values are reported for completeness but are unreliable because n = 120.
- Analysis procedures: The analysis quantifies frequency-specific reconstruction using three-level Haar wavelets, hallucinated and lost pixels using an AE-derived noise floor, and SR-HR latent alignment across seven spatial scales.Statistical comparisons use paired two-sided Wilcoxon signed-rank tests and Cohen’s d with 95% bootstrap confidence intervals from 10,000 resamples.
T PSNR (dB) MS-SSIM LPIPS sec/sample Speedup
Supplementary analyses show that MedVAE-based super-resolution preserves perceptual quality better than bicubic interpolation, while diffusion methods perform comparably in whole-image tumor-region analysis. The MedVAE autoencoder also sets a substantially higher reconstruction ceiling than SD-VAE, whereas MRNet FID estimates are unreliable at the available sample size.
- BraTS ROI analysis: 17.05 dB vs. 17.03 dB: MedVAE SR and SD-VAE SR have comparable whole-image PSNR in BraTS tumor-region analysis.Bicubic reaches 29.91 dB, while low tumor-region PSNR reflects the extreme low-frequency content of diffuse glioma lesions at 4× down-sampled resolution.
- Autoencoder reconstruction: +3.93 to +6.48 dB: the MedVAE autoencoder reconstruction ceiling exceeds the SD-VAE ceiling without any diffusion component.The ceiling is measured using encode-then-decode reconstruction.
- FID limitation: n = 120 < 2,048: MRNet FID values are unreliable because the sample size is below the stated minimum.The limitation applies to the reported MRNet FID estimates.
B Supplementary Figures
Supplementary figures show that MedVAE SR improves both distortion and perceptual fidelity relative to SD-VAE SR across three medical imaging datasets. Additional figures assess pixel-level hallucination and content loss across knee MRI, brain MRI, and chest X-ray.
- Perception-distortion tradeoff: MedVAE SR achieves simultaneously higher PSNR and lower LPIPS than SD-VAE SR on MRNet, BraTS, and MIMIC-CXR.This places MedVAE SR toward the upper-right Pareto frontier in the perception-distortion tradeoff.
- Pixel-level hallucination and content loss: The supplementary analysis evaluates pixel-level hallucination and content loss for MedVAE SR and SD-VAE SR on MRNet knee MRI.Figure S2 reports hallucination and loss rates as percentages of pixels.
- Pixel-level hallucination and content loss: Corresponding pixel-level hallucination and content-loss analyses cover BraTS brain MRI and MIMIC-CXR chest X-ray at 4× super-resolution.Figures S3 and S4 extend the analysis to brain MRI and chest X-ray.
C Supplementary Methods
The supplementary methods specify the VAE architectures, preprocessing, statistical tests, wavelet analysis, hallucination criterion, and latent-similarity procedure used in the experiments. MedVAE and SD-VAE remained frozen throughout LDM experiments.
- VAE architecture details: MedVAE produces 3 × 64 × 64 latents with 4× compression, whereas SD-VAE produces 4 × 32 × 32 latents with 8× spatial compression.MedVAE uses a KL-regularised convolutional autoencoder with residual blocks, group normalisation, and SiLU activations.
- VAE architecture details: Both VAEs used frozen weights throughout all latent diffusion model experiments.
- Dataset preprocessing: All images were resized to 256×256 pixels and normalised to [0, 1], with 4× tasks using 64 × 64 bicubic downsampling followed by upsampling.MRNet training additionally used mild Gaussian noise, and evaluation used no data augmentation.
- Statistical methodology: Wilcoxon tests used per-image PSNR pairs, while Cohen’s d and 95% bootstrap confidence intervals used 10,000 resamples of per-image PSNR differences.Tests were two-sided, with no multiple-comparison correction under the stated a priori hypothesis.
- Wavelet decomposition: Three-level Haar wavelet analysis produced 10 subbands, and per-subband PSNR was averaged across validation images.Each image yielded one LL3 approximation band and LH, HL, and HH detail bands at each decomposition level.
- Hallucination quantification: Hallucinated or lost pixels were identified from signed SR−HR differences exceeding the corresponding autoencoder noise-floor threshold.The threshold was |D| > µ|AE−HR| + 2σ|AE−HR|, with the sign of D determining hallucination versus loss.
D Supplementary Notes · D.1 Supplementary Note S1: CXR Mid-Low Frequency Anomaly
On MIMIC-CXR, SD-VAE slightly outperforms MedVAE in mid-low-frequency wavelet subbands because of its encoder representation, while MedVAE retains a fine-detail advantage. The anomaly is therefore localized to spatial frequency and originates before diffusion.
- D.1 Supplementary Note S1: CXR Mid-Low Frequency Anomaly: −0.61 to −0.81 dB: SD-VAE SR marginally outperforms MedVAE SR at MIMIC-CXR mid-low subbands LH2/HL2.This comparison concerns chest radiograph mid-low spatial-frequency content.
- D.1 Supplementary Note S1: CXR Mid-Low Frequency Anomaly: 0.05–0.10 dB: SD-VAE-AE also scores higher than MedVAE-AE at LH2/HL2 in AE-only images.The same mid-low-frequency pattern appears without downstream diffusion.
- D.1 Supplementary Note S1: CXR Mid-Low Frequency Anomaly: The AE-only replication identifies the mid-low-frequency anomaly as a property of the SD-VAE encoder’s latent representation.It is not attributed to an artefact introduced by the downstream diffusion model.
- D.1 Supplementary Note S1: CXR Mid-Low Frequency Anomaly: A plausible explanation is SD-VAE’s larger channel capacity, specified as 4 × 32 × 32, for encoding broad, slowly varying contrast gradients.These gradients dominate chest-radiograph mid-low-frequency content.
- D.1 Supplementary Note S1: CXR Mid-Low Frequency Anomaly: +0.70 dB: MedVAE retains an advantage over SD-VAE at the finest-scale HH1 subband on CXR.HH1 represents fine detail in the wavelet decomposition.
- D.1 Supplementary Note S1: CXR Mid-Low Frequency Anomaly: MedVAE’s HH1 advantage is consistent with its domain-specific training.The passage contrasts this fine-detail result with SD-VAE’s mid-low-frequency advantage.
D.2 Supplementary Note S2: MRNet Task Description
MRNet comprises 1,370 knee MRI exams and evaluates refinement from blurred, 2× downsampled inputs restored to the original 256 × 256 dimensions. Across diffusion timestep budgets, latent fidelity remains narrowly stable, indicating primary dependence on VAE architecture.
- MRNet Task Description: 1,370 knee MRI exams form the MRNet dataset, with axial 256 × 256 slices acquired at Stanford University Medical Center from 2001–2012.MRNet differs from BraTS and MIMIC-CXR by using a refinement task rather than a 4× bicubic-downsampling setup.
- MRNet Task Description: ∆≤0.025 across diffusion timestep budgets confirms that latent fidelity is determined primarily by the VAE architecture.The comparison uses cosine similarity between SR and HR latents at the global mean (1 × 1) scale.
D.3 Supplementary Note S3: Disentangling Latent Capacity from Domain Specificity
MedVAE’s larger latent space cannot be isolated from domain specificity through direct architectural substitution because its UNet expects different latent dimensions and channels. However, autoencoder ceiling gaps are consistent with and bound the +2.91–+3.29 dB SR gains, supporting domain specificity as the primary driver.
- Latent Capacity Versus Domain Specificity: 3× capacity separates MedVAE’s 12,288-dimensional latent space from SD-VAE’s 4,096-dimensional space.MedVAE uses 3 × 64 × 64 dimensions, while SD-VAE uses 4 × 32 × 32 dimensions.
- Latent Capacity Versus Domain Specificity: MedVAE and SD-VAE are architecturally infeasible to compare by direct substitution because their UNets expect different channel counts and spatial resolutions.MedVAE expects a 6-channel 64 × 64 input, whereas SD-VAE expects an 8-channel 32 × 32 input.
- Latent Capacity Versus Domain Specificity: +2.91–+3.29 dB SR gains are consistent with and bounded by autoencoder ceiling gaps, supporting domain specificity rather than latent capacity as the improvement driver.The comparison is reported as providing strong evidence that domain specificity drives the improvement.
D.4 Supplementary Note S4: Inference Efficiency at T = 16 Steps · E Supplementary Figures (continued)
At 16 DDPM steps, MedVAE SR shows a perception-distortion tradeoff relative to 100 steps across MRNet, BraTS, and CXR. Supplementary figures compare MedVAE and SD-VAE SR with autoencoder reconstructions, high-resolution ground truth, and pixel-difference maps.
- D.4 Supplementary Note S4: Inference Efficiency at T = 16 Steps: T = 16 achieves LPIPS values of 0.218 (MRNet), 0.086 (BraTS), and 0.301 (CXR).These results were obtained with MedVAE SR across all three datasets using 16 DDPM steps.
- D.4 Supplementary Note S4: Inference Efficiency at T = 16 Steps: T = 100 achieves LPIPS values of 0.135 (MRNet), 0.088 (BraTS), and 0.127 (CXR), compared with T = 16.The comparison evaluates MedVAE SR under the two inference schedules.
- D.4 Supplementary Note S4: Inference Efficiency at T = 16 Steps: T = 16 produces outputs closer to the LR conditioning, whereas T = 100 introduces greater perceptual detail at some pixel-level fidelity cost.The schedule comparison therefore reveals a perception-distortion tradeoff.
- D.4 Supplementary Note S4: Inference Efficiency at T = 16 Steps: Progressive distillation offers a complementary route to inference acceleration and is a natural candidate for future MedVAE pipeline integration.This is proposed as future work rather than evaluated in the reported experiment.
- E Supplementary Figures (continued): Figure S5 compares AE reconstruction, MedVAE SR, SD-VAE SR, and HR ground truth for MRNet knee MRI.The figure also shows mean absolute pixel difference |SR −HR| using a hot colormap, where brighter regions indicate larger error.
- E Supplementary Figures (continued): Figure S6 compares AE reconstruction, MedVAE SR, SD-VAE SR, and HR ground truth for MIMIC-CXR chest X-ray 4× super-resolution.The figure includes mean absolute pixel difference |SR −HR| maps in a hot colormap.