Drift-Aware Medical AI
Detects and explains distribution shift in a chest X-ray classifier — with no new labels required.
- Status
- Research / MLOps pipeline — not a clinical tool
- Role
- Solo — MSc dissertation: pipeline, agent, evaluation
- Timeline
- 2026
At a glance
A monitoring pipeline that watches a multi-label chest X-ray model for distribution shift over time, using nine statistical two-sample tests across pixels, embeddings, and predictions. An autonomous agent chains detection to root-cause analysis to a plain-language explanation and a concrete engineering recommendation.
- Nine two-sample tests run on pixel statistics, learned embeddings, and model predictions — no new annotations needed.
- Root-cause analysis ranks metadata and image features by effect size and correctly surfaces injected corruptions.
- The agent produces a human-readable explanation of what shifted and quantifies the impact on calibration.
- Per-disease temperature scaling keeps confidence scores meaningful as the input distribution moves.
- Trained on CheXpert, stress-tested against NIH ChestX-ray14 for genuine cross-source drift.
Key numbers
AUROC ≈ 0.91
Best detector
Single test; ensembling did not improve it
0.877 → 0.841
Cross-source
Macro AUROC, CheXpert → NIH ChestX-ray14
6.6 ms mean
Streaming latency
p95 11.1 ms per batch
Per-disease
Calibration
Beats a single global temperature
Tech stack
- Modelling
- PythonPyTorchDenseNet-style backboneMulti-label classification
- Drift detection
- MMDWassersteinC2STBBSDEntropy tests9 tests total
- Agent
- Detection → analysis → explanation → recommendationRule-based reportingOptional Phi-3
- Serving
- FastAPIKafka streamingpytestContainerised service
The problem
Medical imaging models degrade quietly. A classifier that performs well at release drifts as scanners, protocols, and patient populations change — and in healthcare, silent failure is the most dangerous kind.
The catch: you rarely get fresh labels in production. Drift detection has to work on unlabelled data.
Approach
A backbone model performs multi-label chest X-ray classification. Around it, a monitoring layer runs nine statistical two-sample tests — MMD, Wasserstein distance, a classifier two-sample test, black-box shift detection, entropy-based tests and others — across three representations: raw pixels, learned embeddings, and the model's own output distribution.
When drift is flagged, root-cause analysis ranks candidate drivers (metadata fields, image features, injected corruptions) by effect size, so the output is not just 'something changed' but 'this is what changed and by how much'.
The agent loop
An autonomous agent chains four steps: detect drift, analyse its cause, explain it in plain language, and recommend an engineering response — from temperature re-scaling to flagging a batch for retraining. Explanations are rule-based by default, with an optional small language model (Phi-3) for narrative reporting.
Results
The strongest single detector reached AUROC ≈ 0.91; combining detectors into an ensemble did not help. Root-cause attribution reliably identified injected corruptions as the top drivers. Cross-dataset evaluation showed macro AUROC moving from 0.877 on CheXpert to 0.841 on NIH ChestX-ray14 — a measurable, explainable drop.
The streaming service processes batches at 6.6 ms mean latency (p95 11.1 ms). Per-disease temperature scaling consistently beat a single global temperature for calibration.
Scope
This is a research and monitoring pipeline for an ML operations team, not a clinical decision-support tool. The point is reliability engineering: knowing when a deployed model can no longer be trusted, and why.
Want the details?
The full source, READMEs, evaluation harnesses, and results are on GitHub.
github.com/dariyaDeepanshu/DriftModel