💻 computer science

A Hierarchical Ensemble Inference Pipeline for Robust White Blood Cell Classification Under Domain Shifts

The paper proposes a memory-augmented, hierarchical ensemble inference pipeline utilizing a DinoBloom backbone and k-nearest neighbors retrieval to achieve robust white blood cell classification under domain shifts.

Original authors: Ruyi Dai, Tingkwong Ng, Hao Chen

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ruyi Dai, Tingkwong Ng, Hao Chen

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 Problem: The "Blurry Photo" Dilemma in Blood Tests

Imagine you are a detective trying to identify rare, tiny insects in a massive forest. Most of the time, you see common ants or beetles. But occasionally, you spot a very rare, dangerous species that could change everything.

In medicine, doctors do something similar with blood. They look at White Blood Cells (WBCs) under a microscope to check for diseases like leukemia. However, there’s a huge problem: every "forest" is different. One lab might use a slightly different blue dye (staining), another might use a different microscope (scanner), and the lighting might change. To a computer, these small changes make the same cell look like a completely different species. This is called a "Domain Shift," and it often confuses even the smartest AI.

The Solution: The "Smart Library" Approach

The researchers created a new way for AI to classify these cells. Instead of just asking the AI, "What is this cell?" and hoping it gets it right, they built a three-step system that works like a highly organized library.

Step 1: The Expert Librarian (Fine-tuning with LoRA)

First, they took a very smart, pre-trained AI (called DinoBloom) and gave it "specialized training."

Think of this like taking a person who has read every book in the world and giving them a weekend crash course in Microscopic Entomology. They don't need to relearn how to read; they just need to learn the specific "slang" and details of blood cells. They used a technique called LoRA, which is like giving the AI a small, specialized notebook to write in, rather than forcing it to rewrite its entire massive brain. This makes the training fast and efficient.

Step 2: The Master Catalog (The Feature Bank)

Once the AI became an expert, they didn't just let it make guesses. They had the AI look at thousands of known cells and create a "Feature Bank."

Imagine a massive, high-tech library where every single cell ever seen is filed away. But instead of filing them by name, they are filed by "vibe" (mathematical embeddings). If a cell has a certain texture and a specific shade of purple, it gets placed in a very specific "shelf" in the library.

Step 3: The "Family Tree" Investigation (Hierarchical kNN)

This is the most clever part. When a new, unknown cell comes in, the AI doesn't just shout out a name. It uses a "Coarse-to-Fine" investigation, much like a detective narrowing down a suspect:

The Broad Search: The AI looks at its "Library" and finds the 10 most similar-looking cells (the "neighbors").
The Family Tree Check: Instead of just picking a name, it uses a hierarchy. It first asks: "Is this cell part of the 'Myeloid' family or the 'Lymphoid' family?"
The Narrowing Down: Once it decides it's a "Myeloid" cell, it ignores everything else in the library. It only looks at the "children" of that family. It says, "Okay, I know it's a Myeloid cell, so now I'm only looking for its specific cousins (like Neutrophils or Eosinophils)."

By following this "Family Tree," the AI is much less likely to make a wild, ridiculous mistake (like calling a tiny lymphocyte a giant blast cell).

The Result: A Robust Detective

By combining this Family Tree logic with a Library of examples, the researchers created an AI that is incredibly tough. Even if the "lighting" or "dye" changes (the domain shift), the AI can still look at its library, find the most similar "vibes," and use the hierarchy to stay on the right track.

In short: They didn't just build a better "eye"; they built a better "thinking process" that uses logic and memory to ensure that even the rarest, most critical cells are identified correctly.

Technical Summary: A Hierarchical Ensemble Inference Pipeline for Robust White Blood Cell Classification

1. Problem Statement

Automated White Blood Cell (WBC) classification is critical for large-scale leukemia screening. However, traditional deep learning models face significant challenges in real-world clinical deployment due to:

Domain Shifts: Variability in staining protocols, scanner characteristics, and inter-laboratory differences.
Class Imbalance: The clinical importance of rare subtypes (e.g., blast cells) is often undermined by the overwhelming prevalence of common cell types.
Generalization Issues: Standard CNN and Vision Transformer (ViT) models often struggle with patient-level generalization and robust performance across different datasets.

2. Methodology

The authors propose a three-stage hierarchical ensemble pipeline designed to mitigate domain shifts and class imbalance through retrieval-augmented inference.

Stage I: Fine-tuning (Feature Extraction)

Backbone: The authors utilize DinoBloom, a pretrained backbone, and adapt it using LoRA (Low-Rank Adaptation) to ensure parameter efficiency and transferable embeddings.
Training Objective: They employ a hybrid loss function combining Self-Supervised Learning (SSL) and Supervised Learning. Specifically, they use an EMA (Exponential Moving Average) teacher-student scheme with a DINO-style loss to align features, alongside a class-balanced cross-entropy loss to handle label imbalance.

Stage II: Feature Bank Construction

Instead of relying solely on a final classification layer, the model extracts $\ell_2$ -normalized embeddings from the fine-tuned backbone.
These embeddings, along with their corresponding three-level hierarchical labels (Lineage $\rightarrow$ Subtype $\rightarrow$ Leaf class), are stored in a Feature Bank.

Stage III: Hierarchical kNN Inference

Retrieval: For a query image, the system retrieves its $k$ -nearest neighbors (kNN) from the feature bank using cosine similarity.
Coarse-to-Fine Voting: Inference follows a hierarchical structure. It first predicts the top-level lineage (e.g., Myeloid, Lymphoid, or Blast) via majority voting.
Hierarchical Consistency: For subsequent levels, the search space is restricted only to the "children" of the previously predicted parent. This prevents the model from making biologically impossible errors (e.g., predicting a Lymphoid leaf under a Myeloid lineage), which significantly stabilizes predictions for rare classes.
Ensembling: The final output is stabilized using multi-split majority voting, aggregating predictions across different data splits to improve robustness.

3. Key Contributions

Hierarchical Inference Strategy: A novel coarse-to-fine kNN-based approach that enforces biological consistency and mitigates majority-class bias.
Efficient Adaptation: The use of LoRA fine-tuning on a DinoBloom backbone to produce high-quality, transferable embeddings for retrieval.
Robustness Framework: A combination of SSL-supervised hybrid training and multi-split ensembling to combat domain shifts and long-tailed class distributions.

4. Results

Performance: The proposed method achieved a Macro F1-score of 0.682 on the WBCBench dataset, ranking in the top ten.
Comparison: The DinoBloom + LoRA approach outperformed standard supervised baselines (ResNet-50, Swin-T, ViT-B) and even slightly surpassed full fine-tuning of larger models like ConvNeXt-L.
Ablation Insights:
- Hierarchy Matters: Enabling hierarchical constraints significantly boosted the Macro F1-score (from ~0.62 to ~0.68) compared to standard kNN.
- Ensemble Gains: Increasing the number of ensemble splits provided consistent, incremental improvements in robustness.

5. Significance

This research demonstrates that for highly specialized medical imaging tasks, retrieval-augmented hierarchical inference can be more robust than standard end-to-end classification. By leveraging biological taxonomies to constrain the decision space, the model becomes more reliable at identifying rare, clinically critical cells (like blasts) even when faced with the domain shifts common in real-world pathology labs. This paves the way for more generalizable "foundation models" in hematology.