Diffusion MRI Transformer with a Diffusion Space Rotary Positional Embedding (D-RoPE)

Titolo

Diffusion MRI Transformer con un Embedding Posizionale Rotatorio nello Spazio di Diffusione (D-RoPE)

1. Il Problema

La Risonanza Magnetica con Diffusione (dMRI) è fondamentale per studiare i cambiamenti microstrutturali del cervello, ma l'addestramento di rappresentazioni generali (foundation models) da dati dMRI è stato limitato rispetto ad altre modalità di imaging medico. Le sfide principali identificate sono:

• Natura complessa del segnale: I dati dMRI non sono semplici volumi 3D statici. Sono composti da più volumi cerebrali, ciascuno con caratteristiche di attenuazione dipendenti dalla direzione e dall'intensità dei gradienti di diffusione (vettori b e valori b).
• Dipendenze congiunte: È necessario modellare simultaneamente le dipendenze spaziali (x, y, z) e quelle nello spazio di diffusione (direzione e intensità).
• Variabilità dei protocolli: I protocolli di acquisizione variano notevolmente tra studi (diverso numero di direzioni, diversi valori b), rendendo difficile l'uso di modelli tradizionali che richiedono input fissi o che non gestiscono bene la geometria sferica dello spazio di diffusione.
• Limiti dei modelli esistenti: Gli approcci di deep learning attuali non catturano adeguatamente le proprietà uniche dei segnali di diffusione e spesso falliscono nel generalizzare tra diversi protocolli di acquisizione.

2. Metodologia

Gli autori propongono un'architettura basata su Transformer addestrata con un obiettivo di Masked Autoencoding (MAE).

Architettura del Modello

• Input: I volumi dMRI vengono divisi in patch 3D.
• Embedding Posizionale Assoluto Modificato: A differenza dei ViT standard che usano solo coordinate spaziali, il modello combina:
  • Coordinate spaziali (x, y, z) codificate con un encoding sinusoidale 3D.
  • Coordinate della direzione di diffusione (sfera: ρ, θ, φ) codificate tramite un layer lineare apprendibile.
  • Questi due vettori vengono concatenati e sommati agli embedding delle patch.
• Meccanismo di Attenzione: I blocchi di attenzione alternano l'elaborazione tra le dimensioni spaziali e quelle dello spazio di diffusione, ispirandosi ai transformer per video.

Il Contributo Chiave: D-RoPE

Il cuore dell'innovazione è il Diffusion Space Rotary Positional Embedding (D-RoPE), una generalizzazione del Rotary Positional Embedding (RoPE) standard.

• Motivazione: Il RoPE standard calcola la distanza relativa come una differenza lineare ($m-n$) in una sequenza 1D. Tuttavia, nello spazio di diffusione, due acquisizioni sono definite da un valore $b$ (lineare) e un vettore $b$ (sferico). Una distanza lineare non è appropriata.
• Definizione della Distanza: D-RoPE sostituisce la distanza lineare con una distanza geometrica $D(m, n)$ calcolata tra due punti nello spazio di diffusione $(b, v)$:
  $$D(m, n) = \sqrt{\gamma(b_m - b_n)^2 + \arccos^2(|v_m \cdot v_n|)}$$
  Dove $\gamma$ è un parametro di pesatura e il termine $\arccos$ gestisce la simmetria di 180° delle direzioni opposte (poiché $v$ e $-v$ producono codifiche equivalenti).
• Funzionamento: Questa distanza geometrica viene utilizzata per generare le matrici di rotazione nell'attenzione, permettendo al modello di comprendere le relazioni relative tra diverse direzioni di diffusione e valori b, indipendentemente dall'ordine o dal numero totale di direzioni.

Addestramento (Pretraining)

