BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability

Riepilogo Tecnico di "BrainAnytime: Preaddestramento Cross-Modale Consapevole dell'Anatomia per l'Analisi di Immagini Cerebrali con Disponibilità Arbitraria di Modalità"

Enunciato del Problema

I flussi di lavoro diagnostici clinici per le malattie neurodegenerative, in particolare il Morbo di Alzheimer (AD), seguono un "percorso di escalation delle modalità". I clinici iniziano tipicamente con imaging strutturale di routine (ad esempio, risonanza magnetica T1-pesata), aggiungono selettivamente sequenze (ad esempio, FLAIR, T2) per affinare le diagnosi differenziali e riservano l'imaging molecolare (ad esempio, PET dell'amiloide) per i casi incerti. Di conseguenza, i dati dei pazienti nelle coorti reali sono caratterizzati da sottoinsiemi di modalità eterogenei e spesso incompleti.

I modelli di intelligenza artificiale attuali assumono generalmente modalità di input fisse, richiedendo un cambio architetturale o un ripreaddestramento quando la disponibilità dei dati cambia. Inoltre, i modelli esistenti di preaddestramento su larga scala per immagini cerebrali soffrono spesso di tre limitazioni in questo contesto:

1. Sono tipicamente limitati alla sola risonanza magnetica (MRI), fallendo nell'apprendere ponti espliciti verso la PET dell'amiloide, un biomarcatore critico.
2. Assumono combinazioni di modalità fisse, non corrispondendo alla mancanza dipendente dallo stadio dei dati reali.
3. Il preaddestramento standard con Masked Autoencoder (MAE) si basa su mascheramento spazialmente uniforme, ignorando il fatto che l'AD presenta una vulnerabilità regionale selettiva (ad esempio, l'ippocampo) piuttosto che una patologia uniforme.

Metodologia

Gli autori propongono BrainAnytime, un framework di preaddestramento unificato progettato per supportare l'analisi di immagini cerebrali in condizioni di disponibilità arbitraria di modalità. Il framework è preaddestrato su 34.899 scansioni cerebrali 3D provenienti da cinque dataset (A4, DOD-ADNI, AIBL, BraTS23, NACC) che coprono MRI multi-sequenza e PET dell'amiloide. Integra tre componenti principali all'interno di un'architettura condivisa di autoencoder mascherato 3D:

1. Multi-MAE3D (Autoencoder Mascherato 3D Multi-Modale)

• Architettura: Un encoder Transformer condiviso elabora qualsiasi sottoinsieme di modalità disponibili (T1, T2, FLAIR, PET).
• Tokenizzazione: Ogni modalità utilizza un embedding di patch convoluzionale 3D specifico per la modalità.
• Gestione dei Dati Mancanti: Le modalità mancanti sono completamente mascherate. Per mantenere un conteggio di token costante indipendentemente dalla disponibilità di input, il budget di token visibili viene ridistribuito tra le modalità osservate utilizzando una distribuzione di Dirichlet (Dir($\alpha$)).
• Decodifica: Ogni modalità ha un decoder leggero che ricostruisce tutte le patch. La funzione di perdita di addestramento ($L_{MAE}$) è calcolata solo sulle patch mascherate utilizzando l'errore quadratico medio con normalizzazione per patch.

2. Distillazione Cross-Modale Reciproca (RCMD)

• Obiettivo: Apprendere le corrispondenze tra MRI strutturale e PET molecolare.
• Meccanismo: Una rete insegnante con media mobile esponenziale (EMA) condivide l'architettura dell'encoder. Per i campioni accoppiati MRI-PET, le rappresentazioni a livello di gruppo ($z_{MRI}$ e $z_{PET}$) sono calcolate tramite pooling medio.
• Perdita: Due predittori MLP simmetrici ($g_{M\to P}$ e $g_{P\to M}$) sono addestrati per prevedere le caratteristiche dell'insegnante cross-modale. La perdita ($L_{RCMD}$) massimizza la similarità del coseno tra le caratteristiche previste dallo studente e le caratteristiche dell'insegnante nello spazio della modalità opposta.

3. Mascheramento Curricolare Consapevole della Patologia (PACM)

