🤖 machine learning

Pathology Transport: Optimal-Transport Explanations for Clinical Data, and When Their Heatmaps (Fail to) Localize Disease

This paper introduces an optimal-transport-based generative framework for explaining clinical AI that successfully generates unsupervised malignancy scores and tabular counterfactuals, but critically reveals that while its heatmaps can localize synthetic lesions, they fail to accurately identify real disease regions in medical imaging, highlighting a significant gap between synthetic benchmarks and clinical reality.

Original authors: Lalit Kumar

Published 2026-08-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Lalit Kumar

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

In the world of medical AI, a computer program that predicts whether a patient has a disease is only half the story. Doctors need to know why the machine reached that conclusion. They need to see which specific measurements tipped the scales or which part of an X-ray image looks suspicious. For years, researchers have tried to build these explanations by poking at a finished prediction model, asking it to justify its choices. But a different approach asks a bolder question: what if we could map the journey from health to sickness directly, without ever needing a prediction model to begin with? This idea treats a group of healthy patients and a group of sick patients as two distinct clouds of data points. The goal is to draw a smooth, logical path that connects these two clouds, showing exactly how a healthy body changes to become a diseased one. If such a map exists, it could reveal the precise biological shifts that define an illness, offering a clear, visual explanation of disease that feels natural and grounded in reality.

A researcher set out to test this concept using a mathematical framework known as optimal transport, which essentially finds the most efficient way to move one group of things into the shape of another. They built a system that learns this movement by studying real medical data, first using a dataset of breast cancer biopsies and then moving to chest X-rays. In the first experiment, they fed the computer information about thirty different measurements taken from cell nuclei, such as their size, shape, and texture. The system learned to transport the data points representing healthy cells into the space occupied by malignant cells. The result was a single model that could do three things at once. It could take a specific patient's data and show what that patient would look like if they were sick, highlighting exactly which measurements changed. It could assign a risk score to any patient based on how far their data had to travel to reach the "sick" zone. And it could produce a general map showing which biological features are most important for the disease, without the system ever being told which patients were actually sick during its training.

The results on the breast cancer data were encouraging. The system generated a risk score that correctly identified malignant cases about ninety-one percent of the time, a strong performance for a method that never saw a diagnosis label. More importantly, the explanations it produced made sense to human experts. When the system showed how a healthy cell became cancerous, it pointed to the exact same biological features that pathologists look for, such as irregular nuclear shapes and larger cell sizes. The system did not beat a standard, supervised computer program at predicting the outcome, but it succeeded in creating a clear, navigable path between health and disease that a standard program cannot provide. It offered a way to visualize the disease process itself rather than just a final verdict.

However, the story took a sharp turn when the researcher applied the same method to chest X-rays, trying to see if the system could pinpoint the exact location of pneumonia in a lung image. They trained the model to move images of healthy lungs toward the shape of lungs with pneumonia. When they looked at the resulting heatmaps, which were supposed to highlight the diseased areas, the system failed to localize the problem. Instead of focusing on the specific spot where the infection was, the heatmap spread out across the entire lung, acting more like a general signal that "something is wrong" rather than a precise map of where the trouble lies. To understand why this happened, the researcher created a controlled test by planting a fake, artificial spot of disease into a healthy lung image. On these synthetic, perfect spots, the system's heatmap did work, successfully pointing to the fake lesion.

This success on fake data led to a critical discovery about the limits of the technology. When the researcher tested the system on real X-rays from a large medical challenge, where doctors had drawn actual boxes around the pneumonia, the system's performance collapsed. The heatmaps that looked so promising on the artificial spots failed to find the real disease, performing no better than random chance. The researcher found that the system was reacting to the normal, complex variations in human anatomy, such as the edges of ribs or the curve of the diaphragm, rather than the actual infection. The lesson was stark: a heatmap that looks convincing on a simple, artificial test does not prove the system can find real disease in a complex, messy human body. While the method successfully created a beautiful and logical map of how health transforms into disease for simple data, it has not yet solved the harder problem of pinpointing specific diseases in real medical images without human guidance. The work stands as a powerful demonstration of how to build explainable AI, but also as a necessary warning that what looks like a solution in a controlled simulation may not hold up in the real world.

Technical Summary: Pathology Transport

Problem Statement
The adoption of predictive clinical AI is often hindered not by accuracy, but by a lack of trust and actionability. Standard explainable AI (XAI) methods like SHAP provide feature importance but fail to generate concrete, on-distribution counterfactual examples (i.e., "what would the patient look like if they were healthy?"). Furthermore, while generative models can produce counterfactuals, they often rely on external classifiers for guidance. This paper addresses two core questions: (1) Can a single generative model, trained without a classifier, learn the geometry separating healthy and diseased populations to yield counterfactuals, risk scores, and attributions? (2) Do the resulting generative heatmaps genuinely localize disease pathology in medical images, or do they merely reflect population-level shifts?

