PaperlayerView exact revision on arXiv

Source-linked AI summary

Next-Acceleration-Scale Prediction for Autoregressive MRI Reconstruction

Yilmaz Korkmaz, Vishal M. Patel

arXiv:2605.19354v2eess.IVcs.CV

TL;DR

Extreme undersampling makes MRI reconstruction ill-posed and can cause pixel-domain predictors to lose high-frequency anatomy. The paper uses discrete multi-scale latent next-acceleration-scale prediction with on-policy privileged information distillation, and reports strong perceptual quality and anatomical-detail preservation under large accelerations. Its main scope boundary is that quantized latent representations can limit fidelity, particularly at lower acceleration factors where continuous models may better exploit measurements.

  • Problem

    Extreme acceleration makes MRI reconstruction severely ill-posed, while pixel-domain methods can fail to preserve diagnostically important high-frequency anatomy.

  • Method

    The method casts reconstruction as next-acceleration-scale prediction in a discrete latent hierarchy and distills a fully sampled-context teacher into a student trained on its own rollouts.

  • Results

    Experiments on fastMRI under large acceleration factors show strong perceptual quality and preservation of anatomically meaningful detail, with consistent gains across diverse sampling patterns.

  • Takeaways & Limitations

    Discrete autoregressive reconstruction provides a structured prior for ambiguous, severely undersampled MRI recovery while using privileged fully sampled information only during training.

  • Takeaways & Limitations

    Quantization limits final reconstruction fidelity, and continuous models may remain more beneficial at lower acceleration factors unless higher-fidelity tokenizers are developed.

Abstract

from arXiv · show