• Obiettivo: Prioritizzare l'anatomia vulnerabile alla malattia durante la ricostruzione piuttosto che mascherare in modo uniforme.
• Punteggio di Importanza: Un punteggio statico è derivato da prior neuropatologici (stadi di Braak) mappati tramite l'atlante AAL116. Un punteggio dinamico è ottenuto aggregando le mappe di attenzione da un insegnante di anatomia EMA. Questi sono combinati per generare punteggi di importanza delle patch.
• Piano Curricolare: Viene impiegata una strategia in tre fasi:
  1. Fase iniziale: Mascheramento casuale uniforme.
  2. Fase centrale: Il parametro di temperatura ($\tau$) viene ricotto tramite decadimento coseno, affilando progressivamente la distribuzione della maschera per favorire le regioni ad alta importanza per il mascheramento (cioè, costringendo il modello a ricostruire regioni critiche).
  3. Fase finale: Temperatura fissa.
• Selezione: Le patch sono selezionate utilizzando il trucco Gumbel-top-k, assicurando che le patch ad alta importanza abbiano maggiori probabilità di essere mascherate e quindi ricostruite.

Addestramento e Fine-tuning:
La perdita totale è $L = L_{MAE} + \lambda \cdot \mathbb{1}_{[paired]} \cdot L_{RCMD}$. Per i compiti a valle, viene utilizzato solo l'encoder studente con un task head leggero. Le modalità mancanti rimangono azzerate e bloccate dall'attenzione durante l'inferenza, assicurando che il modello gestisca sottoinsiemi arbitrari senza cambiamenti architetturali.

Contributi Chiave

1. Framework Unificato: Un singolo modello capace di accettare qualsiasi sottoinsieme di modalità (da una singola scansione T1 a un completo lavoro multimodale) senza ripreaddestramento o cambio architetturale.
2. Apprendimento Cross-Modale: Apprende esplicitamente le corrispondenze struttura-molecolare tra MRI e PET tramite distillazione reciproca, colmando il divario tra evidenze strutturali e patologiche.
3. Preaddestramento Consapevole dell'Anatomia: Introduce una strategia di mascheramento curricolare che indirizza lo sforzo ricostruttivo del modello verso la neuroanatomia saliente per l'AD, affrontando le limitazioni del mascheramento uniforme nei MAE standard.
4. Preaddestramento Scalabile: Preaddestrato con successo su un dataset eterogeneo su larga scala (34.899 scansioni) che imita i pattern di mancanza clinica.

Risultati Sperimentali

Il modello è stato valutato su quattro compiti a valle (CN vs AD, CN vs MCI, regressione MMSE, previsione dell'età) attraverso cinque impostazioni di modalità motivate clinicamente (dalla sola T1 a T1+T2+FLAIR+PET).

• Prestazioni: BrainAnytime ha superato i modelli specifici per modalità (ad esempio, 3D ResNet, M3T), i metodi di fusione a doppia modalità, i modelli fondazionali cerebrali su larga scala (BrainIAC, SAM-Brain3D) e gli approcci di apprendimento multimodale incompleti (Flex-MoE, FuseMoE).
• Metriche Chiave:
  • CN vs AD: Ha raggiunto una precisione media del 89,4% (contro l'84,2% del prossimo miglior baseline, Flex-MoE), rappresentando un miglioramento relativo del 6,2% rispetto al baseline più forte con modalità mancanti.
  • CN vs MCI: Ha raggiunto una precisione media del 67,2% (contro il 60,6% di MoE-retriever), un miglioramento relativo del 7,0%.
  • Regressione: Ha dimostrato prestazioni superiori nella previsione di MMSE e età attraverso tutte le impostazioni di modalità.
• Robustezza: Uno studio di robustezza intra-soggetto ha confermato che le prestazioni sono migliorate costantemente all'aggiunta di più modalità (ad esempio, l'F1 CN vs AD è aumentato dal 65,0% con sola T1 all'80,4% con input completo), dimostrando la capacità del modello di sfruttare efficacemente i dati disponibili.
• Ablazione: La rimozione individuale di RCMD o PACM ha portato a cali di prestazioni, confermando la necessità di entrambe la distillazione cross-modale e il mascheramento consapevole dell'anatomia.

Significato

Il documento afferma che BrainAnytime affronta un divario critico tra la realtà clinica (acquisizione di dati a stadi, incompleta) e i paradigmi attuali di intelligenza artificiale (modelli a input fisso). Consentendo a un singolo modello di accettare flessibilmente qualsiasi imaging disponibile, si allinea al percorso di escalation delle modalità cliniche. Il framework dimostra che il preaddestramento con distillazione cross-modale e mascheramento consapevole della patologia produce un modello fondazionale robusto che supera significativamente i baseline esistenti, in particolare in scenari in cui i dati sono incompleti o le modalità mancano. Questo approccio offre una soluzione pratica per il dispiegamento dell'intelligenza artificiale in ambienti clinici reali, dove i dati multimodali completi sono raramente disponibili per ogni paziente.