Multi-Dataset Cross-Domain Knowledge Distillation for Unified Medical Image Segmentation, Classification, and Detection

Sintesi Tecnica: Distillazione della Conoscenza Cross-Dominio su Multi-Dataset per Segmentazione, Classificazione e Rilevamento Unificati di Immagini Mediche

Enunciato del Problema
L'apprendimento profondo ha rivoluzionato l'analisi delle immagini mediche, tuttavia i modelli addestrati su singoli dataset spesso falliscono nel generalizzare a nuovi ambienti clinici a causa di spostamenti di dominio causati da variazioni nei protocolli di imaging, nell'hardware degli scanner e nelle popolazioni di pazienti. Sebbene l'apprendimento per trasferimento e l'aumento dei dati offrano soluzioni parziali, essi fanno spesso affidamento su pre-addestramento da singola fonte o su dati sintetici che non riescono a catturare dettagli anatomici fini e la variabilità inter-dataset. Inoltre, gli approcci esistenti di distillazione della conoscenza (KD) utilizzano tipicamente un singolo insegnante su dati omogenei per un'unica attività (ad esempio, segmentazione), trascurando la natura diversificata e multi-task dei flussi di lavoro clinici, che comprendono segmentazione, classificazione e rilevamento di oggetti attraverso modalità eterogenee (MRI, TC, RX).

Metodologia
Gli autori propongono un framework unificato, agnostico rispetto al task, di tipo insegnante-studente progettato per aggregare conoscenze da molteplici dataset di imaging medico eterogenei. Il framework opera in tre fasi sequenziali:

1. Addestramento dell'Insegnante Adversariale di Dominio:

   • Un insegnante target viene addestrato esclusivamente sul dataset target utilizzando una loss specifica per il task (ad esempio, loss Dice per la segmentazione).
   • Gli insegnanti sorgente vengono addestrati sui rispettivi dataset sorgente ma arricchiti con un discriminatore di dominio. Utilizzando un obiettivo adversariale di dominio, questi encoder sono costretti a produrre caratteristiche invarianti al dominio che confondono il discriminatore, allineando le distribuzioni delle caratteristiche sorgente e target. Questa fase supporta sia l'adattamento supervisionato (utilizzando dati target etichettati) che quello non supervisionato.

2. Costruzione dell'Insegnante Congiunto tramite Fusione a Cross-Attention:

   • Un insegnante congiunto ($\mathcal{T}^*$) viene costruito fondendo le caratteristiche dell'encoder e del collo di bottiglia a più livelli provenienti dall'insegnante target e da tutti gli insegnanti sorgente.
   • Questa fusione utilizza meccanismi di cross-attention per aggregare rappresentazioni complementari da diversi domini. Gli strati dell'encoder e del collo di bottiglia di tutti gli insegnanti costituenti vengono congelati durante questo processo per preservare la conoscenza appresa.
   • Una nuova testata specifica per il task (decoder, classificatore o testata di rilevamento) viene attaccata allo scheletro fuso e addestrata da zero sul dataset target. Ciò consente all'insegnante congiunto di specializzarsi nel task target mantenendo al contempo caratteristiche invarianti al dominio.

3. Distillazione Multi-Livello Guidata da Curriculum:

   • Un modello studente compatto viene addestrato sul dataset target per ereditare la conoscenza dall'insegnante congiunto.
   • L'obiettivo di addestramento combina una loss primaria per il task con termini multipli di distillazione: allineamento delle caratteristiche (MSE), similarità coseno, loss contrastiva (divergenza KL sulle distribuzioni di similarità) ed estensioni specifiche per il task (ad esempio, trasferimento di attenzione, distillazione dei logit, allineamento delle caratteristiche ROI per il rilevamento).
   • Un fattore curriculum ($c_f$) pianifica dinamicamente il peso delle loss di distillazione. L'addestramento inizia con un focus sulla loss del task supervisionato, aumentando gradualmente l'influenza dei termini di distillazione per affinare le rappresentazioni interne dello studente senza destabilizzare l'apprendimento iniziale.

