Latent Diffusion Model – LoDoChallenge (DM4CT)
This repository contains the pretrained latent-space diffusion model used in the benchmark DM4CT: Benchmarking Diffusion Models for Computed Tomography Reconstruction (ICLR 2026).
- Paper: DM4CT: Benchmarking Diffusion Models for Computed Tomography Reconstruction
- Project Page: https://dm4ct.github.io/DM4CT/
- Codebase: https://github.com/DM4CT/DM4CT
🔬 Model Overview
This model learns a prior over CT reconstruction images in a compressed latent space using a denoising diffusion probabilistic model (DDPM).
Unlike the pixel diffusion model, diffusion is performed in the latent space of a pretrained autoencoder.
- Architecture:
- VQ-VAE (image encoder/decoder)
- 2D UNet operating in latent space
- Input resolution (image space): 512 × 512
- Channels: 1 (grayscale CT slice)
- Training objective: ε-prediction (standard DDPM formulation)
- Noise schedule: Linear beta schedule
- Training dataset: Low Dose Grand Challenge (LoDoChallenge)
- Intensity normalization: Rescaled to (-1, 1)
The diffusion model operates purely in latent space and relies on the autoencoder for encoding and decoding. This model is intended to be combined with data-consistency correction for CT reconstruction.
📊 Dataset: Low Dose Grand Challenge
Source: https://www.aapm.org/grandchallenge/lowdosect/
Preprocessing steps:
- Train/test split
- Rescale reconstructed slices to (-1, 1)
- No geometry information is embedded in the model
The model learns an unconditional latent prior over CT slices.
🧠 Training Details
- Optimizer: AdamW
- Learning rate: 1e-4
- Hardware: NVIDIA A100 GPU
- Training scripts: train_latent.py
🚀 Usage
from diffusers import DiffusionPipeline
import torch
pipeline = DiffusionPipeline.from_pretrained(
"jiayangshi/lodochallenge_latent_diffusion"
)
pipeline.to("cuda")
# Generate an unconditional CT slice prior
image = pipeline(batch_size=1).images[0]
image.save("reconstructed_slice.png")
Citation
@inproceedings{shi2026dmct,
title={{DM}4{CT}: Benchmarking Diffusion Models for Computed Tomography Reconstruction},
author={Shi, Jiayang and Pelt, Dani{\"e}l M and Batenburg, K Joost},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=YE5scJekg5}
}