MAGE: Color-Invariant and Spatial Knowledge Distillation for Gastric Neoplasm Classification

Sintesi Tecnica: MAGE – Distillazione della Conoscenza Spaziale e dell'Invarianza al Colore per la Classificazione delle Neoplasie Gastriche

Definizione del Problema
La differenziazione accurata tra adenoma e carcinoma gastrico durante l'endoscopia è critica per il processo decisionale clinico, ma rimane estremamente complessa. Il compito è ostacolato dall'elevata somiglianza inter-classe, dai confini ambigui e dalla tendenza delle reti neurali profonde a fare affidamento su correlazioni spurie (ad esempio, sanguinamento, artefatti di illuminazione o contesto di sfondo) piuttosto che su caratteristiche morfologiche intrinseche. Gli approcci esistenti presentano limitazioni specifiche:

• Metodi basati su ROI: Spesso soffrono di propagazione dell'errore dovuta a una rilevazione o segmentazione imperfetta e perdono il contesto globale circostante.
• Classificazione dell'immagine completa: Manca del necessario focus spaziale e spesso predice basandosi sulle regioni di sfondo piuttosto che sulla lesione stessa.
• Dipendenza dal colore: I modelli tendono frequentemente verso bias di texture specifici del dominio, dove il colore agisce come un confonditore fragile piuttosto che come un indicatore diagnostico primario.

Metodologia: Framework MAGE
Gli autori propongono MAGE (Masked Achromatic Guidance Expert), un framework di addestramento a due rami progettato per imporre un apprendimento morfologico invariante al colore, mantenendo al contempo la capacità di implementazione su immagini complete senza l'ausilio di maschere annotate durante l'inferenza.

1. Esperto Acromatico Mascherato (Ramo Ausiliario):

   • Durante l'addestramento, un ramo "esperto" ausiliario elabora una vista acromatica mascherata dell'input.
   • L'input viene costruito convertendo l'immagine in luce bianca (WLI) in scala di grigi e applicando una maschera di segmentazione ground-truth ($\tilde{x} = m \odot \text{Gray}(x)$), impostando a zero i pixel al di fuori della Regione di Interesse (ROI).
   • Questo ramo è addestrato esclusivamente su caratteristiche strutturali e invarianti al colore, costringendo il modello a imparare pattern discriminativi basati sulla morfologia piuttosto che su indizi cromatici o artefatti di sfondo.
   • Questo ramo utilizza informazioni privilegiate (maschere) ed è congelato durante il processo di distillazione.

2. Modello Principale (Ramo Implementabile):

   • Il ramo principale riceve l'intera immagine WLI grezza (RGB) come input ed è l'unico componente utilizzato durante l'inferenza.
   • Impara a imitare l'esperto attraverso una strategia di distillazione a doppio obiettivo:
     • Distillazione dei Logit di Classificazione: I logit in uscita del modello principale sono regolarizzati tramite divergenza KL rispetto ai logit dell'esperto, trasferendo la fiducia di classificazione dell'esperto.
     • Trasferimento dell'Attenzione Spaziale: Il modello principale è regolarizzato per allineare la propria attenzione delle mappe di feature con quella dell'esperto. Comprimendo la dimensione dei canali e calcolando la distanza di Frobenius al quadrato tra la mappa di attenzione del modello principale e quella dell'esperto (che è vincolata alla ROI), il modello principale è implicitamente guidato a concentrarsi sulla struttura della neoplasia all'interno dell'immagine completa.

3. Scelta dell'Architettura:

   • Il framework utilizza un backbone di Reti Neurali Convoluzionali (CNN) anziché Vision Transformers (ViT). Gli autori sostengono che le CNN forniscano i necessari bias induttivi (invarianza alla traslazione, località) adatti alle dimensioni del dataset e garantiscano la precisa allineamento spaziale richiesto per il calcolo della perdita di attenzione.

