Binary Noise for Binary Tasks: Masked Bernoulli Diffusion for Unsupervised Anomaly Detection

TL;DR

A novel and fast unsupervised anomaly detection approach based on latent Bernoulli diffusion models that achieves state-of-the-art performance compared to other diffusion-based unsupervised anomaly detection algorithms while significantly reducing sampling time and memory consumption.

Abstract

The high performance of denoising diffusion models for image generation has paved the way for their application in unsupervised medical anomaly detection. As diffusion-based methods require a lot of GPU memory and have long sampling times, we present a novel and fast unsupervised anomaly detection approach based on latent Bernoulli diffusion models. We first apply an autoencoder to compress the input images into a binary latent representation. Next, a diffusion model that follows a Bernoulli noise schedule is employed to this latent space and trained to restore binary latent representations from perturbed ones. The binary nature of this diffusion model allows us to identify entries in the latent space that have a high probability of flipping their binary code during the denoising process, which indicates out-of-distribution data. We propose a masking algorithm based on these probabilities, which improves the anomaly detection scores. We achieve state-of-the-art performance compared to other diffusion-based unsupervised anomaly detection algorithms while significantly reducing sampling time and memory consumption. The code is available at https://github.com/JuliaWolleb/Anomaly_berdiff.

Figures (4)

• Figure 1: We train a binarizing autoencoder consisting of an encoder $E$ and a decoder $D$. In the latent space, a diffusion model is trained to restore the binary latent code of healthy images by reversing a Bernoulli noise driven diffusion process. Both the autoencoder and the diffusion model are exclusively trained on healthy data.
• Figure 2: In our proposed denoising process, we aim to restore a healthy representation $z_0$ from an initial input representation $z$. We define a mask $\mathcal{M}$ based on the threshold probability $P$. By stitching the sampled and original latent representations at each timestep, we ensure the preservation of anatomical information from the input.
• Figure 3: On the left, we present Dice scores on the BRATS2020 test set for different noise levels $L$ and probability thresholds $P$. On the right, we show the percentage of the masked entries of $\mathcal{M}$ for the healthy (blue) and diseased (orange) groups for exemplary settings of $P$ and $L$.
• Figure 4: Qualitiative results of all comparing methods for a subject from the BRATS2020 test set, as well as a subject from the OCT2017 dataset suffering from drusen.