Methodology
The author proposes reframing clinical diagnosis as an Optimal Transport (OT) problem between two probability distributions: benign (healthy) and malignant (diseased). The core engine is an OT-coupled Rectified Flow.

Optimal Transport Coupling: Instead of random pairing of source (benign) and target (malignant) samples, the method uses mini-batch optimal transport (via the Hungarian algorithm) to pair samples. This minimizes the transport cost, ensuring trajectories are nearly straight and stable.
Rectified Flow: A velocity field $v_\theta(x, t)$ is learned via regression to transport samples from the benign distribution ( $t=0$ ) to the malignant distribution ( $t=1$ ). The model is trained on straight-line interpolations between the OT-paired endpoints.
Three Interpretability Artifacts:
- Risk Score: The magnitude of the transport displacement ( $\|x - F(x)\|$ ) serves as an unsupervised malignancy score.
- Counterfactuals: Integrating the flow generates the nearest "malignant" version of a benign patient (or vice versa), identifying specific biomarker changes required to flip the diagnosis.
- Attribution: The average displacement vector across a cohort provides a global attribution map of the transition from health to disease.

The method was evaluated on two modalities:

Tabular Data: The Breast Cancer Wisconsin (Diagnostic) dataset (569 biopsies, 30 nuclear biomarkers).
Imaging Data: Chest X-rays (PneumoniaMNIST and RSNA Pneumonia Detection Challenge).

Key Results

Tabular Performance:
- Risk Score: The unsupervised transport score achieved an AUROC of 0.91 (0.93 ± 0.01 across seeds), outperforming a simple distance baseline (0.88) but not surpassing a supervised logistic regression (0.99).
- Counterfactuals: The flow generated valid counterfactuals (84% validity) that were individualized (edit directions varied per patient, cosine similarity 0.64), unlike a simple mean-shift baseline which applied identical edits to all. The generated edits aligned with clinical intuition (e.g., increasing nuclear concavity and perimeter).
- Attribution: The label-free population attribution correlated significantly ( $r \approx 0.5$ ) with coefficients from a supervised logistic regression, correctly identifying malignancy drivers without seeing labels.
- Critical Finding: The OT coupling was decisive; random pairing resulted in off-manifold artifacts and anti-correlated attribution ( $r = -0.03$ ), whereas OT pairing yielded coherent signals.
Imaging Performance & Localization:
- Synthetic Lesions: On a controlled benchmark with planted Gaussian lesions, a population-level transport heatmap failed to localize the lesion (Pointing Game: 0.17, near random). However, a label-free, identity-preserving autoencoder (reconstruction-based) successfully localized synthetic lesions (Pointing Game: 0.52 at 128×128 resolution), outperforming the population flow.
- Real Pathology (The Synthetic-to-Real Gap): When tested on real RSNA radiologist-annotated pneumonia boxes, the label-free autoencoder collapsed to chance performance (Pointing Game: 0.11). The reconstruction error was dominated by normal anatomical variations (ribs, diaphragm) rather than the disease. Only the supervised Grad-CAM remained above chance, though weakly.
- Conclusion on Heatmaps: Label-free heatmaps that appear compelling on synthetic anomalies do not necessarily transfer to real pathology localization.

Contributions and Significance
The paper claims to provide a "honest" and "modest" contribution rather than a state-of-the-art performance boost:

A Reusable Recipe: It demonstrates that a single OT-coupled rectified flow can serve as a complete interpretability engine, generating counterfactuals, risk scores, and attributions without requiring a pre-trained classifier.
Validation of Geometry: It shows that on tabular data, the geometry of the transport map captures clinically relevant structure, recovering known biomarker signatures without labels.
The Synthetic-to-Real Gap: The most significant finding is a cautionary result regarding medical imaging. The author demonstrates that generative heatmaps which successfully localize synthetic lesions often fail to localize real pathology. This gap suggests that current label-free reconstruction methods are insufficient for reliable disease localization in complex real-world scenarios.
Honest Benchmarking: The paper explicitly reports where the method fails (e.g., it does not out-predict logistic regression on tabular data; it fails to localize real pneumonia), framing the work as a controlled investigation into the limits of generative explanations rather than a diagnostic tool.

Limitations and Ethical Considerations
The author emphasizes that the study is a controlled investigation, not a clinical validation. The datasets are small and demographically narrow. The unsupervised scores are uncalibrated and should not be used to gate clinical care. The synthesized "malignant twins" are hypothesis-generating aids, not medical advice, and could be misleading without context. The paper concludes that while the method offers a navigable path between health and disease, it does not yet replace supervised classifiers for prediction or reliable localization in real-world settings.

Technical Summary: Pathology Transport

More like this