MRI reconstruction is an inherently ill-posed inverse problem, since incomplete measurements admit many plausible solutions. This ambiguity becomes more severe under high acceleration, where pixel-domain continuous predictors tend to average over feasible reconstructions and suppress high-frequency anatomy. We address this limitation by moving reconstruction to discrete multi-scale latent space and posing it as autoregressive next-acceleration-scale prediction. Leveraging discrete priors proven effective in visual autoregressive modeling, our method restricts the solution to compact sequences of codebook tokens, enabling sharp reconstructions even from extremely sparse measurements. This discrete autoregressive formulation also aligns naturally with modern large language model post-training techniques. Building on this observation, we introduce on-policy privileged information distillation for visual autoregressive modeling, where a teacher is provided training only privileged context that is unavailable at inference, in our case fully sampled acquisitions, and supervises a student trained on its own rollouts, leading to consistent reconstruction gains. Through extensive experiments on the fastMRI benchmark, we show that our approach delivers improved reconstruction performance across diverse sampling patterns under extreme undersampling. Project website is \href{https://yilmazkorkmaz1.github.io/discrete-mri-reconstruction-opd/}{here}.

1 Introduction

Under extreme undersampling, accelerated MRI reconstruction remains severely ill-posed and pixel-domain methods can lose high-frequency anatomy. This paper addresses the problem with discrete latent next-acceleration-scale prediction and on-policy privileged information distillation.

  • Extreme acceleration leaves MRI reconstruction severely ill-posed, with existing methods often preserving global structure but losing diagnostically important high-frequency anatomy.
  • Discrete visual tokens constrain reconstruction more strictly than direct pixel-level prediction while representing structured anatomical content in compact latent sequences.
  • The proposed framework replaces resolution-wise generation with prediction across acceleration levels in a multi-input discrete latent hierarchy.
  • On-policy privileged information distillation trains a student on its own rollouts while a teacher uses fully sampled acquisitions available only during training.
  • The architecture combines an additive multi-input vector-quantized autoencoder with a shared codebook and a cross-attentive autoregressive transformer for next-scale token prediction.
  • The method is positioned within prior advances in CNN, transformer, Mamba, physics-guided, discrete-latent, and on-policy autoregressive modeling.

3 Preliminaries

Accelerated MRI reconstruction inverts an undersampled Fourier encoding operator using data consistency and a learned prior. The paper proposes learning that prior in a discrete latent space and models latent generation hierarchically across scales.

  • Accelerated MRI recovers the target image from undersampled k-space by inverting an encoding operator that combines coil sensitivities with a partial Fourier transform.
  • Conventional reconstruction learns the prior in continuous pixel space, whereas this method learns it in a discrete latent space.
  • VAR predicts each finer-resolution latent from previously generated coarser latents in a hierarchical autoregressive process.
  • The VAR joint prior factorizes into conditional distributions for successive scales, with each conditional modeled by an autoregressive transformer.

4 Next-Acceleration-Scale Prediction

The method reformulates accelerated MRI reconstruction as autoregressive prediction across acceleration scales in a discrete latent hierarchy. It combines AQ-VAE representations, a cross-attentive transformer, and on-policy privileged-information distillation to improve rollout robustness and suppress hallucinated structures.

  • The framework combines a shared-codebook AQ-VAE, a cross-attentive transformer for next-scale prediction, and post-training on-policy privileged-information distillation.
  • Next-Acceleration-Scale Prediction: Each acceleration level is predicted from all preceding levels, with FS denoting the fully sampled acquisition and K = {32, 16, 8, 4, 2}.The next-scale sequence proceeds from 32× through 16×, 8×, 4×, and 2× to FS.
  • Additive Quantized Variational Autoencoder (AQ-VAE): AQ-VAE builds a multi-input latent hierarchy from multiple acceleration levels, using shared quantized representations that are fused before decoding.Highly accelerated inputs contribute fewer tokens, while lower-acceleration inputs provide progressively richer latent detail.
  • Additive Quantized Variational Autoencoder (AQ-VAE): The hierarchy starts with an 11 × 11 token grid at 32× acceleration and increases by 1 × 1 per level until the 16 × 16 fully sampled scale.
  • Cross-Attentive Transformer Backbone: The transformer injects AQ-VAE encoder features at 64×64, 32×32, and 16×16 resolutions through cross-attention while retaining VAR self-attention and feed-forward components.The network contains 16 transformer blocks.
  • On-Policy Privileged Information Distillation: Post-training consistently improves PSNR and SSIM across sampling patterns and often preserves or improves perceptual quality.
  • On-Policy Privileged Information Distillation: On-policy distillation trains the student on its own sampled latent rollouts while a frozen teacher uses fully sampled MRI information to supervise each next-scale distribution.Reverse KL discourages unsupported predictions and helps suppress hallucinated structures.

5 Results

At R = 32 on fastMRI, the method is evaluated across three contrasts and four undersampling patterns, showing strong Cartesian performance and consistently favorable feature-space perceptual quality. Qualitative comparisons indicate sharper boundaries and better-preserved fine structures than smoother competing reconstructions.

  • Evaluation setting: At R = 32, evaluation covers T1-weighted, T2-weighted, and FLAIR contrasts across ES Cartesian-X, ES Cartesian-Y, Radial, and Gaussian-VD masks.The masks represent distinct artifact patterns and reconstruction regimes.
  • Baselines and metrics: The comparison includes CNN, transformer, unrolled physics-informed, diffusion-based, and Mamba-based reconstruction baselines.Reported metrics include PSNR, SSIM, and complementary feature-space measures such as LPIPS.
  • Quantitative results: Under ES Cartesian-X and ES Cartesian-Y, the method achieves the best PSNR on T1 and FLAIR and the best LPIPS across all three contrasts.The strongest reported Cartesian improvements occur on FLAIR.
  • Quantitative results: FLAIR PSNR improves from 18.53 to 21.29 for Cartesian-X and from 18.15 to 20.96 for Cartesian-Y relative to MambaRecon.In radial and Gaussian-VD settings, some continuous baselines achieve higher PSNR and SSIM, while the proposed method remains strongest or highly competitive in LPIPS.
  • Perceptual quality: The method achieves the best average AlexNet-LPIPS, VGG-LPIPS, and DISTS for every mask type.These feature-space metrics provide a complementary view when pixel-level metrics favor smoothing.
  • Qualitative results: Qualitative results show sharper tissue boundaries and better-preserved fine structures, whereas competing methods can show smoothing and loss of high-frequency detail despite higher per-image PSNR or SSIM.The comparisons are shown for ES Cartesian-Y and other sampling patterns with zoomed-in regions and per-image metrics.

6 Ablation Experiments

Ablation experiments show that privileged-information distillation improves fidelity and rollout robustness, while the main architectural components each contribute to performance. The reported effects are evaluated across contrasts and sampling patterns at R = 32.

  • Distillation: On-policy privileged information distillation consistently improves PSNR and SSIM across all sampling patterns and contrasts.The gains are especially clear in Cartesian settings, while LPIPS remains largely stable with only minor regressions in some cases.
  • Distillation: Distillation improves rollout robustness and reduces anatomically implausible token predictions without materially degrading perceptual quality.The teacher uses privileged fully sampled information during post-training, while inference-time inputs remain unchanged.
  • Component ablations: Removing cross-attention, token hierarchy, or the trainable encoder consistently reduces PSNR and SSIM and increases LPIPS relative to Base with argmax decoding.The ablations are averaged over all contrasts and sampling patterns on the fastMRI validation set at R = 32.
  • Experimental protocol: Table 5 reports base and distilled models across different sampling patterns on the fastMRI test set at R = 32.The table directly organizes the post-training comparison by sampling pattern.

7 Discussion and Limitations

The framework’s scope is bounded by the fidelity of its discrete tokenizer and by the extreme 32× evaluation regime. Its modular tokenizer design offers a path for future improvement, while interpretability of token behavior remains future work.

  • Scope boundary: At 32× acceleration, the results demonstrate robustness in a highly ambiguous regime rather than clinical readiness.The acquired measurements may be insufficient for consistently reliable diagnostic interpretation in routine clinical use.
  • Tokenizer limitation: The final reconstruction cannot exceed the representational precision of the learned tokenizer and codebook.This quantization bottleneck is especially relevant at lower acceleration factors, where continuous physics-informed models may exploit measurements more directly.
  • Future direction: Improved tokenizers can be used as drop-in replacements as discrete representation learning advances.The framework’s modularity provides a direct path for improving the discrete representation component.
  • Future direction: Discrete representations may expose token usage, hierarchy, and error-propagation patterns, but this interpretability analysis is left for future study.The paper presents this as a potential benefit rather than an evaluated result.
  • Distillation perspective: Continuous information can serve as privileged training-only context through cross-attention to guide student rollouts toward better target-structure preservation.The privileged information is unavailable at inference.

8 Conclusion

The paper presents a discrete autoregressive MRI reconstruction framework that predicts acceleration scales in a multi-scale latent token hierarchy. Its on-policy privileged information distillation uses fully sampled information during training to supervise student rollouts, achieving strong perceptual quality and preserving meaningful anatomy across sampling patterns.

  • The framework casts accelerated MRI recovery as next-acceleration-scale prediction in a multi-scale latent token hierarchy.
  • An additive multi-input AQ-VAE and cross-attentive transformer let measured acquisitions guide token prediction at every scale.
  • On-policy privileged information distillation trains a student on its own autoregressive rollouts while a teacher accesses fully sampled information.
  • Experiments on fastMRI under large acceleration factors show strong perceptual quality and preservation of anatomically meaningful detail across diverse sampling patterns.

Supplementary Material

The supplementary material details the AQ-VAE architecture, acquisition-aware conditioning, multi-scale tokenization, fusion, adversarial training, and optimization setup. It also describes the supporting quantitative and qualitative experiments.

  • The supplementary material covers AQ-VAE architecture and training, undersampling patterns, extended quantitative results, and qualitative results across mask types.
  • AQ-VAE uses six reconstruction levels from 32× acceleration through fully sampled input, with shared encoding, quantization, fusion, and decoding.
  • The conditional encoder adapts representations to acceleration factor and sampling pattern using label-informed FiLM-style modulation.
  • Each scale is quantized at its own token resolution, mapped to a common latent resolution, and fused through summed scale contributions.
  • Codebook training uses EMA updates and the rotation trick instead of explicit codebook regression and straight-through estimation.
  • Training combines SSIM, LPIPS, and least-squares adversarial objectives, with fixed reconstruction, perceptual, and adversarial weights.

10 Undersampling Patterns

The experiments evaluate four undersampling masks with the same nominal acceleration but different k-space geometries and artifact profiles. These include directional Cartesian aliasing, variable-density sampling, and radial streaking.

  • The study compares Equispaced Cartesian-X, Equispaced Cartesian-Y, Gaussian Variable Density, and Radial masks.
  • Equispaced Cartesian-Y undersamples along the phase-encoding direction and produces structured directional fold-over aliasing.
  • Equispaced Cartesian-X swaps the undersampling orientation of Cartesian-Y, reorienting artifacts while preserving the Cartesian acquisition structure.
  • Gaussian Variable Density samples central k-space more frequently than peripheral regions, preserving global structure while changing the reconstruction regime.
  • Radial sampling acquires angular spokes through k-space and produces globally distributed streaking from limited angular coverage.
  • Supplementary Figure 1 shows representative examples of all four undersampling patterns at 32× acceleration.

11 Additional Qualitative Results

The supplementary qualitative results compare reconstruction methods under Radial, Cartesian-X, Cartesian-Y, and Gaussian-VD undersampling. Each comparison reports per-image metrics and includes zoomed-in regions.

  • Radial undersampling results provide per-image metrics and zoomed-in regions for qualitative method comparison.
  • Equispaced Cartesian-X results provide per-image metrics and zoomed-in regions for qualitative method comparison.
  • Equispaced Cartesian-Y results provide per-image metrics and zoomed-in regions for qualitative method comparison.
  • Gaussian-VD results provide per-image metrics and zoomed-in regions for qualitative method comparison.

12 Additional Quantitative Results

The supplementary results report reconstruction performance and mean feature-space perceptual metrics on fastMRI at acceleration R = 32. They cover Gaussian-VD reconstruction and multiple sampling-method labels, with lower values preferred for all perceptual metrics.

  • Supplementary Table 1 reports reconstruction performance on fastMRI using Gaussian-VD at R = 32.
  • Supplementary Table 2 reports mean feature-space perceptual metrics across T1, T2, and FLAIR on fastMRI brain data at R = 32.Lower is better for all metrics.
  • The listed sampling-method labels are Cartesian-X, Cartesian-Y, Gaussian-VD, and Radial.
Loading 2605.19354v2…