Reconstruct Anything Model: a lightweight general model for computational imaging

1. Il Problema

Il campo dell'imaging computazionale affronta problemi inversi, ovvero il recupero di un'immagine latente $x$ a partire da misurazioni degradate $y$, modellate come $y \sim p(y|Ax)$, dove $A$ è un operatore lineare che descrive la fisica dell'acquisizione (es. trasformata di Radon per la TAC, trasformata di Fourier per la risonanza magnetica) e $p$ è la distribuzione del rumore (Gaussiano, Poisson, o misto).

Le soluzioni esistenti si dividono in due categorie principali, entrambe con limiti significativi:

• Algoritmi iterativi (es. Plug-and-Play, Diffusione): Utilizzano denoiser pre-addestrati all'interno di schemi iterativi. Sono computazionalmente costosi, lenti e spesso introducono sfocature o artefatti.
• Architetture "Unrolled" (srotolate): Reti neurali progettate per mimare i passi di un algoritmo di ottimizzazione. Sebbene efficaci, sono specifiche per un singolo problema, richiedono un addestramento costoso per ogni nuova configurazione e non scalano bene tra diversi domini (es. da scale di grigio a colori, o da MRI a TAC).

L'obiettivo è creare un modello leggero, non iterativo e generale che possa risolvere una vasta gamma di problemi inversi con prestazioni state-of-the-art, adattandosi a nuove distribuzioni senza ground-truth.

2. Metodologia: Reconstruct Anything Model (RAM)

RAM è un'architettura neurale non iterativa basata su DRUNet (una rete convoluzionale residua), potenziata da meccanismi specifici per integrare la fisica dell'acquisizione direttamente nella rete.

Architettura Chiave

L'architettura introduce tre modifiche fondamentali rispetto a un backbone standard come DRUNet:

1. Modulo di Stima Prossimale (Proximal Estimation Module):
   Invece di usare direttamente $A^\top y$ (che può essere sfocato) o la pseudo-inversa $A^\dagger y$ (sensibile al rumore), RAM calcola un passo prossimale:
   $$\text{prox}_{\lambda f}(A^\top y) = \arg \min_u \lambda \|Au - y\|^2 + \|u - A^\top y\|^2$$
   Questo passo bilancia l'inversione dell'operatore e la regolarizzazione, trasformando il problema in una forma più adatta al denoising iniziale.

2. Modulo Krylov Subspace (KSM):
   Questo è il cuore dell'adattamento. Invece di "srotolare" iterativamente un algoritmo, RAM utilizza un modulo che apprende le combinazioni lineari dei vettori dello spazio di Krylov generati dall'operatore $A^\top A$.
   A ogni scala della rete, le mappe di caratteristiche vengono decodificate nello spazio delle immagini, elaborate attraverso iterazioni di Krylov (concatenando potenze di $A^\top A$ applicate alle caratteristiche e a $A^\top y$) e poi ricalcolate. Questo permette alla rete di "imparare" le operazioni di ottimizzazione necessarie senza doverle eseguire esplicitamente passo dopo passo.

3. Condizionamento Multiscala e sul Rumore:

   • Multiscala: Per problemi mal posti su griglie fini, RAM condiziona l'architettura su operatori definiti su griglie più grezze ($A_s$), ispirandosi ai metodi multigriglia, per stabilizzare l'inversione.
   • Rumore: La rete riceve in input i parametri del rumore (livello Gaussiano $\sigma$ e guadagno Poisson $\gamma$) concatenati come mappe di caratteristiche costanti, permettendo di gestire modelli di rumore misti (Poisson-Gaussian).

Adattabilità ai Canali

L'architettura supporta nativamente immagini in scala di grigi (1 canale), complesse (2 canali, per MRI) e a colori (3 canali). Solo i livelli di input e output e i blocchi di codifica/decodifica del KSM vengono adattati per il numero di canali; il resto dei pesi è condiviso.

3. Addestramento e Fine-tuning