4. Funzione di Perdita:
   La perdita totale $L$ combina:

   • Binary Cross-Entropy (BCE) pesata rispetto alle etichette ground truth.
   • Divergenza KL scalata per temperatura tra i logit dell'esperto e del modello principale.
   • Termine di regolarizzazione dell'attenzione (distanza di Frobenius tra le mappe di attenzione).

Contributi Chiave

1. Esperto Ausiliario Solo per l'Addestramento: Introduzione di un esperto acromatico mascherato che promuove l'apprendimento morfologico invariante al colore sopprimendo il contesto di sfondo e le informazioni cromatiche durante l'addestramento.
2. Framework di Distillazione Spaziale: Un meccanismo per trasferire gli indizi strutturali centrati sulla ROI e l'attenzione spaziale dall'esperto a un classificatore di immagini complete implementabile, consentendo l'inferenza in tempo reale senza richiedere maschere di segmentazione.
3. Prestazioni e Robustezza: Dimostrazione di miglioramenti costanti rispetto a solidi baseline basati su rilevamento (es. YOLO) e classificazione (es. Swin-Transformer, GastroNet), insieme a una maggiore spiegabilità e robustezza alle perturbazioni del colore.

Risultati Sperimentali
Gli esperimenti sono stati condotti su un dataset privato di endoscopia gastrica (Seoul St. Mary's Hospital) e su un dataset esterno del colon (PICCOLO).

• Prestazioni di Classificazione: MAGE ha raggiunto un AUC di 0,786 e un F1 di 0,733 sul set di test gastrico, superando i baseline inclusi GastroNet-5M (AUC 0,744), Swin-Tiny (AUC 0,719) e le pipeline basate su YOLO.
• Spiegabilità: MAME ha dimostrato un allineamento superiore con le maschere ground-truth nelle visualizzazioni Grad-CAM.
  • Attention-in-Mask (AiM): MAGE ha ottenuto 0,5303 (rispetto allo 0,3176 del baseline EffNet).
  • Pointing Game Accuracy (PGA): MAGE ha ottenuto 0,7324 (rispetto allo 0,4355 del baseline).
  • Ciò indica che il modello si affida maggiormente alla morfologia della neoplasia piuttosto che a pattern globali spuri.
• Robustezza al Colore: Sotto shift di tonalità sintetici, MAGE ha mostrato una degradazione delle prestazioni significativamente minore rispetto ai baseline EfficientNet e ResNet50 pre-addestrati su GastroNet, confermando l'efficacia dell'addestramento invariante al colore.
• Generalizzazione Cross-Domain: Sul dataset esterno PICCOLO, MAGE ha raggiunto l'AUC più elevato (0,9434), suggerendo che le caratteristiche morfologiche apprese sono trasferibili attraverso diversi domini anatomici.
• Studi di Ablazione: I risultati hanno confermato che il design asimmetrico (Esperto in Scala di Grigi + Principale in RGB) è critico. L'esperto ha performato meglio sugli input in scala di grigi, mentre il modello principale richiedeva l'input RGB completo per un'implementazione ottimale.

Significatività e Rivendicazioni
Il paper sostiene che MAGE affronti il problema dello "shortcut learning" nella classificazione delle neoplasie gastriche, forzando esplicitamente il modello a imparare caratteristiche strutturali invarianti al colore attraverso un segnale di addestramento privilegiato. Decoupling (disaccoppiando) il requisito della maschera per l'addestramento dal requisito della maschera per l'inferenza, MAGE offre una soluzione pratica che migliora sia l'accuratezza diagnostica che l'interpretabilità clinica. Gli autori sottolineano che il loro metodo fornisce non solo metriche di classificazione superiori, ma anche mappe di attenzione interpretabili, essenziali per l'affidabilità clinica.

Limitazioni
Gli autori notano che lo studio si concentra sulla classificazione e sulla robustezza senza addestrare un modulo dedicato di rilevamento/segmentazione, il che potrebbe limitare l'output diretto di localizzazione in certi flussi di lavoro clinici. Il lavoro futuro suggerito consiste nel sfruttare il classificatore appreso per una localizzazione leggera (es. basata su CAM) e nel condurre una validazione multicentrica più ampia.