Transformer Causality Regularization for Dynamic Inverse Problems

1. Problem Statement

The paper addresses Dynamic Inverse Problems (DIPs), where the goal is to reconstruct a time-evolving object $\vartheta(t)$ from indirect, time-dependent measurements $\psi(t)$.

• Challenges:
  • Under-sampling: In many applications (e.g., dynamic CT), measurements per time step are sparse, making the problem ill-posed.
  • Causality Violation: Many existing data-driven methods (like Unrolled Adversarial Regularizers) compute reconstructions for all time steps simultaneously. This allows future data to influence past reconstructions, violating the physical principle of causality (the state at time $t'$ should depend only on $t < t'$, not $t > t'$).
  • Error Accumulation: Purely autoregressive prediction models (like standard Transformers) tend to accumulate errors over time, leading to degraded reconstruction quality in later frames.
  • Lack of Theoretical Guarantees: Purely data-driven approaches often lack the convergence and regularization guarantees provided by classical variational methods.

2. Methodology: Transformer Causality Regularization (TCR)

The authors propose TCR, a hybrid framework that integrates a learned causality function into a classical variational regularization scheme.

A. Core Concept

The method solves a minimization problem for each time step $t'$:
$$\min_{\vartheta(t')} \left\{ D[A_{t'}\vartheta(t') - \psi(t')] + \alpha_{t'} P[\vartheta(t') - \Lambda(\{\vartheta(t) : 0 \le t < t'\})] \right\}$$

• $D$ (Data Discrepancy): Ensures the solution matches the measured data (e.g., $L_2$ or $L_1$ norm).
• $P$ (Penalty): Enforces the solution to be close to a prior prediction.
• $\Lambda$ (Causality Function): A learned function that predicts the state at $t'$ based only on previous states $\{\vartheta(t)\}_{t<t'}$. This acts as a learned prior.

B. Architecture: Spatial-Temporal Transformer

To implement $\Lambda$, the authors design a specialized neural network:

• Hybrid Design: Combines CNNs (for spatial feature extraction and inductive bias) with Transformers (for temporal sequence modeling).
  • Encoder: 3D CNNs process spatial and temporal dimensions.
  • Transformer: Uses Rotary Position Embeddings (RoPE) and Causal Attention Masks. The causal mask ensures the model only attends to past time steps, strictly enforcing causality.
  • Decoder: CNN-based UNETR to reconstruct the spatial image from the transformer output.
• Two-Stage Training:
  1. Refinement Model ($\Lambda_{re}$): Refines initial reconstructions for the first two time steps (which serve as the seed).
  2. Prediction Model ($\Lambda_{pre}$): Autoregressively predicts the next frame given the sequence of previous reconstructions.

C. Optimization Algorithms

The variational problem is solved using iterative algorithms depending on the penalty choice:

• L2-TCR: Solved via Gradient Descent (Landweber iteration).
• L1-TCR: Solved via Accelerated Proximal Algorithms (FISTA) or Primal-Dual Hybrid Gradient (PDHG) if Total Variation (TV) is added.
• Key Advantage: Because the causality term $\Lambda$ does not depend on the current variable $\vartheta(t')$, standard convergence proofs for variational regularization apply directly.

3. Key Contributions

1. Causality as a Regularizer: The first framework to explicitly encode the principle of causality into a variational regularization term using a learned Transformer prior.
2. Hybrid Architecture: A novel spatial-temporal Transformer architecture that separates spatial processing (CNN) from temporal processing (Transformer with causal masking), bridging the gap between Bochner space theory and deep learning.
3. Theoretical Guarantees: By embedding the learned prior into a variational framework, the method inherits classical convergence and stability guarantees, unlike end-to-end deep learning methods.
4. Error Mitigation: The variational step corrects the accumulating errors typical of pure autoregressive Transformers by enforcing data consistency at every time step.

4. Experimental Results

The method was tested on Dynamic Computerized Tomography (CT) using both simulated phantoms and real experimental data ("rolling stones").

• Comparison Baseline: Compared against Unrolled Adversarial Regularizer (UAR), extended to dynamic settings (both static-trained and dynamic-trained versions).
• Performance Metrics (PSNR/SSIM):
  • TCR vs. UAR: L1-TCR outperformed both static and dynamic UAR versions in PSNR and SSIM, particularly in highly under-sampled scenarios (3 angles per time step).
  • Temporal Consistency: TCR produced more continuous time-series reconstructions. While UAR quality remained stable (due to access to future data), TCR showed a slight quality drop over time but maintained better structural integrity than pure Transformer predictions.
  • Feature Recovery: In sparse data cases (3 angles), TCR successfully recovered features (e.g., small circles, lines) that were distorted or lost in pure Transformer predictions.
• Generalization:
  • Tested on the "rolling stones" dataset (30 time steps, individual stone motions) which was out-of-distribution relative to the training data (10 time steps, rigid affine motion).
  • TCR successfully generalized to this new motion pattern and longer sequence length, demonstrating robustness.

5. Significance and Future Work

• Theoretical Bridge: The work successfully transfers the mathematical framework of Lebesgue-Bochner spaces (used for time-dependent PDEs) into a deep learning context, providing a rigorous foundation for dynamic inverse problems.
• Real-Time Applicability: By enforcing causality, TCR is suitable for real-time applications where future data is unavailable, unlike non-causal methods.
• Limitations & Future Directions:
  • Initial Data Requirement: The method requires high-quality initial reconstructions for the first two time steps. Future work aims to improve the refinement model to handle sparser initial data.
  • Nonlinearity: Current implementation assumes linear forward operators. Future research will explore applying TCR to nonlinear problems (e.g., MRI, Electrical Impedance Tomography).
  • Applications: Potential extension to medical imaging of moving organs and nano-CT.

In summary, TCR represents a significant step forward in dynamic inverse problems by combining the predictive power of Transformers with the stability and theoretical guarantees of variational regularization, while strictly adhering to the physical principle of causality.