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

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.

Paper

Similar papers

Peer review

Reviewer YPrH4/10 · confidence 4/52024-07-06

Summary

This paper presents a diffusion model for MRI acceleration. In which, an autoregressive image diffusion (AID) model is proposed to sequentially generate MRI image conditions on a given prior image sequences. This method is evaluated on the accelerated MRI reconstruction task using the public available dataset, fastMRI dataset.

Strengths

The figures are very helpful for the understanding of the paper. The paper is well-organized and easy to follow.

Weaknesses

1. Lack of quantitive comparison with other methods for MRI reconstruction. 2. The paper should highlight the technical contribution, as most of the context of the methods is from the existing works. 3. The paper claims the contribution on 3D and dynamic MRI, but expects more such results, not just for 2D images. 4. The paper deals with 3D reconstruction but the results only show 2D slices. It is not clear the consistency from other views like Coronal and Sagittal views.

Questions

The method is conditioned on previous images, how about the computational cost, as the method will use all previous slices? The efficiency of the proposed method should be horrible as it is a drawback of the autoregressive model, is this correct? As the cost is high, can the AIM model be constructed in the latent space?

Rating

4

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

1. The method lacks novelty as diffusion models are commonly used in MRI acceleration tasks 2. The method is more like for normal images but not for medical images like MRI; For medical studies, we care more about he anatomy plausibility while the image metrics, like MSE, may not be a good measurement for this purpose; 3. 3D medical images are not videos; we care more about the structure consistency from any of the three views.

Authorsrebuttal2024-08-09

Thank you for your feedback. Regarding your points, it's worth noting that VarNet (2018 by Hammerick et al.) predates CSGM (2021 by Jalal et al.). Furthermore, Jalal et al.'s NIPS paper includes a direct comparison between CSGM and VarNet, demonstrating that CSGM outperforms VarNet. Therefore, additional comparisons may not provide new insights. Additionally, I’d like to clarify that reconstruction differs from pure generation. The consistency of the reconstruction task is addressed in Appendix E of the manuscript, where the reconstruction results are thoroughly presented. The reason we chose CSGM as the baseline for comparison is detailed in Section 1 of the general rebuttal. Furthermore, VarNet was cited in our manuscript's introduction. It's important to emphasize that VarNet is a supervised method for MRI reconstruction. In contrast, the diffusion-based methods we explore offer more flexibility and reuse of a model. Specifically, eight VarNet models need to be trained for eight different MRI reconstruction setups, whereas only a single diffusion model is required to handle all setups, demonstrating greater efficiency and adaptability.

Authorsrebuttal2024-08-13

Dear Reviewer, Thank you sincerely for the time and effort you have dedicated to reviewing our paper as a volunteer. As today marks the deadline for author-reviewer interactions, we hope that our detailed explanations have addressed your concerns effectively. If you require any further clarification on any points, please don’t hesitate to reach out to us. We greatly appreciate your valuable contributions to this process. Wishing you a nice day!

Reviewer oZEE5/10 · confidence 4/52024-07-12

Summary

The paper proposes an autoregressive diffusion model, where each image in an MRI sequence is generated by a diffusion model, but the noise predictions of the diffusion model are autoregressively conditioned on previous MRI images in the sequence. Essentially, previously introduced single-image diffusion based MRI reconstruction [16] is upgraded to use sequence information with an autoregressive formulation. Improvements over the single-image model are demonstrated both qualitatively and quantitatively in MRI images.

Strengths

* The autoregressive formulation makes use of image sequence information and clearly outperforms single-image diffusion. * The problem formulation incorporating sequential information into single-image diffusion model predictions with the use of VQVAE encodings appears to be novel and powerful. * Various domain-specific experiments prove the effectiveness of the model in the MRI setting.

Weaknesses

Evaluation could be much more detailed: * While results are compared to the single image diffusion based approach "Guide", this model is internal. Direct comparisons to previous MRI reconstruction approaches would be much more convincing. Is there a particular reason why this could not be done? * The VQVAE and transformer-based autoregressive model appears fairly complex. What would be the result of applying **only** this model to predict the image sequence? i.e. having the autoregressive model directly predict $x_n$ given $[x_0, \dots, x_{n-1}]$, with no diffusion model? It is clear that the sequence information makes the autoregressive-diffusion combination superior to single-image diffusion. But we cannot infer the effect of diffusion here: How important is the existence of the diffusion part of the model? How superior is the autoregressive-diffusion combination compared to pure autoregressive? * Pure generative performance on standard natural image datasets is not evaluated. While authors do state this limitation in their paper, it is still a shame to be missing these results. Also, a few typos to be fixed: 1. Line 57: "As in the clinical practice of MRI, we often involves acquiring..." This start is grammatically incorrect. Maybe: "As the clinical practice of MRI often involves acquiring..."? 2. Figure 1: "... noisy image that sampled from ..." -> "... noisy image that **is** sampled from ..." 3. Line 186: "... OpenAI's guide diffusion codebase ..." -> "... OpenAI's guide**d** diffusion codebase ..."