• Addestramento Supervisionato: RAM viene addestrato simultaneamente su un vasto set di compiti (deblurring, MRI, TAC, inpainting, super-risoluzione) utilizzando dataset come LSDIR, fastMRI e LIDC-IDRI. La funzione di perdita è una somma pesata delle perdite per ogni compito, normalizzata in base al rapporto segnale-rumore (SNR) per evitare bias verso compiti più facili.
• Fine-tuning Self-Supervised (Senza Ground-Truth): Un contributo cruciale è la capacità di adattare il modello pre-addestrato a nuovi problemi o dataset reali senza dati di riferimento (ground-truth). Utilizzando solo le misurazioni $y$, il modello viene fine-tunato minimizzando una perdita di consistenza delle misurazioni ($L_{MC}$, es. SURE o splitting loss) e una perdita di nullspace ($L_{NULL}$, es. Equivariant Imaging) per recuperare le informazioni mancanti. Questo permette di ottenere risultati eccellenti con poche immagini (anche una sola).

4. Risultati Sperimentali

Il paper presenta risultati su una vasta gamma di compiti, dimostrando prestazioni superiori o competitive rispetto a metodi iterativi e architetture unrolled molto più pesanti.

• Prestazioni In-Distribution: Su compiti di deblurring (Gaussiano e Motion), MRI e TAC, RAM ottiene risultati SOTA (State-of-the-Art) in termini di PSNR e SSIM.
  • Supera i metodi iterativi come DPIR e DDRM (basati su diffusione) in termini di velocità e qualità.
  • Confrontato con uDPIR (una versione unrolled di DRUNet), RAM ottiene prestazioni simili o migliori con un costo computazionale drasticamente inferiore.
• Efficienza Computazionale:
  • RAM è 3.7 volte più veloce dell'architettura unrolled uDPIR.
  • È 120 volte più veloce dei metodi basati su diffusione (DDRM).
  • Richiede meno memoria e operazioni FLOP (circa 360 GFLOPS vs 2234 GFLOPS per uDPIR).
• Prestazioni Zero-Shot e Out-of-Distribution:
  • Il modello funziona bene su compiti non visti durante l'addestramento (es. MRI multi-coil, TAC con rumore Poisson puro).
  • Viene esteso con successo a problemi non lineari (Phase Retrieval) e ciechi (Blind Deblurring), dimostrando una robustezza sorprendente.
• Fine-tuning Self-Supervised:
  • In scenari reali (es. microscopia Cryo-EM, imaging a basso fotone, dati satellitari Sentinel-2), il fine-tuning su una singola immagine o poche decine di misurazioni senza ground-truth porta a miglioramenti significativi rispetto ai metodi basati su denoiser generici, superando le basi PnP e Diffusione.

5. Contributi Chiave e Significato

1. Unificazione di Due Mondi: RAM riesce a fondere l'interpretabilità e la robustezza delle architetture "unrolled" (che conoscono la fisica $A$) con l'efficienza e la potenza delle moderne architetture di visione (come DRUNet/UNet), eliminando la necessità di iterazioni esplicite.
2. Modello Fondamentale per l'Imaging: Dimostra che un singolo modello leggero (36M parametri) può gestire compiti eterogenei (da MRI a microscopia) con prestazioni SOTA, sfidando l'idea che ogni problema inverso richieda un modello specializzato.
3. Adattabilità Senza Ground-Truth: La capacità di fine-tunare il modello in modo self-supervised su dati reali, anche con un solo campione, risolve uno dei maggiori colli di bottiglia nell'imaging scientifico e medico, dove i dati etichettati sono rari o inesistenti.
4. Efficienza: Offre prestazioni competitive con un costo computazionale e di memoria ridottissimi rispetto alle alternative iterative o basate su diffusione, rendendo l'elaborazione in tempo reale più fattibile.

In conclusione, RAM rappresenta un passo avanti verso un approccio "foundation model" per l'imaging computazionale, dove lo sforzo principale è posto sullo sviluppo di modelli base robusti e tecniche di adattamento self-supervised, piuttosto che sulla creazione di soluzioni specifiche per ogni singolo problema inverso.