Transformer Causality Regularization for Dynamic Inverse Problems

Titolo: Regularizzazione della Causalità con Transformer per Problemi Inversi Dinamici

1. Il Problema

Il lavoro si concentra sui problemi inversi dinamici (DIP), ovvero problemi in cui la quantità di interesse evolve nel tempo e deve essere ricostruita da osservazioni indirette (es. tomografia computerizzata dinamica, risonanza magnetica).
Le sfide principali identificate sono:

• Sottocampionamento: In genere, sono disponibili poche misurazioni per ogni passo temporale, rendendo il problema fortemente sottodeterminato.
• Violazione della Causalità: Molti metodi di apprendimento profondo attuali (come le reti neurali ricorrenti o i modelli "unrolled" che elaborano l'intera sequenza simultaneamente) permettono ai dati futuri di influenzare la ricostruzione di stati passati. Questo viola il principio di causalità, secondo cui lo stato di un oggetto al tempo $t'$ dipende solo dai suoi stati precedenti ($t < t'$) e non da quelli futuri ($t > t'$).
• Limiti dei metodi esistenti: I metodi puramente basati sui dati tendono ad accumulare errori di previsione nel tempo, specialmente per movimenti complessi, mentre i metodi variazionali classici richiedono regolarizzatori "handcrafted" che potrebbero non catturare bene le dinamiche temporali.

2. Metodologia: Transformer Causality Regularization (TCR)

Gli autori propongono un metodo ibrido chiamato Transformer Causality Regularization (TCR) che integra un principio di causalità appreso all'interno di un quadro di regolarizzazione variazionale classica.

• Formulazione Variazionale: Il problema di ricostruzione al tempo $t'$ è formulato come la minimizzazione di un funzionale che combina:

  1. Un termine di discrepanza dei dati ($D$) tra la proiezione forward e le misurazioni.
  2. Un termine di penalità ($P$) che misura la differenza tra la soluzione corrente e una previsione causale $\Lambda$.
     $$\min_{\vartheta(t')} \{ D[A_{t'}\vartheta(t') - \psi(t')] + \alpha_{t'} P[\vartheta(t') - \Lambda(\{\vartheta(t) : 0 \le t < t'\})] \}$$
     Qui, $\Lambda$ è una funzione appresa che predice lo stato futuro basandosi solo sugli stati passati.

• Architettura Ibrida CNN-Transformer:

  • Encoder CNN: Utilizza convoluzioni 3D per estrarre caratteristiche spaziali e temporali iniziali, sfruttando l'inductive bias delle CNN per le immagini.
  • Transformer Temporale: Utilizza un'architettura Transformer (basata su Rotary Position Embeddings - RoPE) per modellare le dipendenze temporali.
  • Maschera Causale: Viene applicata una maschera di attenzione causale durante l'addestramento e l'inferenza. Questo impedisce al modello di "vedere" i futuri passi temporali, garantendo che la previsione al tempo $t'$ dipenda esclusivamente dai tempi $t < t'$.
  • Decoder CNN: Un decoder basato su CNN (simile a UNETR) ricostruisce l'immagine finale dallo spazio latente del transformer.

• Algoritmo di Addestramento:

  • Vengono addestrati due modelli: uno per rifinire le ricostruzioni iniziali dei primi due passi temporali ($\Lambda_{re}$) e uno per prevedere il passo successivo dato un input sequenziale ($\Lambda_{pre}$).
  • L'addestramento utilizza tecniche di teacher forcing probabilistico e rollout per stabilizzare la previsione auto-regressiva.

• Risoluzione: Per ogni passo temporale, la previsione del transformer funge da prior (priorità) in un problema di ottimizzazione variazionale (risolto tramite discesa del gradiente per $L_2$ o algoritmi proximali per $L_1$ e TV), garantendo la consistenza con i dati misurati.

3. Contributi Chiave

1. Integrazione della Causalità: È il primo approccio che incorpora esplicitamente il principio di causalità come regolarizzatore appreso all'interno di un framework variazionale per problemi inversi dinamici.
2. Trasferimento Teorico: Poiché la soluzione finale è ottenuta risolvendo un problema variazionale classico, le proprietà teoriche di regolarizzazione e convergenza (note per i metodi variazionali in spazi di Lebesgue-Bochner) si trasferiscono direttamente al metodo proposto.
3. Architettura Spazio-Temporale: Unisce i vantaggi delle CNN (per lo spazio) e dei Transformer (per la sequenza temporale variabile e le dipendenze a lungo termine) in un'unica architettura coerente.
4. Generalizzazione: Il metodo non dipende dall'operatore forward specifico, rendendolo applicabile a diverse modalità di imaging.

4. Risultati

Gli esperimenti sono stati condotti su dati simulati (fantasmi dinamici) e dati reali (dataset "rolling stones" in tomografia computerizzata).

• Confronto con UAR (Unrolled Adversarial Regularizer):

  • Il TCR supera le versioni statiche e dinamiche dell'UAR (un metodo basato su GAN) in termini di PSNR e SSIM, specialmente nel caso di forte sottocampionamento (3 angoli per passo temporale).
  • Mentre l'UAR può accedere a dati futuri durante la ricostruzione (violando la causalità), il TCR mantiene la coerenza temporale.
  • Il TCR riduce la propagazione dell'errore nel tempo: le ricostruzioni TCR rimangono più stabili e continue rispetto alle previsioni pure del transformer, che tendono a degradare dopo molti passi.

• Dati Reali:

  • Applicato al dataset "rolling stones" (30 passi temporali, movimento non visto durante l'addestramento), il TCR dimostra una buona capacità di generalizzazione su dati fuori distribuzione (out-of-distribution), recuperando dettagli che le sole previsioni del transformer perderebbero.
  • L'aggiunta di un termine di Variazione Totale (TV) al termine di regolarizzazione $L_1$ migliora ulteriormente la qualità visiva e la consistenza dei dati.

5. Significato e Impatto

Questo lavoro rappresenta un passo significativo verso l'uso di modelli di deep learning avanzati (Transformer) in contesti scientifici rigorosi dove la consistenza fisica e la causalità sono fondamentali.

• Teorico: Colma il divario tra i metodi variazionali classici (garantiti teoricamente) e i metodi data-driven moderni, offrendo un framework che mantiene le garanzie di convergenza.
• Pratico: Offre una soluzione robusta per l'imaging dinamico in tempo reale (es. monitoraggio di processi industriali o imaging medico di organi in movimento), dove non è possibile attendere la fine della sequenza per ricostruire i primi frame.
• Futuro: Il metodo è promettente per applicazioni in tomografia a raggi X, risonanza magnetica e imaging a nano-scala, aprendo la strada a soluzioni che rispettano le leggi fisiche del tempo senza sacrificare la potenza predittiva dell'IA.

In sintesi, il paper dimostra che combinare l'intelligenza artificiale con principi fisici fondamentali (come la causalità) all'interno di framework matematici solidi porta a ricostruzioni più accurate, stabili e fisicamente plausibili rispetto ai metodi puramente basati sui dati.