💻 computer science

Safeguarding AI in Medical Imaging: Post-Hoc Out-of-Distribution Detection with Normalizing Flows

This paper proposes a post-hoc normalizing flow-based method for out-of-distribution detection in medical imaging that integrates seamlessly with pre-trained models without retraining, achieving superior performance on the MedOOD and MedMNIST benchmarks compared to existing approaches.

Original authors: Dariush Lotfi, Mohammad-Ali Nikouei Mahani, Mohamad Koohi-Moghadam, Kyongtae Ty Bae

Published 2026-04-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Dariush Lotfi, Mohammad-Ali Nikouei Mahani, Mohamad Koohi-Moghadam, Kyongtae Ty Bae

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The Overconfident AI Doctor

Imagine you have a brilliant AI doctor that has studied thousands of X-rays of adult brains. It is an expert at spotting tumors in adults. It has seen so many adult brains that it can diagnose them instantly.

But here is the danger: This AI doesn't know what it doesn't know.

If you show this AI a picture of a child's brain (which it has never seen) or a CT scan instead of an MRI, the AI might still look at it and say, "I am 99% sure this is a healthy adult brain!" It is confident, but it is completely wrong. In the real world, this could lead to a missed diagnosis or a dangerous error.

This is called the Out-of-Distribution (OOD) problem. The AI is encountering data that is "out of its distribution" (outside its training experience), but it fails to flag it as strange.

The Solution: A "Security Guard" That Doesn't Touch the Doctor

The researchers in this paper wanted to build a safety net. They needed a way to tell the AI, "Hey, this input looks weird to me, don't trust your answer!"

Usually, to fix an AI's confidence issues, you have to retrain the whole system. But in medicine, you can't just retrain a model whenever you want. These models are often locked down by strict regulations, and retraining them is expensive and risky.

So, the authors created a "Post-Hoc" (after-the-fact) Security Guard.

The Analogy: Imagine the AI Doctor is a chef who has cooked the same soup for 10 years. You don't want to fire the chef or teach them a new recipe (retraining). Instead, you hire a taste-tester (the Security Guard) who stands next to the chef. The chef serves the soup, and the taste-tester checks it. If the soup looks like it was made with the wrong ingredients (e.g., it's actually a smoothie), the taste-tester yells, "Stop! This isn't soup!"
The Magic: This taste-tester works without changing the chef's recipe or the chef's hands. It just watches the output and decides if it's safe.

How It Works: The "Normalizing Flow"

How does this taste-tester know what is "weird"?