• Obiettivo: Masked Autoencoding (MAE). Una parte delle patch (sia spaziali che di direzione) viene mascherata e il modello deve ricostruire i volumi originali.
• Strategie di Mascheramento: Sono state testate tre strategie: mascheramento puramente spaziale (75%), mascheramento puramente direzionale (50% delle direzioni intere) e una strategia alternata (epoch-wise).
• Dataset: Pretraining su HCP Young Adult (HCP-YA), valutazione su HCP Development, HCP Aging e ADNI.

3. Contributi Chiave

1. Nuova Architettura Transformer: Un modello capace di gestire acquisizioni dMRI arbitrarie (qualsiasi numero di direzioni e valori b) senza necessità di risampling o upsampling a un protocollo comune.
2. D-RoPE: Un nuovo meccanismo di embedding posizionale che generalizza il RoPE allo spazio sferico della diffusione, catturando le interdipendenze geometriche tra direzioni e valori b.
3. Rappresentazioni Trasferibili: Dimostrazione che le rappresentazioni apprese tramite self-supervised learning sono robuste e trasferibili a diversi compiti downstream, superando i limiti dei modelli specifici per task.

4. Risultati

I risultati sono stati valutati sia sulla ricostruzione che su compiti downstream (regressione dell'età, classificazione del sesso, classificazione MCI, previsione ADAS).

• Ricostruzione (MAE):

  • La combinazione di D-RoPE e mascheramento alternato ha ottenuto le metriche migliori (PSNR, SSIM, FID).
  • L'uso di D-RoPE ha permesso di ricostruire non solo la struttura anatomica ma anche le intensità generali, catturando meglio le relazioni tra direzioni e valori b rispetto ai modelli senza D-RoPE.
  • Le mappe di Anisotropia Frazionata (FA) e Diffusività Media (MD) derivate dalle ricostruzioni mostrano errori bassi, confermando che le informazioni microstrutturali sono preservate nello spazio latente.

• Compiti Downstream:

  • Classificazione MCI (Deterioramento Cognitivo Lieve): Il modello con rappresentazioni "congelate" (frozen latents) e un head MLP ha raggiunto una accuratezza bilanciata (BA) del 64.7%, superando i metodi basati su feature manuali (61.9%) e i ViT standard addestrati da zero.
  • Predizione ADAS (Scala di Valutazione Cognitiva): Il modello ha ottenuto un coefficiente di correlazione ($\rho$) di 0.38, superiore a tutti i baselines, inclusi i modelli completamente supervisionati.
  • Robustezza: Il modello ha mantenuto alte prestazioni anche quando il numero di direzioni di diffusione variava tra pretraining e task downstream, dimostrando la capacità di generalizzare a protocolli diversi senza ri-addestramento specifico.

5. Significato e Implicazioni

• Foundation Model per dMRI: Questo lavoro rappresenta un passo significativo verso la creazione di un "foundation model" universale per la dMRI, capace di armonizzare analisi tra coorti con protocolli di acquisizione diversi.
• Efficienza dei Dati: Le rappresentazioni pre-addestrate permettono di ottenere prestazioni competitive o superiori con dati di addestramento limitati per i task clinici, un vantaggio cruciale in ambito medico dove i dati etichettati sono scarsi.
• Nuovi Biomarcatori: La capacità del modello di estrarre pattern microstrutturali predittivi senza dover adattare modelli biofisici espliciti suggerisce che l'approccio potrebbe rivelare nuovi biomarcatori per il declino cognitivo e le malattie neurodegenerative.
• Limitazioni: Il calcolo di D-RoPE ha un costo computazionale leggermente superiore ($O(lh^2)$ contro $O(lh)$ del RoPE standard) perché la distanza geometrica non è facilmente fattorizzabile in matrici di chiavi e query come nel caso lineare.

In sintesi, il paper introduce un approccio innovativo che risolve il problema della modellazione geometrica complessa dei dati dMRI, permettendo l'apprendimento di rappresentazioni generali robuste che superano i metodi tradizionali sia nella ricostruzione che nella diagnosi clinica.