SO-Mamba: State-Ownership Mamba for Unrolled MRI Reconstruction

Riepilogo Tecnico: SO-Mamba per la Ricostruzione MRI Srotolata

Enunciato del Problema
La ricostruzione accelerata della Risonanza Magnetica (MRI) mira a recuperare immagini di alta qualità da misurazioni dello spazio k sottocampionate, preservando al contempo la coerenza anatomica su grandi regioni spaziali. Sebbene l'apprendimento profondo abbia fatto progredire questo campo attraverso prior basati sui dati, le architetture attuali affrontano limitazioni specifiche nei solutori srotolati accoppiati alla consistenza dei dati (DC). Le Reti Neurali Convoluzionali (CNN) faticano con le dipendenze anatomiche a lungo raggio a causa di campi ricettivi limitati, mentre i Transformer per la visione soffrono di una complessità computazionale quadratica che ostacola la scalabilità verso input clinici ad alta risoluzione.

I Modelli a Spazio di Stato Selettivi (SSM), come Mamba, offrono una modellazione a lungo raggio con complessità lineare, rendendoli regolarizzatori attraenti per la ricostruzione srotolata. Tuttavia, esiste un divario critico di progettazione nel modo in cui l'evidenza di ricostruzione viene assegnata all'interno del calcolo dello spazio di stato di un solutore srotolato accoppiato alla DC. In tali solutori, ogni stadio opera su una diversa iterazione di ricostruzione. Trattare uniformemente tutta l'evidenza – collocando sia il contenuto persistente "residente" che gli aggiornamenti dipendenti dallo stadio "non residenti" nello stesso percorso di contenuto ricorrente – può intrecciare ruoli che richiedono una gestione distinta. Nello specifico, non è chiaro quale evidenza debba diventare il contenuto ricorrente (la traiettoria dello stato) rispetto a quale evidenza debba semplicemente modulare come tale contenuto viene scritto/letto o correggere l'output senza entrare nel ciclo ricorrente.

Metodologia: SO-Mamba
Gli autori propongono SO-Mamba (State-Ownership Mamba), un regolarizzatore che formalizza questo problema di allocazione come proprietà dello stato. Il principio fondamentale separa i ruoli funzionali dell'evidenza di ricostruzione all'interno di ogni stadio srotolato in tre percorsi distinti:

1. Residenza Ricorrente: Evidenza che diventa il contenuto diretto del calcolo dello spazio di stato con scansione selettiva.
2. Accesso all'Interfaccia di Stato: Evidenza che controlla come lo stato viene scritto e letto (modulando le interfacce B e C).
3. Correzione dell'Output Non-Stato: Evidenza che contribuisce all'output finale senza entrare nel percorso del contenuto ricorrente.

Componenti Chiave:

• Router di Proprietà dello Stato (SOR): Questo modulo prende la rappresentazione in feature dell'iterazione di ricostruzione corrente e la divide in due flussi:
  • Vettore Residente ($L_k$): Generato tramite una proiezione binomiale fissa e non adattiva del vettore. Questo flusso è assegnato alla residenza ricorrente, fornendo i token per il percorso dei contenuti di Mamba. Rappresenta layout anatomici stabili e coerenti spazialmente.
  • Evidenza Non Residente ($G_k$): Composta dall'evidenza rimanente (deviazioni locali dal vettore) e da specifiche caratteristiche non residenti. Questo flusso è escluso dalla fonte di contenuto ricorrente ($G_k \nrightarrow u_k$). Invece, viene instradato a:
    • Modulazione dell'Interfaccia di Stato: $G_k$ prevede parametri affini per modulare le interfacce di stato B e C, adattando come viene elaborato il contenuto residente.
    • Uscita di Output Non-Stato: $G_k$ passa attraverso un blocco di Rifinitura Non-Stato (NSR) e viene concatenato con la lettura di Mamba per correggere l'output finale.
