Sintesi Tecnica: Pre-addestramento di Reti di Multiple Instance Learning con Distillazione Multi-Teacher da Modelli Foundation di Patologia su Slide

Definizione del Problema

Il Multiple Instance Learning (MIL) si è affermato come il paradigma dominante per l'analisi delle Whole Slide Images (WSI) nella patologia computazionale. In questo framework, una slide viene trattata come un "sacchetto" (bag) di istanze a livello di patch, dove un encoder di tile congelato estrae le feature e un aggregatore MIL apprende a combinare queste istanze in una predizione a livello di slide. Nonostante la diversità delle architetture degli aggregatori (ad es. basati su attenzione, Transformer, grafi, modelli state-space), persiste un limite fondamentale: gli attuali aggregatori MIL vengono tipicamente addestrati da zero per ogni task a valle.

Questo approccio "da zero" si affida a etichette limitate a livello di slide per apprendere simultaneamente rappresentazioni distintive delle istanze e strategie di aggregazione efficaci. Di conseguenza, questi modelli soffrono spesso di ottimizzazione instabile, overfitting su correlazioni spurie e scarsa trasferibilità tra i task. Mentre l'analisi delle immagini naturali ha adottato con successo il paradigma pre-addestra-e-fine-tuna utilizzando dataset su larga scala (ad es. ResNet, ViT), applicare questo approccio al MIL in patologia incontra due barriere critiche:

1. Scarsità di Dati: A differenza delle immagini a livello di tile che possono essere ritagliate in milioni, i dataset di alta qualità disponibili pubblicamente a livello di slide sono limitati a miglia di o decine di migliaia di campioni, insufficienti per addestrare modelli foundation su larga scala da zero.
2. Capacità del Modello: Gli aggregatori MIL convenzionali sono moduli leggeri e specifici per il task con pochi parametri. Mancano della capacità di beneficiare del convenzionale pre-addestramento supervisionato nello stesso modo dei grandi encoder su larga scala, e il fine-tuning diretto di grandi encoder di slide per diversi task clinici può essere computazionalmente costoso e instabile nell'ottimizzazione.

Metodologia

Gli autori propongono un framework di distillazione multi-teacher per pre-addestrare aggregatori MIL leggeri trasferendo la conoscenza da modelli foundation di patologia su slide pre-addestrati. Il cuore del metodo consiste nell'utilizzare due modelli encoder di slide all'avanguardia, TITAN e CARE, come modelli "teacher" complementari per guidare l'addestramento di varie architetture MIL "student".

1. Configurazione della Distillazione Multi-Teacher

• Teachers:
  • TITAN: Fornisce semantica whole-slide generalista e trasferibile appresa da un pre-addestramento multimodale su larga scala.
  • CARE: Offre rappresentazioni regionali informate morfologicamente e molecolarmente.
  • La combinazione di questi teacher fornisce segnali di supervisione diversificati, riducendo il bias inerente all'uso di una singola rappresentazione.
• Students: Nove diverse architetture MIL (tra cui ABMIL, CLAM, TransMIL, WiKG, DAGMIL, GDFMIL, 2DMamba, AMDMIL e AEMMIL) fungono da studenti. I modelli student sono inizializzati con pesi casuali e addestrati per imitare gli embedding a livello di slide prodotti dai teacher congelati.
• Costruzione dei Dati: Per evitare l'irregolarità computazionale delle intere WSI, il pre-addestramento viene condotto su Regioni di Interesse (ROI) ad alta risoluzione dal dataset TCGA-UT-8K. Ogni ROI è suddivisa in una griglia regolare di $16 \times 16$ patch, garantendo dimensioni dei batch costanti e un consumo di memoria stabile.

2. Perdita di Distillazione Normalizzata per Dispersione Angolare (ADN Loss)

Una sfida centrale nella distillazione multi-teacher è che diversi teacher possono esibire livelli variabili di concentrazione nelle loro distribuzioni di rappresentazione sull'ipersfera. Se un teacher produce embedding strettamente raggruppati mentre un altro è più disperso, i loro segnali di supervisione possono avere scale differenti, portando a un'ottimizzazione sbilanciata.

Per affrontare questo problema, gli autori introducono la Angular Dispersion Normalized (ADN) Distillation Loss.

• Meccanismo: La perdita misura la distanza angolare tra la predizione dello studente e il target del teacher. Fondamentalmente, questo errore è normalizzato dalla dispersione angolare corrente ($\sigma_k$) della distribuzione di rappresentazione del particolare teacher.
• Formula: Per un teacher $k$, la perdita è calcolata come:
  $$L^{(k)}_{distill} = \frac{1}{B} \sum_{i=1}^{B} \frac{\Theta^{(k)}_i}{\sigma_k + \epsilon}$$
  dove $\Theta^{(k)}_i$ è la distanza angolare e $\sigma_k$ è la deviazione standard stimata degli embedding del teacher attorno al loro centro di aggiornamento.
