Autoregressive Image Diffusion: Generation of Image Sequence and Application in MRI

TL;DR

The paper tackles accelerated MRI by addressing the trade-off between image quality and speed under undersampled k-space. It proposes Autoregressive Image Diffusion (AID), an autoregressive extension of diffusion models that models sequences of images and integrates MRI likelihood via a Bayesian posterior sampling scheme. AID uses a Unet-based Temporal-Spatial Conditioning architecture with causal attention and derives a posterior-sampling algorithm that combines DDIM-like steps with data fidelity. Empirical results on fastMRI brain and cardiac cine data show improved sequence coherence and reconstruction accuracy, with reduced hallucinations compared to standard diffusion baselines, highlighting AID's potential for dynamic and multi-contrast MRI applications.

Abstract

Magnetic resonance imaging (MRI) is a widely used non-invasive imaging modality. However, a persistent challenge lies in balancing image quality with imaging speed. This trade-off is primarily constrained by k-space measurements, which traverse specific trajectories in the spatial Fourier domain (k-space). These measurements are often undersampled to shorten acquisition times, resulting in image artifacts and compromised quality. Generative models learn image distributions and can be used to reconstruct high-quality images from undersampled k-space data. In this work, we present the autoregressive image diffusion (AID) model for image sequences and use it to sample the posterior for accelerated MRI reconstruction. The algorithm incorporates both undersampled k-space and pre-existing information. Models trained with fastMRI dataset are evaluated comprehensively. The results show that the AID model can robustly generate sequentially coherent image sequences. In MRI applications, the AID can outperform the standard diffusion model and reduce hallucinations, due to the learned inter-image dependencies. The project code is available at https://github.com/mrirecon/aid.

Figures (18)

• Figure 1: The interaction between the images in conditioning sequence occurs in the DiTBlock, which has a causal attention module to ensure $x_n$ is conditioned on previous images $x_{<n}$. During training, the net predicts the noise for each noisy image that is sampled from the target sequence given the conditioning sequence in parallel. During generation, the net iteratively refines the noisy input to produce a clean image, which is then appended to the conditioning sequence.
• Figure 2: (a): A sequence of images from dataset is shown in the first row and is used as conditioning to generate retrospective samples that are shown in the second row. (b): With the given sequence in (a) as a warm start, prospective samples extending it are shown.
• Figure 3: Prospective samples with cold start. The initial images generated in the cold start are not sequentially coherent, but as the sampling process continues, the model progressively generates more sequentially coherent and realistic images.
• Figure 4: (a): The folded single-coil image caused by two-times undersampling mask. (b): The comparison of unfolding ability by the autoregressive and the standard diffusion model, i.e., AID (top) and Guide (bottom). Reference image is reconstructed from k-space data without undersampling. The error is the difference between the mean, $x_\mathrm{MMSE}$, and the reference image. The "Mean+std" is the mean highlighted with confidence interval, which indicates the reconstruction by AID is more trustworthy in the region of folding artifacts.
• Figure 5: E: equispaced, R: random. (a): PSNR and (b): NRMSE of the images reconstructed from the twelve-times undersampled k-space data using the autoregressive diffusion model (AID), the standard diffusion model (Guide), and the baseline method CSGM. PSNR higher is better, and NRMSE lower is better.