• Aggiornamento Mamba-3 Consapevole della Proprietà: Il blocco Mamba-3 elabora il vettore residente come token di input, utilizzando contemporaneamente il flusso non residente per condizionare le transizioni dello spazio di stato. Ciò garantisce che la traiettoria dello stato ricorrente rimanga dominata da contenuti anatomici stabili, mentre gli aggiornamenti dinamici influenzano il sistema attraverso interfacce controllate e correzioni dell'output.
• Diagnostica di Perdita a Banda Esterna a Due Livelli: Per validare la separazione dei ruoli, gli autori introducono una metrica diagnostica che misura l'energia "a banda esterna" (alta frequenza) in due punti: la traiettoria grezza dello stato nascosto ($H_{Leak}$) e la lettura post-scansione ($R_{Leak}$). Ciò separa l'archiviazione dello stato nascosto dall'espressione della lettura, verificando che l'evidenza non residente non corrompa la traiettoria dello stato ricorrente ma sia efficacemente espressa a livello di output.

Contributi Chiave

1. Formulazione del Problema: Il documento identifica l'"allocazione dello stato locale allo stadio" come un problema di progettazione distinto nella ricostruzione MRI srotolata basata su Mamba e lo formula come proprietà dello stato.
2. Architettura SO-Mamba: Un regolarizzatore innovativo basato su Mamba-3 che assegna esplicitamente l'evidenza alla residenza ricorrente, all'accesso all'interfaccia di stato e alla correzione dell'output non-stato all'interno di ogni stadio.
3. Meccanismo SOR: L'introduzione del Router di Proprietà dello Stato, che costruisce un vettore residente e instrada l'evidenza non residente alla modulazione dell'interfaccia e alla correzione dell'output, imponendo il contratto di proprietà.
4. Valutazione Completa: Valutazione su cinque benchmark pubblici (fastMRI, CC359, ACDC, M4Raw, Prostate158) che coprono diverse anatomie e pattern di campionamento, affiancata da una nuova analisi diagnostica del comportamento dello stato.

Risultati Sperimentali

• Prestazioni: SO-Mamba supera costantemente le baseline basate su CNN, Transformer e Mamba (inclusi UNet, TransUNet, ReconFormer, LMO e HiFi-Mamba) su cinque dataset.
  • Su fastMRI (AF=4), ottiene 34.65 dB PSNR e 0.855 SSIM.
  • Su CC359 (AF=4), ottiene 36.69 dB PSNR e 0.954 SSIM.
  • Si posiziona al primo posto anche nelle configurazioni radiali (ACDC) e multi-coil (M4Raw).
• Efficienza: SO-Mamba raggiunge la migliore accuratezza tra i modelli efficienti utilizzando i FLOPs più bassi (45.59G) rispetto a HiFi-Mamba (67.87G) e LMO (484.98G) su fastMRI con accelerazione 8×.
• Studi di Ablazione:
  • La rimozione dell'accesso all'interfaccia di stato o dell'uscita di output degrada le prestazioni.
  • Consentire all'evidenza non residente di entrare nel percorso del contenuto ricorrente (violando il contratto di proprietà) causa un calo netto delle prestazioni (da 34.65 dB a 31.71 dB), confermando la necessità della separazione.
  • La diagnostica di perdita mostra che la rimozione dell'uscita di output forza il contenuto a banda esterna nella lettura ricorrente (aumentando $R_{Leak}$ di circa 5 volte), mentre il design completo di SO-Mamba mantiene lo stato ricorrente pulito esprimendo le correzioni necessarie a livello di output.

Significato e Affermazioni
Il documento afferma che SO-Mamba affronta un vincolo architettonico fondamentale nella ricostruzione MRI srotolata: la necessità di distinguere tra l'evidenza che dovrebbe formare la traiettoria dello stato persistente e l'evidenza che dovrebbe influenzare solo la dinamica del sistema o l'output. Imponendo la proprietà dello stato, il metodo previene l'intreccio di strutture anatomiche stabili con aggiornamenti dipendenti dallo stadio, portando a confini più netti e mappe di errore più uniformi.

Gli autori sottolineano che il Router di Proprietà dello Stato (SOR) è una regola di assegnazione architettonica piuttosto che una decomposizione semantica di anatomia versus artefatti. Il significato risiede nel dimostrare che gli SSM a complessità lineare possono essere efficacemente adattati per la ricostruzione srotolata gestendo attentamente il flusso delle informazioni, raggiungendo prestazioni all'avanguardia con elevata efficienza computazionale. Il lavoro suggerisce che i futuri miglioramenti potrebbero risiedere nell'instradamento adattivo della proprietà e in contesti più ampi di acquisizione clinica, ma il contributo attuale è l'istituzione del principio di proprietà dello stato stesso.