• Effetto: Questa normalizzazione assicura che un teacher con embedding strettamente raggruppati (piccolo $\sigma_k$) non domini gli aggiornamenti del gradiente semplicemente perché la sua magnitudo di errore grezzo è piccola. Bilancia l'influenza dei teacher con diverse caratteristiche di rappresentazione.

3. Adattamento a Valle (Downstream Adaptation)

Dopo il pre-addestramento, le teste di proiezione specifiche per il teacher vengono scartate. I pesi dell'aggregatore MIL pre-addestrato ($\theta^*$) servono come inizializzazione riutilizzabile per i task a valle. Gli autori valutano due protocolli di adattamento:

• Full-Parameter Fine-tuning: Sia l'aggregatore MIL che la testa di classificazione specifica per il task vengono aggiornati utilizzando le etichette a livello di slide del task a valle.
• Linear Probing: L'aggregatore MIL pre-addestrato è congelato e viene addestrata solo la testa di classificazione lineare, valutando direttamente la qualità delle rappresentazioni apprese a livello di slide.

Contributi Chiave

1. Pre-addestramento MIL Fattibile: Il paper dimostra un percorso pratico per pre-addestrare aggregatori MIL leggeri nonostante la scarsità di dati etichettati su larga scala a livello di slide, sfruttando la distillazione della conoscenza dai modelli foundation esistenti.
2. Framework Multi-Teacher: Introduce un framework che utilizza la conoscenza complementare di TITAN e CARE, dimostrando che combinare teacher eterogenei può produrre una generalizzazione migliore rispetto alla distillazione con un singolo teacher.
3. ADN Loss: La proposta della Angular Dispersion Normalized Loss affronta la sfida specifica di bilanciare i segnali di supervisione da teacher con diverse distribuzioni di rappresentazione.
4. Valutazione Completa: Il metodo è valutato sistematicamente attraverso 15 task a valle su 6 dataset di benchmark, coprendo sottotipizzazione del cancro, predizione di biomarcatori e rilevamento di mutazioni.

Risultati Sperimentali

Gli autori hanno condotto estesi esperimenti confrontando i loro modelli MIL con inizializzazione tramite distillazione contro:

• Modelli addestrati da zero.
• Fine-tuning diretto dei modelli teacher (TITAN e CARE).

Risultati Chiave:

• Guadagni di Performance: L'inizializzazione pre-addestrata migliora costantemente gli aggregatori MIL rispetto all'addestramento da zero. Ciò è particolarmente evidente nei casi di linear probing e few-shot.
• Superamento dei Teacher: In molti casi, i leggeri aggregatori MIL con inizializzazione distillata superano i pesanti encoder di slide dei teacher (TITAN e CARE) quando adattati ai task a valle, suggerendo che la conoscenza distillata è efficacemente compressa e specializzata per il task di aggregazione.
• Efficienza Few-Shot: Nei setting few-shot (1 a 16 shot per classe), i modelli inizializzati tramite distillazione mostrano guadagni significativi (ad es. fino a 20,55 punti percentuali su BRACS-Coarse) rispetto all'addestramento da zero, evidenziando una migliore efficienza dei dati e stabilità dell'ottimizzazione.
• Agnosticismo dell'Architettura: I benefici sono stati osservati in varie architetture (attenzione, Transformer, grafi, state-space), sebbene alcune architetture (come GDFMIL e 2DMamba) abbiano mostrato performance più variabili a seconda del task e del protocollo di adattamento.
• Efficienza della Memoria: A differenza dei grandi modelli foundation che soffrono di esplosione della memoria all'aumentare del numero di patch, i leggeri aggregatori MIL mantengono un uso stabile della memoria anche con un gran numero di patch (ad es. 30k patch), rendendoli scalabili per l'analisi di WSI gigapixel.

Significato e Rivendicazioni

Il paper sostiene che questo lavoro fornisce un'alternativa pratica all'attuale addestramento MIL da zero. Trasferendo i priori a livello di slide dai modelli foundation di grandi dimensioni agli aggregatori leggeri tramite distillazione, l'approccio consente:

• Priori a Livello di Slide Trasferibili: I modelli MIL possono acquisire rappresentazioni stabili e generalizzabili senza richiedere massicci dataset etichettati per il proprio pre-addestramento.
• Efficienza: I modelli risultanti rimangono computazionalmente efficienti e scalabili, adatti ad ambienti clinici con risorse computazionali o budget di annotazione limitati.
• Robustezza: Il metodo migliora la velocità di convergenza e riduce l'overfitting, specialmente nei regimi di scarsità di dati come il few-shot learning.

Gli autori posizionano questo lavoro come un passo verso la trasformazione della patologia computazionale, rendendo il paradigma "pre-addestra-e-fine-tuna" accessibile al MIL, una componente critica dell'analisi WSI che è stata precedentemente limitata all'addestramento da zero. Si suggerisce come lavoro futuro la scalabilità del pre-addestramento dei dati e la valutazione del framework su un intervallo ancora più ampio di task a valle.