Image-conditioned latent rectified flow models for 3D medical anomaly localisation

Reconstruction-based methods offer a promising solution for unsupervised anomaly detection in medical imaging tasks. These methods train generative models on healthy data alone and identify anomalies as deviations between an input image and its pseudo-healthy reconstruction. The downfall of these methods is their dependence on two assumptions that often fail in practice: that models cannot reproduce unseen pathologies, yet can faithfully reconstruct healthy tissue. A recent image-conditioned diffusion approach explicitly addresses these issues by training a model to restore synthetic anomalies inserted into healthy images. However, it operates in 2D pixel space, discarding inter-slice context and incurring high computational cost.

We address both limitations by performing image-conditioned restoration in a 3D latent space using a pretrained VAE and rectified flow, capturing volumetric context whilst drastically reducing computational overhead. To mitigate false positives introduced by VAE compression, we propose using the restoration change which measures the difference between the pseudo-healthy latent restoration and the VAE reconstruction of the original, rather than the standard reconstruction error. We further experiment with applying the synthetic anomaly training task directly in latent space to improve sensitivity to low-contrast anomalies.

We perform extensive experiments across various medical imaging benchmarks, including brain MRI and the newly released AADD dataset, comparing against reconstruction-based, feature-modelling, attention-based and self-supervised anomaly detection methods. Our image-conditioned rectified flow models establish a new state-of-the-art, with an ensemble of models trained with pixel-space and latent-space anomalies yielding the strongest overall performance.

These results demonstrate how incorporating 3D context enables better anomaly detection performance whilst also being ∼10 times faster. The difference in performance between models trained using latent-space and pixel-space anomalies suggests that further broadening of the anomaly imputation process could continue to improve the robustness of these models. Such improvements are certainly necessary, as the AADD benchmark is far from being saturated. Code is available at https://github.com/matt-baugh/img-cond-latent-rflow-model-ad.