Il framework è progettato per essere flessibile, supportando architetture UNet, TResUNet, MedViT, EfficientNet, Faster R-CNN e RF-DETR attraverso task di segmentazione, classificazione e rilevamento.

Contributi Chiave

1. Framework Unificato: Un'unica architettura scalabile in grado di gestire segmentazione, classificazione a livello di immagine e rilevamento di oggetti attraverso dataset medici eterogenei.
2. Architettura dell'Insegnante Congiunto: Un metodo innovativo per fondere caratteristiche a più livelli da molteplici insegnanti sorgente tramite cross-attention, abilitando un trasferimento robusto di rappresentazioni cross-dominio senza richiedere un singolo modello pre-addestrato monolitico.
3. Distillazione Basata su Curriculum: Una strategia di distillazione multi-livello che combina loss di segmentazione, contrastive, di allineamento e di similarità, governata da una pianificazione curriculum per garantire un'ottimizzazione stabile.
4. Valutazione Completa: Esperimenti estesi su tre task, tre modalità (MRI, TC, RX) e numerosi dataset (inclusi BrainMetShare, ISLES, BraTS, LiTS, KiTS, COVIDx-CXR, OASIS, ADNI e DeepLesion).

Risultati Sperimentali
Il metodo proposto è stato valutato rispetto a due baseline: modelli specifici per dataset (addestrati indipendentemente su ciascun dataset) e modelli multi-testata multi-dataset (scheletro condiviso con testate specifiche per dataset).

• Segmentazione: Su dataset MRI e TC, i modelli studente hanno costantemente superato entrambe le baseline. Per i task basati su TC, lo studente ha raggiunto miglioramenti medi dell'IoU di circa il 2,5% rispetto alle baseline specifiche per dataset e di circa il 4,8% rispetto ai modelli multi-testata, con riduzioni significative della Distanza di Hausdorff (HD95), indicando previsioni di confine più nitide.
• Classificazione: Nei task di classificazione delle malattie polmonari e della demenza, i modelli studenti distillati hanno raggiunto una maggiore accuratezza rispetto alle baseline specifiche per dataset e multi-testata. La distillazione late-$k$ (trasferimento di conoscenza dagli strati più profondi) ha generalmente prodotto risultati superiori.
• Rilevamento di Oggetti: Su dataset TC polmonari, il framework ha migliorato la Precisione Media (mAP). L'insegnante congiunto ha spesso raggiunto le prestazioni più elevate, mentre lo studente distillato ha mantenuto la maggior parte di questi guadagni, offrendo un modello più compatto adatto al deployment.
• Studi di Ablazione: I risultati hanno confermato che (i) l'adattamento target supervisionato migliora l'accuratezza dei confini; (ii) le prestazioni scalano con il numero di dataset sorgente; e (iii) l'estrazione di caratteristiche multi-livello (encoder + collo di bottiglia) produce risultati migliori rispetto all'estrazione a singolo livello.
• Visualizzazione: Le visualizzazioni t-SNE hanno dimostrato che i modelli studenti distillati apprendono cluster di caratteristiche più compatti e meglio separati rispetto alle baseline specifiche per dataset in tutti e tre i task.

Significato e Affermazioni
Il documento afferma che la distillazione della conoscenza su multi-dataset funge da approccio scalabile e agnostico rispetto al task per migliorare le prestazioni attraverso domini eterogenei di imaging medico. Aggregando rappresentazioni invarianti al dominio da fonti diverse e distillandole in uno studente compatto, il framework affronta le sfide della scarsità di dati, dei costi di annotazione e degli spostamenti di dominio. Gli autori sottolineano che il loro approccio produce miglioramenti coerenti rispetto a baseline solide, dimostrando una maggiore robustezza agli spostamenti distribuzionali e una generalizzazione superiore. L'opera si posiziona come un'estensione sostanziale della ricerca precedente focalizzata sulla segmentazione, ora validata sull'intero spettro dei task di analisi delle immagini mediche.