Zero-Shot Anomaly Detection in Brain MRIs using Diffusion Models

Introduction

In this tutorial, we will explore how to use Denoising Diffusion Probabilistic Models (DDPMs) for zero-shot anomaly detection on brain MRIs. By training a diffusion model to generate healthy brain MRIs, we can detect anomalies (tumors, lesions) in unseen images by evaluating the reconstruction error when the model attempts to map pathological images onto the healthy data manifold.

Architecture Overview

graph LR
    Path(Pathological MRI) --> AddNoise(Add Noise to t=500)
    AddNoise --> Denoise(Denoise to t=0)
    Denoise --> Heal(Healthy Reconstruction)
    Path --> Diff(Subtract)
    Heal --> Diff
    Diff --> Anom(Anomaly Map)
        

Prerequisites

• Python 3.10+
• PyTorch 2.0+
• MONAI 1.3+
• A dataset of healthy brain MRIs (e.g., from OASIS or fastMRI)

Step 1: Setting up the Diffusion Model with MONAI

We leverage MONAI's generative module to initialize a 2D/3D Diffusion Model.

import torch
from monai.networks.nets import DiffusionModelUNet

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Initialize a 2D U-Net for Diffusion
model = DiffusionModelUNet(
    spatial_dims=2,
    in_channels=1,
    out_channels=1,
    num_channels=(64, 128, 256),
    attention_levels=(False, True, True),
    num_res_blocks=2,
    num_head_channels=32,
).to(device)

print(f"Model parameters: {sum(p.numel() for p in model.parameters())}")

Step 2: Training the Model on Healthy Data

Train the DDPM strictly on healthy brain slices. The model learns the probability distribution of healthy anatomy.

from monai.generative.inferers import DiffusionInferer
from monai.generative.schedulers import DDPMScheduler

scheduler = DDPMScheduler(num_train_timesteps=1000)
inferer = DiffusionInferer(scheduler)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# Simplified training loop
for epoch in range(epochs):
    for batch in dataloader:
        images = batch["image"].to(device)
        noise = torch.randn_like(images).to(device)
        timesteps = torch.randint(0, scheduler.num_train_timesteps, (images.shape[0],), device=device).long()
        
        noisy_images = scheduler.add_noise(images, noise, timesteps)
        
        optimizer.zero_grad()
        prediction = inferer(inputs=noisy_images, diffusion_model=model, timesteps=timesteps)
        loss = torch.nn.functional.mse_loss(prediction, noise)
        loss.backward()
        optimizer.step()

Step 3: Anomaly Detection via Image-to-Image Translation

To detect anomalies, we add noise to a test image (which may contain a tumor) up to timestep $t$, and then denoise it back to $t=0$. For medical images, using Simplex noise instead of standard Gaussian noise often yields better "healing" of large anomalies because it preserves low-frequency spatial structure during the forward pass.

def get_anomaly_map(pathological_image, t=500, use_simplex=True):
    # Add noise to timestep t
    if use_simplex:
        # Pseudo-code for Simplex noise generation (requires external library)
        # noise = generate_simplex_noise_like(pathological_image)
        pass 
    else:
        noise = torch.randn_like(pathological_image)
        
    noise = torch.randn_like(pathological_image) # fallback
    noisy_img = scheduler.add_noise(pathological_image, noise, torch.tensor([t]).to(device))
    
    # Denoise back to 0
    reconstructed = noisy_img
    for step in reversed(range(0, t)):
        reconstructed = inferer.sample(
            input_noise=reconstructed, 
            diffusion_model=model, 
            scheduler=scheduler, 
            timesteps=torch.tensor([step]).to(device)
        )
        
    # Anomaly map is the absolute difference, often smoothed with a Gaussian blur
    anomaly_map = torch.abs(pathological_image - reconstructed)
    return anomaly_map

Conclusion

Diffusion models offer an exciting frontier for unsupervised anomaly detection in medical imaging, bypassing the need for expensive pixel-level annotations. Advanced techniques like Simplex Noise and latent diffusion models (LDMs) continue to push the boundaries of what is possible in zero-shot medical AI.