Time-Embedded Algorithm Unrolling for Computational MRI

TL;DR

This work tackles inverse problems in MRI by introducing time-embedded algorithm unrolling, where iteration-time information is embedded into both the proximal operator and the data-fidelity term. Building on Vector AMP and diffusion-model conditioning, the approach uses time-dependent proximal operators and time-varying data-fidelity weights, realized with time-embedded U-Nets and ResNets. Empirical results on the fastMRI knee and brain datasets show state-of-the-art performance in accuracy and artifact suppression under limited data, with improved robustness across unroll depths and acceleration factors. The method generalizes to other unrolling schemes (e.g., VSQP and ADMM) and achieves favorable efficiency, making it appealing for clinical MRI where data are often limited.

Abstract

Algorithm unrolling methods have proven powerful for solving the regularized least squares problem in computational magnetic resonance imaging (MRI). These approaches unfold an iterative algorithm with a fixed number of iterations, typically alternating between a neural network-based proximal operator for regularization, a data fidelity operation and auxiliary updates with learnable parameters. While the connection to optimization methods dictate that the proximal operator network should be shared across unrolls, this can introduce artifacts or blurring. Heuristically, practitioners have shown that using distinct networks may be beneficial, but this significantly increases the number of learnable parameters, making it challenging to prevent overfitting. To address these shortcomings, by taking inspirations from proximal operators with varying thresholds in approximate message passing (AMP) and the success of time-embedding in diffusion models, we propose a time-embedded algorithm unrolling scheme for inverse problems. Specifically, we introduce a novel perspective on the iteration-dependent proximal operation in vector AMP (VAMP) and the subsequent Onsager correction in the context of algorithm unrolling, framing them as a time-embedded neural network. Similarly, the scalar weights in the data fidelity operation and its associated Onsager correction are cast as time-dependent learnable parameters. Our extensive experiments on the fastMRI dataset, spanning various acceleration rates and datasets, demonstrate that our method effectively reduces aliasing artifacts and mitigates noise amplification, achieving state-of-the-art performance. Furthermore, we show that our time-embedding strategy extends to existing algorithm unrolling approaches, enhancing reconstruction quality without increasing the computational complexity significantly.

Figures (17)

• Figure 1: Descriptions of (a) the existing unrolled architecture and (b) the proposed time-embedded unrolled architecture.
• Figure 2: Illustrations of (a) Positional Encoder generating time-dependent scaling ($\alpha^t_i$) and shift ($\beta^t_i$) features, (b) Time-embedding module for ResNet and U-Net, and (c) Architectures of ResNet and U-Net showing how time-embedded features are applied.
• Figure 3: Qualitative comparisons between the standard shared ($\spadesuit$) and unshared ($\clubsuit$) $\mathcal{R}(\cdot)$ optimization methods (VS, ADMM) and the proposed time-embedded unrolled algorithms with $T\!=\!10$ unrolls in U-Net. Top: Results for $R=4$ using PD-FS data. Middle: Results for $R=6$ using PD data. Bottom: Results for $R=8$ using Axial T2-W data. The proposed methods reduce artifacts (yellow arrow) that the shared and unshared methods fail to eliminate (pink arrow).
• Figure 4: Qualitative comparisons using ResNet (instead of U-Net in Fig. \ref{['fig:result_unet']}). The proposed methods reduce artifacts (yellow arrow) that the shared and unshared methods fail to eliminate (pink arrow).
• Figure 5: The intermediate visual results of the proposed method with ResNet-TE proximal operator at each iteration in a coronal PD slice.