Foundation-VAE

 

We make a progressive stride toward training-free medical VAEs by leveraging a critical observation: a single Foundation VAE, pretrained at scale on natural images and videos, can serve as a unified interface for CT Reconstruction, Augmentation, and Generation. (1) CT Reconstruction: A Foundation VAE pretrained at scale on natural images and videos reconstructs a 3D CT volume via its frozen encoder $E$ and decoder $D$. (2) CT Augmentation: Through zero shot transfer to CT, the reconstructed volumes provide a boundary enhanced training view, improving downstream segmentation especially on surface accuracy. (3) CT Generation: In the fixed latent space of the same \textit{Foundation VAE}, we train a conditional latent diffusion model to synthesize anatomically consistent healthy and abnormal CT volumes, controlled by organ masks and clinical findings.

Paper

Foundation VAE for CT Reconstruction, Augmentation, and Generation
Qi Chen^1,*, Shuhan Ding^2,*, Yu Gu³, Nan Liu², Jiang Bian³, Alan L. Yuille¹, Zongwei Zhou¹, and Jingjing Fu^3
1 Johns Hopkins University
² Duke-NUS Medical School
³ Microsoft Research
^* Equal contribution

paper | code | huggingface

We have summarized publications related to Medical VAE in Awesome Medical VAE .

0. Installation

git clone https://github.com/qic999/Foundation-VAE.git
cd Foundation-VAE

1. Reconstruction

We transfer a Foundation VAE pretrained on natural images/videos to 3D CT reconstruction with both encoder and decoder frozen. This reconstruction operator suppresses acquisition noise while preserving clinically relevant anatomical boundaries, making it suitable as a stable CT interface across heterogeneous scanners and protocols.

Demo

MSD Task06 Lung: reconstruction and segmentation comparison.

MSD Task07 Pancreas: reconstruction and segmentation comparison.

Data

• You can evaluate on MSD dataset
• Our released reconstruction assets: https://huggingface.co/qicq1c/Foundation-VAE/tree/main/Reconstruction

2. Augmentation

Our augmentation strategy is reconstruction-based augmentation: use reconstructed CT volumes as an additional training view for downstream tasks. Because reconstruction is boundary-stable and largely preserves label-defining geometry, segmentation trained on reconstructed CTs is comparable to, and often better than, training on raw CTs, with clear gains on boundary-sensitive metrics.

This stage is annotation-free with respect to reconstruction itself and can be directly plugged into standard segmentation pipelines.

Released augmentation models/assets: https://huggingface.co/qicq1c/Foundation-VAE/tree/main/Augmentation

3. Generation

In the same fixed Foundation VAE latent space, we train a conditional latent diffusion model for controllable 3D CT generation.

Conditioning

• Anatomy masks for spatial grounding
• Disease masks for pathology control
• Radiology report embeddings for semantic control

Key design

• Frozen Foundation VAE encoder/decoder as latent interface
• Mask latents concatenated during denoising for structure consistency
• Lightweight 3D consistency attention across slices for coherent volumetric anatomy/pathology

This enables controllable synthesis of healthy and abnormal CT volumes under unified latent modeling.

Demo

Three-view (axial/coronal/sagittal) generated CT with report conditioning.

Anatomical and pathological grounding comparison.

Released generation models/assets: https://huggingface.co/qicq1c/Foundation-VAE/tree/main/Generation

Citation

@inproceedings{chenfoundation,
  title={Foundation VAE for CT Reconstruction, Augmentation, and Generation},
  author={Chen, Qi and Ding, Shuhan and Gu, Yu and Liu, Nan and Bian, Jiang and Yuille, Alan and Zhou, Zongwei and Fu, Jingjing},
  booktitle={Forty-third International Conference on Machine Learning}
}