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

Riassunto Tecnico: Un Pipeline di Inferenza Ensemble Gerarchica per la Classificazione Robusta dei Globuli Bianchi sotto Domain Shift

Problema

La classificazione automatizzata dei globuli bianchi (WBC) è fondamentale per lo screening su larga scala della leucemia. Tuttavia, l'implementazione nel mondo reale è ostacolata dai cosiddetti "domain shifts" (spostamenti di dominio), causati da variazioni nei protocolli di colorazione, diverse caratteristiche degli scanner e variabilità tra i laboratori. Questi fattori degradano le prestazioni dei modelli di deep learning tradizionali. Inoltre, il compito è complicato da un forte sbilanciamento delle classi (class imbalance), dove le cellule critiche per la diagnosi (come i blasti) sono rare rispetto ad altre tipologie.

Metodologia

Gli autori propongono un framework di inferenza basato su un'architettura gerarchica e un sistema di ensemble, articolato in tre fasi principali:

1. Fine-tuning (Stage I): Utilizzano un backbone DinoBloom (un modello pre-addestrato tramite self-supervised learning) e lo adattano tramite LoRA (Low-Rank Adaptation). Per migliorare la robustezza, adottano uno schema teacher-student basato su EMA (Exponential Moving Average) e una funzione di perdita ibrida che combina la perdita DINO (self-supervised) e la cross-entropy bilanciata (supervised).
2. Costruzione del Feature Bank (Stage II): Una volta addestrato il modello, vengono estratti gli embedding $\ell_2$-normalizzati di tutte le immagini del database e memorizzati in un "feature bank" insieme alle loro etichette gerarchiche.
3. Inferenza kNN Gerarchica (Stage III): Invece di una classificazione end-to-end standard, il sistema esegue un'inferenza "coarse-to-fine" (dal generale al particolare). Per una nuova immagine, il sistema recupera i $k$-vicini più prossimi (kNN) tramite similarità del coseno e procede attraverso tre livelli di una gerarchia ematologica (Lineage $\rightarrow$ Sottotipo $\rightarrow$ Classe foglia). La previsione di ogni livello è vincolata dai risultati del livello precedente, garantendo coerenza biologica e riducendo gli errori tra classi diverse.

Per aumentare ulteriormente la robustezza, viene applicato un ensemble di più split (cross-validation) tramite voto di maggioranza.

Contributi Chiave

• Strategia di inferenza gerarchica: Un approccio che mitiga il bias verso le classi maggioritarie e riduce gli errori nelle classi rare sfruttando la struttura tassonomica delle cellule.
• Efficienza con LoRA: L'uso di LoRA per il fine-tuning di DinoBloom permette di ottenere rappresentazioni altamente trasferibili e robuste ai cambiamenti di dominio con un costo computazionale ridotto.
• Robustezza al Domain Shift: L'integrazione di kNN e ensemble di split permette al modello di gestire meglio le variazioni tra dataset diversi rispetto ai classificatori standard.

Risultati

Il metodo è stato valutato sul dataset WBCBench (ISBI 2026 Challenge):

• Il modello proposto (DinoBloom2) ha raggiunto un Macro F1-score di 0,682, posizionandosi tra i primi dieci nella fase di test finale.
• I risultati dimostrano che l'uso di DinoBloom con LoRA supera i modelli basati su ConvNeXt o ViT addestrati con fine-tuning completo.
• Studi di ablazione: L'analisi ha confermato che l'inferenza gerarchica fornisce un incremento di prestazioni molto più significativo rispetto al semplice ensemble di split, dimostrando che la coerenza gerarchica è cruciale per gestire lo sbilanciamento delle classi.

Significato

Questo lavoro rappresenta un passo avanti significativo verso l'automazione clinica affidabile. Dimostra che, per compiti medici critici dove i dati sono sbilanciati e i contesti operativi variano, un approccio che combina apprendimento self-supervised, adattamento efficiente dei parametri (LoRA) e vincoli biologici (gerarchia) è superiore ai modelli di classificazione pura, offrendo una maggiore stabilità e precisione diagnostica.