The Feature Space (The "Vibe" Check):
Most old methods tried to look at the raw pixels (the individual dots of the image) to see if they looked different. This is like trying to spot a fake painting by counting the number of brushstrokes. It's easy to fool.
The new method looks at the features. This is like looking at the style of the painting. Is it Impressionist? Is it Abstract?
The AI Doctor extracts the "vibe" or "essence" of the image (high-level features). The Security Guard then checks this vibe.
The Normalizing Flow (The "Shape-Shifter"):
The researchers used a mathematical tool called a Normalizing Flow. Think of this as a magical, reversible machine that squishes and stretches data.
- The machine learns what "normal" adult brain vibes look like. It learns to squish all the "normal" vibes into a tight, neat cluster.
- When a weird image (like a child's brain) comes in, its "vibe" doesn't fit into that neat cluster. It gets stretched out or lands in a weird, empty space.
- The machine calculates a Likelihood Score. If the score is low, it means, "This vibe is very unlikely to be a normal adult brain."

The Results: A New Benchmark

The researchers didn't just test this on one thing. They built a new playground called MedOOD.

The Playground: They created 21 different types of "weird" scenarios, including:
- Motion Artifacts: Blurry images (like taking a photo while shaking).
- Population Shifts: Images of children or people from different parts of the world.
- Modality Shifts: Switching from MRI to CT scans.
- Organ Shifts: Showing an image of a liver instead of a brain.

The Score:
When they tested their "Security Guard" against other methods:

On the MedOOD dataset, it scored 84.6% accuracy in spotting the weird stuff, beating the previous best methods (which were around 80%).
On the MedMNIST (a standard medical image test), it scored 93.8%, crushing the competition.

Why This Matters for You

Safety First: This method acts as a "kill switch" for AI errors. If the AI sees something it shouldn't, the system can flag it for a human doctor to review, preventing mistakes.
No Re-training Needed: Hospitals can use this immediately. They don't need to buy new computers or retrain their expensive, regulated AI models. They just add this "Security Guard" module to the existing pipeline.
Feature vs. Pixel: The paper proved that looking at the "vibe" (features) is much better than looking at the raw pixels. It's like recognizing a person by their face rather than by the color of their shirt.

The Catch (Limitations)

The system isn't perfect. It struggles a little with subtle changes.

The Analogy: If you take a photo of an adult brain and just slightly lower the brightness or zoom in a tiny bit, the "vibe" is still very similar to the original. The Security Guard might get confused and think, "Hmm, this looks mostly normal."
However, for big, obvious changes (like a child's brain or a CT scan), it is incredibly effective.

Summary

This paper introduces a smart, non-invasive safety layer for medical AI. It acts like a vigilant security guard that checks if an image "feels right" based on what the AI has learned, without ever touching the AI's brain. This makes AI in hospitals safer, more reliable, and ready for the real world.

1. Problem Statement

In AI-driven medical imaging, deep learning models often fail to recognize Out-of-Distribution (OOD) data—inputs that differ significantly from the training distribution (e.g., different patient demographics, imaging modalities, or acquisition protocols). When encountering OOD data, these models frequently produce high-confidence but incorrect predictions, posing severe risks to clinical reliability and patient safety.

Existing OOD detection methods face two major barriers to clinical adoption:

Retraining Requirements: Many approaches require retraining the base model or modifying its architecture/weights, which is impractical in regulated clinical environments where pre-trained models are already approved.
Pixel-Space Limitations: Traditional density estimation methods operating directly on pixel space often fail to distinguish OOD samples due to sensitivity to low-level image statistics rather than semantic content.

2. Methodology

The authors propose a post-hoc OOD detection framework based on Normalizing Flows (specifically RealNVP) that operates in the feature space of pre-trained models without altering their weights.

Feature Extraction: Instead of using raw pixels, the method extracts high-level feature representations from a pre-trained base model (e.g., ResNet-18, ResNet3D-18). It aggregates features from multiple hierarchical levels (via average pooling and concatenation) to capture rich semantic information.
Normalizing Flow Model:
- The extracted feature vectors are fed into a Real Non-Volume Preserving (RealNVP) flow.
- The architecture consists of four masked affine coupling layers, each preceded by an ActNorm layer for stabilization.
- Each coupling layer uses Multilayer Perceptrons (MLPs) with 1024 hidden dimensions to model complex scale and translation functions.
- The model transforms input features into a latent space with a standard normal distribution, allowing for exact likelihood estimation.
Training & Inference:
- Unsupervised Training: The flow is trained only on In-Distribution (ID) samples to maximize the log-likelihood of ID features. No OOD examples are required during training.
- Post-Hoc Integration: The base classifier's weights remain frozen. The flow model acts as an auxiliary module that processes feature embeddings to generate an OOD score (likelihood). Low likelihood indicates an OOD sample.
Implementation: The system uses the Adam optimizer (learning rate $10^{-4}$ ) for 100 epochs. It is implemented in PyTorch using the OpenOOD framework.

3. Key Contributions

Post-Hoc Normalizing Flow Framework: A novel approach that integrates seamlessly with existing pre-trained medical imaging models without retraining or architectural changes, making it suitable for regulated clinical deployment.
MedOOD Dataset: Introduction of a curated benchmark dataset containing 21 OOD datasets grouped into five clinically relevant categories:
- Transformation Shifts: Artifacts (motion, ghost, noise), downsampling, scaling, etc.
- Population Shifts: Pediatric data, geographically distinct cohorts (e.g., Sub-Saharan Africa).
- Modality Shifts: Different MRI sequences (FLAIR, T1ce) and CT scans.
- Diagnostic Shifts: Different pathologies (MS, stroke, epilepsy).
- Organ Shifts: Non-brain MRIs (abdominal, spine).
Comprehensive Evaluation: Extensive experiments on both the new MedOOD dataset and the standard MedMNIST benchmark, demonstrating superior performance over state-of-the-art (SOTA) methods.

4. Experimental Results

The method was evaluated using standard metrics: AUROC, FPR@95, AUPR_IN, and AUPR_OUT.

Performance on MedOOD:
- Achieved an AUROC of 84.61%, outperforming SOTA post-hoc methods like ViM (80.65%) and MDS (80.87%).
- Achieved an FPR@95 of 67.74%, significantly lower (better) than ViM (76.77%) and MDS (76.13%).
- Statistical significance was confirmed via DeLong's test and bootstrap resampling ( $P < 0.05$ ).
- The method showed perfect discrimination for "Organ Shifts" and strong robustness for "Transformation Shifts," though it faced challenges with subtle "Population Shifts" (e.g., African cohort data).
Performance on MedMNIST:
- Achieved an AUROC of 93.8%, surpassing ViM (88.08%) and ReAct (87.05%).
- Attained the highest AUPR_IN (84.17%) and AUPR_OUT (97.9%).
Feature Space vs. Pixel Space:
- A critical ablation study demonstrated that operating in feature space is superior to pixel space.
- Feature Space AUROC: 93.8%
- Pixel Space AUROC: 70.59%
- Pixel space models failed to separate ID and OOD samples due to entanglement in t-SNE visualizations, whereas feature space showed clear semantic separation.
Efficiency:
- Inference time per sample is 17.9 ms, comparable to other post-hoc methods (e.g., ViM at 15.8 ms) and only slightly slower than lightweight score-based methods like MSP.
- Offline training cost is a one-time process (~128 seconds), making it practical for deployment.

5. Significance and Conclusion

This work addresses a critical gap in medical AI safety by providing a regulatory-compliant, post-hoc solution for OOD detection.

Clinical Relevance: By not requiring model retraining, it allows hospitals to deploy OOD safeguards on existing, approved AI systems immediately.
Robustness: The use of normalizing flows in feature space captures semantic distinctions better than pixel-based density estimation, effectively handling diverse distribution shifts (modalities, artifacts, demographics).
Future Impact: The introduction of MedOOD provides a standardized benchmark for future research. The authors suggest that while the method is robust, future work should focus on improving detection of subtle population shifts and integrating multi-scale feature representations.

The proposed framework serves as a practical "safety net" for clinical imaging workflows, capable of flagging unreliable inputs before they reach a clinician, thereby reducing the risk of diagnostic errors.