Questions

I've addressed my questions related to evaluation in the "Weaknesses" section. If the evaluation concerns were addressed, I'd certainly have an even more positive opinion of the paper. I have one more: * One of the contributions is listed as the technique for efficiently optimizing the autoregressive loss in parallel. To me, this seems rather straightforward and directly based on autoregressive/diffusion model training processes. Would you be willing to share the particular challenges you have encountered in this process? I might be missing something here and I'd be glad to learn.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors address the limitations of their work in the paper sufficiently.

Reviewer oZEE2024-08-13

Thank you for the rebuttal!

I thank the authors for their detailed rebuttal effort. Based on the positive comparison to CSGM and some results on new datasets shown in the PDF attached to the rebuttal, I am willing to bump my rating up from 4 to 5. I find the work the authors put into their rebuttal quite respectable. I still think the total contribution is relatively low for NeurIPS: the paper is an application of a sequence-level autoregressive approach and a diffusion model, in my eyes a combination of existing approaches applied to the MRI domain. I think it is a perfectly acceptable paper for a more domain-specific conference. As I've said before, I would've also liked to see how direct prediction from sequence-level (not voxel-level) autoregression with a simple neural network would have fared compared to the diffusion approach, but that is perhaps for another work, and would not significantly affect the contribution. Thank you for your work.

Authorsrebuttal2024-08-14

Thanks for your feedback and appreciation for your reconsideration of the rating. Although the MRI-focused community at NeurIPS may be small, we hope that our work can help bridge the gap between machine learning and healthcare, showing the relevance of this interdisciplinary research at NeurIPS.

Reviewer NC3G7/10 · confidence 2/52024-07-16

Summary

This paper introduces an autoregressive image diffusion (AID) model for generating image sequences and accelerating MRI reconstruction. The model combines autoregressive and diffusion approaches to leverage inter-image dependencies, aiming to improve reconstruction from undersampled k-space data in MRI. It was trained on the fastMRI dataset using 4 NVIDIA A100 GPUs for 440,000 iterations with the Adam optimizer. The model architecture incorporates several key components, including DiTBlock, DDIM, and VQVAE. Experiments demonstrate that AID outperforms standard diffusion models in terms of PSNR and NRMSE metrics, particularly for twelve-times undersampled data. The model shows a reduction in hallucinations in reconstructed images compared to standard models. The paper provides a detailed explanation of the model's theoretical foundations, describing the autoregressive factorization of the joint distribution of image sequences and how the diffusion process is applied to each conditional probability in the factorization. The authors derive the training loss for the AID model from a common diffusion loss and present an algorithm for sampling the posterior for accelerated MRI reconstruction using AID. The model is evaluated on its ability to generate images with varying amounts of initial information, including both retrospective and prospective sampling approaches. The paper demonstrates the model's effectiveness in unfolding aliased single-coil images and shows improved reconstruction quality across various sampling masks and undersampling factors. The authors discuss the potential applications of the model in other medical imaging tasks and acknowledge limitations, proposing future work to address them.

Strengths

The paper presents a novel combination of autoregressive and diffusion models for image sequence generation, which is well-motivated, particularly for medical imaging applications like MRI reconstruction. The authors provide a comprehensive set of experiments demonstrating improved performance over standard diffusion models, effectively leveraging inter-image dependencies to enhance reconstruction quality. There is a clear demonstration of reduced hallucinations in reconstructed images using the AID model. The paper offers a detailed theoretical foundation for the proposed model, deriving the training loss and sampling algorithm in a clear and reproducible manner. The experimental methodology is robust, using appropriate datasets and metrics for evaluation across various sampling patterns and undersampling factors. The authors include both qualitative and quantitative assessments of the model's performance, providing visual examples that effectively illustrate the improvements in image quality. The model's ability to generate coherent image sequences is demonstrated through retrospective and prospective sampling. The paper discusses the potential broader impact of the model on medical imaging applications and acknowledges limitations, proposing future work and showing scientific integrity. The model architecture is clearly explained and illustrated with helpful diagrams, demonstrating flexibility in handling different types of undersampled k-space data. The authors provide a thorough comparison with a standard diffusion model baseline and discuss the model's potential for incorporating pre-existing information from other imaging modalities. The paper explores the model's performance in both image space and latent space and provides insights into the model's uncertainty estimation capabilities.

Weaknesses

