Note

The code for this project can be found on GitHub. The code is fully documented. Some images may take a while to load on this webpage.

Introduction

This post describes a Trustworthy Machine Learning approach to medical image diagnostics, built using Bayesian Deep Learning. Designed specifically with the UK's National Health Service (NHS) in mind, the system processes Optical Coherence Tomography (OCT) scans to classify Age-Related Macular Degeneration (AMD) while simultaneously estimating its own uncertainty.

With an ageing UK population, ophthalmology clinics—particularly those at major centres like Moorfields Eye Hospital—are facing unprecedented backlogs. Automating scan triage is essential, but clinical adoption requires systems that know when they don't know. By replacing standard deterministic network heads with Bayesian layers and applying Variational Inference using Pyro, our model quantifies both the uncertainty inherent in the scan (e.g. poor image quality) and its own internal doubt.

The model uses a Vision Transformer (ViT) backbone pretrained on ImageNet, fine-tuned on an AMD dataset (normal, early AMD, late AMD). By explicitly decoupling uncertainty into aleatoric (data noise) and epistemic (model ignorance) components, the system addresses the notorious "overconfidence" issue common in standard deep neural networks operating on out-of-distribution (OOD) medical data. This capability is critical for safe NHS deployment: if the model flags a scan as "highly uncertain," it is automatically routed to a human consultant, preventing dangerously confident but incorrect automated diagnoses.

Background

The Overconfidence Problem in Deep Learning

Standard neural networks represent weights as single point estimates. When presented with anomalous data or images fundamentally different from their training distribution (Out-of-Distribution or OOD), standard models often produce high-confidence softmax outputs. This occurs because softmax simply squashes arbitrary logits into a

1[0, 1]

Bayesian Neural Networks (BNNs)

A Bayesian Neural Network treats its weights $w$ not as fixed values, but as probability distributions $P(w)$. Instead of finding a single optimal set of weights, learning involves finding the posterior distribution of the weights given the training data $P(w|D)$. Because calculating the exact posterior is computationally intractable for deep networks, we use Variational Inference (VI) to approximate it using a simpler distribution $q_\theta(w)$.

The network is trained by maximising the Evidence Lower Bound (ELBO), which balances the data likelihood against the Kullback-Leibler (KL) divergence between the approximate posterior and a prior distribution $P(w)$:

Decomposing Uncertainty

Uncertainty in predictions arises from two distinct sources:

1. Aleatoric Uncertainty (Data Uncertainty): Arises from inherent noise in the observations (e.g., sensor noise, poor image quality, ambiguous pathologies). It cannot be reduced by collecting more training data. We capture this by training the network to directly predict the variance of its outputs ($\sigma^2$) using a Heteroscedastic Loss function.
2. Epistemic Uncertainty (Model Uncertainty): Arises from a lack of knowledge about the best model parameters, often due to sparse training data in certain regions of the feature space. This can be reduced with more data. We capture this using Monte Carlo (MC) Dropout at inference time, measuring the variance in predictions across multiple stochastic forward passes.

Implementation / Methodology

Data Pipeline

The system is designed to train on the Moorfields AMDP (Age-Related Macular Degeneration Patient) Dataset, a world-leading collection of OCT scans from the UK. To ensure the code can run anywhere, the pipeline features an intelligent fallback mechanism: if Kaggle authentication is unavailable, it automatically generates a synthetic, noisy OCT-like dataset with simulated pathologies (drusen and fluid).

To rigorously test uncertainty quantification, the model is evaluated across three distinct regimes:

• Standard Test Set: Clean OCT scans (normal, early AMD, late AMD) from the primary dataset.
• Shifted (Noisy) Test Set: The standard test set aggressively corrupted with random rotations and severe Gaussian noise. This simulates low-quality clinic sensors and tests the model's Aleatoric Uncertainty.
• Out-of-Distribution (OOD): CIFAR-10 natural images (e.g., aeroplanes, dogs), which the model has never seen and which contain no AMD pathology. This perfectly tests Epistemic Uncertainty—verifying that the model registers high doubt when fundamentally confused.

