Unsupervised brain imaging 3D anomaly detection and segmentation with transformers

Abstract

Pathological brain appearances may be so heterogeneous as to be intelligible only as anomalies, defined by their deviation from normality rather than any specific set of pathological features. Amongst the hardest tasks in medical imaging, detecting such anomalies requires models of the normal brain that combine compactness with the expressivity of the complex, long-range interactions that characterise its structural organisation. These are requirements transformers have arguably greater potential to satisfy than other current candidate architectures, but their application has been inhibited by their demands on data and computational resources. Here we combine the latent representation of vector quantised variational autoencoders with an ensemble of autoregressive transformers to enable unsupervised anomaly detection and segmentation defined by deviation from healthy brain imaging data, achievable at low computational cost, within relative modest data regimes. We compare our method to current state-of-the-art approaches across a series of experiments with 2D and 3D data involving synthetic and real pathological lesions. On real lesions, we train our models on 15,000 radiologically normal participants from UK Biobank and evaluate performance on four different brain MR datasets with small vessel disease, demyelinating lesions, and tumours. We demonstrate superior anomaly detection performance both image-wise and pixel/voxel-wise, achievable without post-processing. These results draw attention to the potential of transformers in this most challenging of imaging tasks.

Keywords: Anomaly detection; Transformer; Unsupervised anomaly segmentation; Vector quantized variational autoencoder.

Graphical abstract

Fig. 1

Our method uses a VQ-VAE to learn the latent discrete representation of brain data. This latent representation is transformed into a 1D sequence that is learned by the autoregressive transformer.

Fig. 2

Anomaly segmentation method. A) the sequence obtained from the VQ-VAE is fed to the transformer with an “begin of sentence” token prepended. For each position of the sequence, the transformer will predict the value of the next element. Using the output probability of each real value, we apply a threshold (in this example, we use a threshold of 0.05) to identify which one is anomalous. A binary mask (the “resampling mask”) is created to indicate which value is below the threshold and should be corrected. B) For each position indicated in the resampling mask, we use the transformer to obtain values that have a higher probability of occurrence and we create a healed sequence. C) The healed 1-dimensional sequence is reshaped and processed by the VQ-VAE decoder to create a reconstruction without anomalies.

Fig. 3

Using the spatial information from the resampling mask to improve segmentation. First, we reshape the resampling mask back to the format of the VQ-VAE latent space. Then, we upsample it to have the input image shape and we smooth it with a Gaussian filter. Finally, we use this mask to filter the residuals maps obtained from the difference between the inputted image and its healed reconstruction.

Fig. 4

To predict the probability of the value in the red square, the transformer using the ordering of the left image (raster ordering, left → right, top → bottom) mostly uses the information of the image background as context (blue squares). If the transformer uses the ordering of the right image (raster ordering, right → left, bottom → top), it will have a richer context, with more information about the brain, that could help make a more accurate prediction about the value in the red square.

Fig. 5

Residual maps on the synthetic examples from the variational autoencoder and different steps of our approach.

Fig. 6

Different orderings used to transform the 2D latent representation into a 1D sequence.

Fig. 7

Performance with synthetic anomalies with different intensity values. We also performed the analysis including an additive Gaussian noise into the anomalies. The performance is measure by the best achievable DICE-score.

Fig. 8

Log-likelihood distribution of the classes of examples evaluated by our ensemble of models, in-distribution, near out-of-distribution (near OOD), and far out-of-distribution (far OOD). The model assigned higher log-likelihoods for examples similar to the training set, intermediary values for examples with small synthetic lesions and lower values for examples of different classes.

Fig. 9

Residual maps on the real lesions from the variational autoencoder, the f-AnoGAN, and our transformer-based method.

Fig. 10

Anomaly detection image-wise on 3D data. In this experiment, we use the log-likelihood obtained from the transformers and the lesion size from the binary mask predicted by our models to train a one-class SVM and classify subjects with multiple sclerosis diagnosis in their records as out of distribution.