The evaluation is primarily limited to medical imaging datasets, lacking comparison on standard image datasets. While the theoretical justification for the model is provided, it could be expanded further. The paper does not thoroughly discuss potential negative societal impacts or ethical considerations of the technology, and computational requirements and efficiency compared to standard methods are not extensively discussed. The paper lacks comparison with other state-of-the-art approaches beyond standard diffusion models and does not provide standard generative model metrics like FID or Inception Score. The model's sensitivity to hyperparameter choices, particularly sequence length, is not thoroughly explored, and there is limited discussion on the model's scalability to larger or more diverse datasets. The paper does not explore the model's performance on other medical imaging modalities beyond MRI and lacks a detailed analysis of the model's failure cases or limitations. There is no discussion on the interpretability of the model's decisions or outputs and no exploration of the model's robustness to adversarial attacks or noise in the input data. The authors do not discuss the potential privacy implications of using the model in medical settings or provide a comparison of training times or computational resources required versus other methods. The paper lacks discussion on the model's ability to handle multi-modal or multi-contrast imaging data. It does not explore the potential for transfer learning or fine-tuning the model on different datasets. There is limited discussion on the model's performance in low-resource or edge-computing scenarios and no exploration of the model's ability to handle out-of-distribution or rare pathological cases. The authors do not discuss the potential integration of their model with existing clinical workflows or systems.

Questions

1. What is the computational cost of the proposed method compared to standard diffusion models, both in training and inference? 2. How sensitive is the model to the choice of hyperparameters, particularly the sequence length? 3. Can the model be extended to handle multi-modal or multi-contrast imaging data? 4. What are the privacy implications of using this model in clinical settings? How does the model's performance scale with increasing dataset size or diversity? Can the model be adapted for other medical imaging modalities beyond MRI? 5. What is the potential for integrating this model into existing clinical workflows or systems? How does the model handle cases where there are significant anatomical variations between sequential images? 6. Can the model be extended to generate 3D volumetric data or time-series data? How do different k-space sampling patterns impact the model's performance? How does the model perform when there are motion artifacts or other types of image degradation? 7. Can the model be used for other tasks, such as image segmentation or anomaly detection in medical imaging?

Rating

7

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors acknowledge the limitation of not evaluating the model on common image datasets like ImageNet or CIFAR-10 and note the lack of standard generative model metrics like FID and Inception Score in their evaluation. The paper does not explicitly discuss potential negative societal impacts of the technology. The evaluation is primarily focused on MRI reconstruction, limiting insights into the model's generalizability, and the authors do not provide a comprehensive comparison with other state-of-the-art methods in image generation. The paper lacks a detailed analysis of the model's computational efficiency and resource requirements. There is limited exploration of the model's performance on diverse pathological cases or rare conditions, and the authors do not discuss the potential limitations of the autoregressive approach in certain scenarios. The paper does not address the interpretability of the model's decision-making process or its robustness to adversarial attacks or input perturbations. The authors do not explore the privacy implications of using the model in clinical settings or analyze the model's performance in low-resource or edge computing environments. There is no discussion on the potential for transfer learning or domain adaptation of the trained model, and the authors do not address the scalability of the model to larger or more diverse datasets. The paper does not explore the model's ability to handle multi-modal or multi-contrast imaging data or discuss the integration of the model with existing clinical workflows or systems. The authors do not provide an analysis of the model's failure cases or edge scenarios, and the paper lacks exploration of the model's performance on other medical imaging modalities beyond MRI. There is no discussion on the potential ethical considerations of using AI-generated medical images, and the authors do not address the model's ability to handle out-of-distribution or anomalous cases in medical imaging.

Authorsrebuttal2024-08-13

Dear Reviewer, Thank you sincerely for the time and effort you have dedicated to reviewing our paper voluntarily. As today marks the deadline for author-reviewer interactions, we hope that our detailed explanations have addressed your concerns effectively. If you require any further clarification on any points, please don’t hesitate to reach out to us. We greatly appreciate your valuable contributions to this process. Wish you a nice day!

Reviewer NC3G2024-08-14

Thank you

Dear Authors, Thank you for your kind words. I appreciate you taking the time to address my comments and questions. Your explanations are helpful and have clarified all the points. I believe the revised manuscript is significantly improved. Please don’t hesitate to reach out if you need further clarification on my comments.

Reviewer YPrH2024-08-09

Thanks the authors for their efforts in the reply. My concerns remain. 1. the baseline method is very old and there are many better ones, like VarNet. 2. We can clearly see the cross-slice artifacts from other views, this further demonstrate this is not sutable model for 3D generation tasks.

Authorsrebuttal2024-08-13

Dear Reviewer, Thank you sincerely for the time and effort you have dedicated to reviewing our paper as a volunteer. As today marks the deadline for author-reviewer interactions, we hope that our detailed explanations have addressed your concerns effectively. If you require any further clarification on any points, please don’t hesitate to reach out to us. We greatly appreciate your valuable contributions to this process. Wishing you a nice day!

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC