Deep Generative Classification of Blood Cell Morphology

TL;DR

CytoDiffusion is introduced, a diffusion-based generative classifier that faithfully models the distribution of blood cell morphology, combining accurate classification with robust anomaly detection, resistance to distributional shifts, interpretability, data efficiency and uncertainty quantification that surpasses clinical experts.

Abstract

Accurate classification of haematological cells is critical for diagnosing blood disorders, but presents significant challenges for machine automation owing to the complexity of cell morphology, heterogeneities of biological, pathological, and imaging characteristics, and the imbalance of cell type frequencies. We introduce CytoDiffusion, a diffusion-based classifier that effectively models blood cell morphology, combining accurate classification with robust anomaly detection, resistance to distributional shifts, interpretability, data efficiency, and superhuman uncertainty quantification. Our approach outperforms state-of-the-art discriminative models in anomaly detection (AUC 0.990 vs. 0.918), resistance to domain shifts (85.85% vs. 74.38% balanced accuracy), and performance in low-data regimes (95.88% vs. 94.95% balanced accuracy). Notably, our model generates synthetic blood cell images that are nearly indistinguishable from real images, as demonstrated by an authenticity test in which expert haematologists achieved only 52.3% accuracy (95% CI: [50.5%, 54.2%]) in distinguishing real from generated images. Furthermore, we enhance model explainability through the generation of directly interpretable counterfactual heatmaps. Our comprehensive evaluation framework, encompassing these multiple performance dimensions, establishes a new benchmark for medical image analysis in haematology, ultimately enabling improved diagnostic accuracy in clinical settings. Our code is available at https://github.com/CambridgeCIA/CytoDiffusion.

Figures (10)

• Figure 1: Representation of the diffusion-based classification process. An input image $\boldsymbol{x}_0$ is first encoded into a latent space using an encoder $\mathcal{E}$. Gaussian noise $\boldsymbol{\epsilon} \sim \mathcal{N}(0,\boldsymbol{I})$ is then added to create a noisy latent representation $\boldsymbol{z}_t$. This noisy representation is fed through a diffusion model for each possible class condition $c$. The model predicts the noise $\boldsymbol{\epsilon}_\theta$ for each condition. The classification decision is made by selecting the class that minimises the error between the predicted noise $\boldsymbol{\epsilon}_\theta$ and the true noise $\boldsymbol{\epsilon}$.
• Figure 2: Bayesian psychometric analysis of model and expert performance. A. Psychometric function of CytoDiffusion’s performance on our custom dataset with model confidence as the index of discriminability. The 95 per cent credibility interval of the threshold of the function, estimated at the point of 80 per cent correct unscaled by lapse and guess rate, is shown as an error bar, and the posterior distributions of threshold and width are presented in the inset axes. Note tight bounds on the parameters. B. Psychometric function of the performance of a single human expert, derived from a randomly selected subset of 200 cases for which a consensus ground truth is available, with model confidence as the index of discriminability. C. Psychometric function of the expert shown in B, with consensus expert confidence as the index of discriminability. D & E. Width and threshold parameter estimates and their 95 per cent credibility intervals for each of the six human experts with model confidence (D) and consensus expert confidence (E) as the indices of signal strength.
• Figure 3: Kernel density estimate figures comparing the anomaly detection performance of ViT-B/16 (top row) with CytoDiffusion (bottom row) for blasts (left column) and erythroblasts (right column). The horizontal axis represents the certainty measure, normalised to $[0, 1]$ by dividing by the maximum certainty for each model and dataset. The AUC values indicate the model's ability to distinguish between normal and abnormal samples.
• Figure 4: Model performance comparison under low-data conditions. Shaded areas represent the standard deviation from five independent training sessions.
• Figure 5: Counterfactual heatmap visualisation. Left: Original eosinophil image. Centre-right: Counterfactual heatmap ($\boldsymbol{H}_{\text{neutrophil}}$) illustrating the areas that differ when the model considers the image as a neutrophil. (For an example of how a neutrophil can look, see Figure \ref{['fig:heatmap_summary']}.) Far right: Overlay of the thresholded heatmap on the original image.