Source-linked AI summary
Fast Unsupervised Brain Anomaly Detection and Segmentation with Diffusion Models
Walter H. L. Pinaya, Mark S. Graham, Robert Gray, Pedro F Da Costa, Petru-Daniel Tudosiu, Paul Wright, Yee H. Mah, Andrew D. MacKinnon, James T. Teo, Rolf Jager, David Werring, Geraint Rees, Parashkev Nachev, Sebastien Ourselin, M. Jorge Cardoso
TL;DR
Medical anomaly segmentation needs methods that avoid transformer sequence-order and sampling limitations while reducing inference time. This paper trains a VQ-VAE and DDPM on healthy brain data, detects unlikely latent states, and heals them through reverse diffusion before deriving pixel-space residuals. Across synthetic and real CT and MRI lesions, the method performs competitively with autoregressive approaches and substantially reduces inference time.
Problem
Autoregressive transformer approaches for medical anomaly detection require one-dimensional image sequences, can accumulate sampling errors, and have significant inference times.
Method
The method compresses images with a VQ-VAE, learns healthy latent representations with a DDPM, identifies unlikely Markov-chain states, and heals anomalies through masked reverse diffusion.
Results
The method performed competitively with transformers on synthetic and real data, with better performance than a single transformer in most cases and much reduced inference times.
Takeaways & Limitations
Faster diffusion inference supports the method’s potential use for time-constrained clinical anomaly detection and segmentation.
Abstract
from arXiv · showhide
Deep generative models have emerged as promising tools for detecting arbitrary anomalies in data, dispensing with the necessity for manual labelling. Recently, autoregressive transformers have achieved state-of-the-art performance for anomaly detection in medical imaging. Nonetheless, these models still have some intrinsic weaknesses, such as requiring images to be modelled as 1D sequences, the accumulation of errors during the sampling process, and the significant inference times associated with transformers. Denoising diffusion probabilistic models are a class of non-autoregressive generative models recently shown to produce excellent samples in computer vision (surpassing Generative Adversarial Networks), and to achieve log-likelihoods that are competitive with transformers while having fast inference times. Diffusion models can be applied to the latent representations learnt by autoencoders, making them easily scalable and great candidates for application to high dimensional data, such as medical images. Here, we propose a method based on diffusion models to detect and segment anomalies in brain imaging. By training the models on healthy data and then exploring its diffusion and reverse steps across its Markov chain, we can identify anomalous areas in the latent space and hence identify anomalies in the pixel space. Our diffusion models achieve competitive performance compared with autoregressive approaches across a series of experiments with 2D CT and MRI data involving synthetic and real pathological lesions with much reduced inference times, making their usage clinically viable.
1 Introduction
Neuroimaging lesion segmentation is clinically valuable but traditionally labor-intensive and dependent on expertise. Diffusion models are proposed as a non-autoregressive alternative to transformers that avoids sequence-order bias and accumulated sampling errors while supporting fast anomaly detection and segmentation.
- Motivation: Manual lesion segmentation is time-consuming and depends on human expertise, motivating accurate automatic methods for clinical use.Generative models can learn normal-data distributions and highlight pathological deviations without expensive labels or anomalous training images.
- Limitations of autoregressive models: Autoregressive transformers model images as ordered one-dimensional sequences, imposing a unidirectional bias that is perceptually unnatural for brain images.Prior work used ensembles with differently ordered latent sequences to improve performance.
- Limitations of autoregressive models: Teacher-forced training differs from inference based on previously sampled elements, allowing prediction errors to accumulate during likelihood computation and anomaly healing.These accumulated errors may affect the quality and coherence of anomaly-corrected images.
- Proposed alternative: The proposed approach uses DDPMs to remove sequence-order bias and accumulated healing errors while targeting clinically viable inference times.DDPMs iteratively reverse a forward corruption process and are applied to VQ-VAE latent representations for high-dimensional neuroimaging.
2 Background
The method combines VQ-VAE compression with a DDPM that learns healthy latent representations through forward diffusion and reverse denoising. Its probabilistic training objective and accelerated samplers support efficient modeling of high-dimensional brain data.
- Compression model: A VQ-VAE compresses images into compact quantized latent representations, reducing the computational complexity of the diffusion model.The encoder maps images to latents, quantization selects codebook vectors, and the decoder reconstructs observations.
- Denoising diffusion probabilistic models: The DDPM learns the distribution of healthy brain-image latents using a forward Markov process that gradually destroys structure and a reverse process that recovers it.Forward transitions add Gaussian noise according to a fixed variance schedule; reverse transitions predict the cumulative noise.
- Denoising diffusion probabilistic models: The diffusion model can sample noisy latents at arbitrary timesteps in closed form, enabling efficient access to intermediate states.The closed-form representation expresses each noisy state using the clean latent and Gaussian noise.
- Efficient sampling: DDIM sampling can reduce the number of reverse steps by 10× ∼50×, accelerating diffusion inference.
- Training objective: Training maximizes a variational lower bound on log-likelihood, with reverse-step errors compared against corresponding diffusion posteriors.The simplified objective trains the model to predict noise added at uniformly sampled timesteps.
3 Proposed anomaly segmentation method
The method detects unlikely healthy-latent transitions, replaces anomalous latent values through masked reverse diffusion, and converts the healed representation into residual maps for pixel-level segmentation.
- Anomaly identification: The method uses the DDPM’s variational lower-bound terms across its Markov chain to identify latent values unlikely under healthy training data.Unlikely values are replaced with more probable DDPM values, while latent spatial information filters residual maps.
- Anomaly identification: Healthy inputs produce low reverse-step KL divergence, whereas anomalous inputs produce deviations from the expected Gaussian transitions.The method computes these values after encoding test images and diffusing their latent representations across the Markov chain.
- Anomaly identification: The method calibrates thresholds on healthy validation images and focuses on intermediate diffusion steps, where different timesteps represent different image features.The study reports t = [400, 600] as less noisy for its thresholding procedure.
- Healing: Masked reverse diffusion inpaints highlighted latent regions while preserving the rest of the original latent representation as context.The process starts at t = 500 in synthetic lesion tests and progressively denoises only masked regions.
- Residual segmentation: The healed latent is decoded into pixel space, and smoothed upsampled masks weight pixel residuals to identify anomalous regions with high final values.
4 Experiments
Experiments evaluated the DDPM-based method on synthetic and real 2D MRI and CT lesions, measuring segmentation performance, anomaly detection, and inference time. The method performed competitively while remaining under one minute in the reported CT setting and scaling more easily across resolutions than transformers.
- Synthetic anomalies: The synthetic-data evaluation trained on 9,000 HeadCT images and tested 100 images corrupted with sprites.
- Synthetic anomalies: The complete approach achieved the best synthetic-dataset results among its evaluated processing steps, while the single model outperformed the transformer and slightly exceeded the ensemble on ⌈DICE⌉ but had slightly lower AUPRC.
- Synthetic anomalies: DDPM anomaly detection on corrupted images obtained AUCROC=0.827 and AUPRC=0.702, compared with AUCROC=0.921 and AUPRC=0.707 for the transformer-based approach.
- Real MRI lesions: Real-lesion evaluation used healthy UK Biobank FLAIR images and tested small-vessel disease, tumours, demyelinating lesions, and white-matter hyperintensities.
- Inference time: Unlike transformers, whose required forward passes can grow with sequence length, DDPM forward-pass counts remain constant across resolutions, facilitating scaling to higher-resolution latent spaces.
5 Conclusions
The study proposes DDPMs for unsupervised brain anomaly detection and segmentation. Across synthetic and real data, the method performed competitively with transformers, with faster inference supporting time-constrained clinical use.
- The proposed method uses DDPMs for unsupervised brain anomaly detection and segmentation.
- The method performed competitively with transformers on synthetic and real data and showed better performance than a single transformer in most cases.
- Faster inference, especially with DDIMs, supports use in scenarios with prediction-time constraints and may help bring anomaly detection into clinical workflows.