1# Shifted transforms for evaluating OOD / Aleatoric uncertainty
2shifted_transform = transforms.Compose([
3    transforms.Resize((IMG_SIZE, IMG_SIZE)),
4    transforms.RandomRotation(90),
5    transforms.ToTensor(),
6    AddGaussianNoise(0., 0.5),
7    transforms.Normalize(...)
8])

Bayesian ViT Architecture

The architecture leverages a hybrid approach to combine the strong feature extraction of Vision Transformers with the probabilistic rigour of BNNs:

1vit_small_patch16_224

1pyro-ppl

The classification head consists of physical linear layers where weights and biases are stochastic

1PyroSample

Heteroscedastic Classification Loss

To train the aleatoric variance head, we use a bespoke Heteroscedastic Classification Loss. Rather than computing cross-entropy on deterministic logits, we treat the predicted logits as a Gaussian distribution $\mathcal{N}(\mu, \sigma^2)$. We approximate the expected cross-entropy over this distribution using Monte Carlo sampling via the reparameterization trick:

1# Reparameterization: z = mu + sigma * epsilon
2expected_loss = 0.0
3for _ in range(num_samples):
4    epsilon = torch.randn_like(pred_mu)
5    sampled_logits = pred_mu + std * epsilon
6    expected_loss += cross_entropy(sampled_logits, target)
7loss = expected_loss / num_samples

MC Dropout Inference

During evaluation, passing an image through the network $T$ times yields a distribution of predictions. We bypass standard

1.eval()

Analysis

In our prototype evaluation workflow, the model demonstrated the full end-to-end pipeline, producing metrics across the three evaluation regimes.

(Note: The metrics below reflect a rapid diagnostic training run. In a fully converged production model with a massive medical dataset, absolute accuracy would be significantly higher, but the relative uncertainty patterns remain the primary focus.)

Expected Calibration Error (ECE)

A perfectly calibrated model possesses an Expected Calibration Error (ECE) of 0.0, meaning its predicted confidence matches its empirical accuracy. Our Bayesian ViT achieved an ECE of 0.0513 (or 5.1%), indicating strong calibration even during rapid diagnostic runs. This was visualised using a Reliability Diagram, proving that the model aligns tightly with the ideal identity line.

Uncertainty Decomposition vs. Data Corruption

The core thesis of the project is that uncertainty should behave predictably when input data deviates from the training manifold:

1. Aleatoric Uncertainty represents noise in the input itself.
2. Epistemic Uncertainty represents model ignorance.

By rigorously evaluating the model on Standard, Noisy, and OOD data, we plot the distributions of both uncertainty metrics over the test cohorts. The analysis pipeline successfully generates grouped histograms and class-wise correlation matrices, paving the way for adaptive, uncertainty-aware thresholding in clinical triage queues.

OOD Detection

Using strictly the predicted epistemic uncertainty, we compute the Area Under the Receiver Operating Characteristic curve (AUROC) to distinguish between in-distribution (Medical) and out-of-distribution (CIFAR-10) scans. The pipeline dynamically calculates this separating power, quantifying exactly how well the model "knows what it doesn't know."

Conclusion

We have successfully implemented a comprehensive framework for Uncertainty-Aware Medical Diagnostics using Bayesian Deep Learning. By augmenting a Vision Transformer with Pyro-based Bayesian linear layers and deploying a Heteroscedastic Classification Loss, we established a pipeline capable of disentangling aleatoric and epistemic uncertainty.

Our evaluation suite rigorously tests the model's calibration (ECE), out-of-distribution detection capabilities (AUROC), and classification robustness under severe data shifts. While standard deterministic models suffer from silent, overconfident failures on corrupted clinical data, this bayesian architecture explicitly quantifies its doubt.

The implications for the UK software landscape are substantial. As the NHS digitises to manage immense backlogs, completely autonomous AI is often deemed too risky for frontline diagnostics. This Bayesian approach bridges the gap: it allows safe automation of routine scans while reliably escalating ambiguous or anomalous cases to human specialists. Future work will focus on scaling the model to the full high-resolution Moorfields AMDP cohort, and integrating these uncertainty-aware outputs directly into automated NHS clinical triage routing systems.