Source-linked AI summary

Stochastic Finite-Transition Learning and Exploration for Accelerated MRI Reconstruction

Muhammad Usama Mirza, Esin Öztürk Işık, Tolga Çukur

arXiv:2610.13095v1eess.IVcs.CVeess.SP

TL;DR

Accelerated MRI reconstruction must recover unacquired k-space information, while common generative methods use fine-grained dynamics and few-step approximations can compromise fidelity. FLEX-MRI directly learns finite transitions defined by k-space support changes and adds support-relaxed learning and stochastic exploration. Across two multi-coil MRI datasets, it improves reconstruction quality over competing generative methods while operating roughly two orders of magnitude faster than diffusion-based baselines.

  • Problem

    Fine-grained generative recovery requires many network evaluations, while few-step methods approximate larger transitions from finer-scale dynamics and can compromise reconstruction fidelity.

    Source: §I · pp. 1–2
  • Method

    FLEX-MRI directly learns finite transitions defined by progressive k-space support changes, using support-relaxed learning and inference-time stochastic k-space perturbation without an auxiliary generative sampler.

    Source: §I · pp. 1–2, C. Finite-Transition Modeling · p. 3
  • Results

    Across two multi-coil MRI datasets, FLEX-MRI consistently improved reconstruction quality over competing generative methods while operating roughly two orders of magnitude faster than diffusion-based baselines.

    Source: §VII · p. 9
  • Takeaways & Limitations

    Stochastic k-space perturbation induces nonzero transition-level variation within the degrees of freedom unresolved by acquired measurements, without an auxiliary generative sampler.

    Source: §I · p. 2, §VII · p. 9

    Experiments focus on supervised training for two-dimensional Cartesian MR acquisitions; broader acquisition settings and scanner and protocol variation remain to be evaluated.

    Source: §VI · p. 9

Abstract

from arXiv · show

Generative MRI reconstruction maps an initial source distribution to the target image distribution along a learned recovery trajectory. Common generative methods, such as diffusion and flow, typically learn fine-grained local dynamics whose integration defines this global mapping, hence requiring many network evaluations at inference. Existing few-step sampling approaches can lower computational burden by approximating large transformations from local dynamics, albeit at the expense of reconstruction fidelity. FLEX-MRI instead partitions recovery into a small number of finite intervals between retained states defined by progressive changes in k-space support, and directly learns the reconstruction change across each interval. Because such finite transitions can admit multiple plausible outcomes, FLEX-MRI combines support-relaxed finite-transition learning with stochastic exploration. During training, randomized k-space sampling masks define finite-transition targets with varying support, while the prediction loss remains unrestricted by each individual mask. At inference, FLEX-MRI perturbs a random subset of currently estimated k-space coefficients, allowing the learned predictor to stochastically explore alternative transition paths. Each predicted state is then reconciled with the acquired measurements through data consistency. We theoretically show that this stochastic exploration induces nonzero transition-level variation through the learned recovery map and data consistency without requiring an auxiliary generative sampler. Experiments on multi-coil MRI datasets show that FLEX-MRI outperforms competing methods while maintaining computational efficiency comparable to few-step baselines.

I. INTRODUCTION

Common generative MRI methods retain fine-grained recovery dynamics, making inference costly, while few-step approximations can compromise fidelity. FLEX-MRI directly learns finite transitions defined by k-space support changes and adds stochastic exploration without an auxiliary sampler.

I. INTRODUCTION

  • Accelerated MRI leaves unacquired k-space coefficients to be recovered, and undersampled measurements can admit multiple compatible image estimates.Generative priors constrain reconstruction using the distributional characteristics of high-quality MR images.Source: §I · p. 1
  • Diffusion and flow methods integrate fine-grained local recovery dynamics, typically requiring hundreds of sequential network evaluations for diffusion and retaining substantial inference cost despite MRI-specific conditioning.Few-step methods reduce evaluations by approximating larger transitions from those finer-scale dynamics, which can compromise reconstruction fidelity.Source: §I · p. 1

Contributions:

  • FLEX-MRI defines retained states through progressive k-space support changes and directly learns each finite transition, making the transitions used at few-step inference the training targets.Randomized masks provide targets with varying support, while the loss is not restricted to coefficients selected by each mask.Source: §I · p. 1, Contributions · p. 2
  • At inference, random k-space perturbations let the predictor explore alternative transition outcomes before data consistency restores agreement with acquired measurements.This approach represents transition-level variation without an auxiliary generative sampler.Source: §I · p. 2, Contributions · p. 2
  • The analysis shows that stochastic exploration induces a non-degenerate transition residual within the measurement-null subspace.The induced residual follows the same mean-plus-residual decomposition as finite-transition variation.Source: Contributions · p. 2, §I · p. 2

A. Many-Step Generative Recovery

Many-step generative recovery defines a global source-to-target mapping by integrating fine-grained local dynamics across intermediate states. Repeated evaluations allow refinement and measurement reconciliation but require substantial computation.

A. Many-Step Generative Recovery

  • Diffusion defines fine-grained local recovery dynamics between a source distribution, typically white Gaussian noise, and a target distribution of high-quality MR images.Acquired measurements can be incorporated during recovery through data-consistency projections or acquisition-aware conditioning.Source: A. Many-Step Generative Recovery · p. 2
  • Integration over many intermediate states defines the global source-to-target mapping and requires repeated network evaluations during reconstruction.Fine discretization offers repeated opportunities to refine uncertain estimates and reconcile them with acquired measurements.Source: A. Many-Step Generative Recovery · p. 2
  • Task-aligned trajectories can connect image states associated with undersampled and fully sampled acquisitions through progressive k-space changes, but retain fine-grained dynamics and do not necessarily reduce transition count.These trajectories use intermediate states derived from the reconstruction problem rather than a generic noise-to-image process.Source: A. Many-Step Generative Recovery · p. 2

B. Few-Step Generative Recovery

Few-step diffusion and flow methods lower network-evaluation costs while retaining recovery dynamics learned at a finer scale. Larger inference steps can compromise fidelity, and flow commonly follows a deterministic path.

B. Few-Step Generative Recovery

  • Diffusion acceleration uses implicit sampling, distillation, or consistency-based formulations to reduce network evaluations while retaining fine-grained recovery dynamics.These methods approximate finite transitions from dynamics learned at a finer scale, and larger steps can make the Gaussian reverse-kernel approximation increasingly inaccurate.Source: B. Few-Step Generative Recovery · p. 2
  • With larger diffusion sampling steps, the Gaussian reverse-kernel approximation can become increasingly inaccurate, contributing to fidelity loss under aggressive acceleration.Implicit samplers, distillation, and consistency-based methods each reduce sampling cost through approximations tied to finer-scale recovery.Source: B. Few-Step Generative Recovery · p. 2
  • Flow models integrate a continuous local velocity field in relatively few solver steps, so large steps execute transitions not learned directly and can compromise fidelity.A given current state has a single transport direction, so integration commonly follows a deterministic path and earlier errors may persist.Source: B. Few-Step Generative Recovery · p. 2

C. Finite-Transition Modeling

Finite-transition modeling directly learns mappings across retained intervals, but common formulations use an auxiliary sampler for transition variation and linear interpolation that mismatches binary k-space acquisition. FLEX-MRI instead derives states, supervision, and exploration from the acquisition process.

C. Finite-Transition Modeling

  • Common finite-transition modeling directly learns conditional mappings across retained intervals, but uses an auxiliary generative sampler to represent transition-level variation, adding cost within each transition.Its common linear interpolation of source and target endpoints does not represent intermediate binary k-space sampling patterns.Source: C. Finite-Transition Modeling · p. 2
  • FLEX-MRI defines retained states through progressive k-space support changes and directly supervises the finite transition across each interval.Randomized k-space masks provide targets with varying support.Source: C. Finite-Transition Modeling · pp. 2–3
  • Support-relaxed learning evaluates prediction loss without restricting it to coefficients selected by each transition mask, while inference-time k-space perturbation explores alternative outcomes before data consistency.The method combines direct finite-transition learning and stochastic exploration without an auxiliary generative sampler.Source: C. Finite-Transition Modeling · pp. 2–3

III. THEORY

The theory first frames few-step MRI recovery as finite transitions, then analyzes how stochastic k-space perturbations realize transition-level variation without auxiliary sampling.

III. THEORY

  • Few-step recovery is analyzed by first identifying variation introduced when intermediate states are omitted, then relating that variation to stochastic k-space perturbation without auxiliary sampling.The analysis also establishes correspondence between the induced variation and that required by the finite-transition process.Source: §III · p. 3

A. Finite-Transition View of Few-Step MRI Recovery

Few-step MRI recovery retains a small subset of states from a finer trajectory, reducing evaluations while making each retained transition represent a broader conditional distribution. Variation across omitted states must therefore be represented within each retained transition.

A. Finite-Transition View of Few-Step MRI Recovery

  • Fine-grid recovery uses hundreds to thousands of evaluations, whereas few-step recovery retains J ≪ T states and requires fewer network evaluations.The retained states are indexed along the original fine-grained trajectory.Source: A. Finite-Transition View of Few-Step… · p. 3
  • A retained transition is a conditional increment that decomposes into a mean displacement and residual variation, with acquired coefficients preserved.The residual lies in the measurement-null subspace because its acquired-coefficient component is zero.Source: A. Finite-Transition View of Few-Step… · p. 3
  • Marginalizing over omitted intermediate states adds variation across their possible values and mean displacements to the retained transition’s residual.Variation represented across several fine steps must therefore be represented within the residual of the retained transition.Source: A. Finite-Transition View of Few-Step… · p. 3
  • Flow methods map a fixed state to one next state, while transition matching represents conditional variation with an auxiliary generative model and inner sampling.The paper identifies a need for few-step methods that preserve transition-level variation without increasing computational burden.Source: A. Finite-Transition View of Few-Step… · p. 3

B. FLEX-MRI

FLEX-MRI directly learns finite transitions defined by progressive changes in k-space support. It pairs support-relaxed training with stochastic inference and analyzes how the resulting variation persists through data consistency.

B. FLEX-MRI

  • FLEX-MRI defines retained states through progressive k-space undersampling and directly learns the finite transition across each retained interval.Randomized masks provide varying-support targets, while prediction loss covers the complete update rather than only coefficients selected by each mask.Source: B. FLEX-MRI · pp. 3–4, Figure 1 · p. 4

B.1 Support-Relaxed Finite-Transition Learning:

  • Support-relaxed learning removes the transition mask from prediction error, while randomized masks generate targets spanning unacquired k-space.Across independent forward-process realizations, mask selection probability for a coefficient approaches unity as the number of realizations increases.Source: §B.1 · p. 5
  • At inference, random k-space perturbation gives the finite-transition predictor alternative inputs, and data consistency restores acquired measurements after prediction.The inference perturbation can remove coefficients in the acquired set as well as unacquired coefficients; data consistency subsequently restores the acquired measurements.Source: §B.1 · p. 5

B.3 Propagation of Conditional Variation:

  • For 0 < γj < 1, the residual after the predictor and data consistency has nonzero conditional covariance for almost every admissible predictor Jacobian.The resulting residual occupies the measurement-null subspace, matching the mean-plus-residual role of finite-transition variation.Source: §B.3 · p. 6

A. Datasets

The experiments use multi-coil fastMRI and M4Raw data across multiple contrasts and retrospective undersampling rates. The supplied visual descriptions identify representative fastMRI reconstructions and performance metrics, but do not provide outcome values.

A. Datasets

  • Experiments use multi-coil complex k-space from high-field fastMRI and low-field M4Raw, covering T1, T2, and FLAIR contrasts.The M4Raw split comprises 128, 30, and 25 subjects and 9,000 cross-sections; fastMRI uses 240, 60, and 120 subjects and 12,600 cross-sections.Source: A. Datasets · p. 6
  • The datasets are retrospectively undersampled in two dimensions at rates R = 4–8 using a normal-density distribution.FastMRI is coil-compressed to five virtual coils, while M4Raw uses its native four-channel acquisition.Source: A. Datasets · p. 6
  • The representative fastMRI FLAIR reconstruction at R = 8 is shown with a fully sampled reference, an inset region of interest, and an error map.The caption identifies the comparison elements but does not state a visual winner.Source: Figure 3 · p. 7
  • The fastMRI performance table reports test-set PSNR (dB) and SSIM (%) as mean ± standard deviation across R = 4–8.The caption says boldface marks the highest-performing method but supplies no metric values.Source: Table II · p. 7

B. Competing Methods

FLEX-MRI achieves the highest PSNR and SSIM across acceleration rates on fastMRI and M4Raw, and in cross-dataset testing. It retains few-step inference efficiency and moderate memory requirements.

B. Competing Methods

  • The comparison covers nine methods spanning low-rank, adversarial, diffusion, flow-matching, and transition-matching reconstruction.FLEX-MRI uses a U-Net predictor, R = 6 endpoint, J = 10 retained transitions, and a perturbation rate of γ_j = 0.1; methods used identical data splits and undersampling patterns.Source: B. Competing Methods · p. 7

A. Comparison Studies

  • 2.7dB PSNR and 4.4% SSIM: FLEX-MRI improves over diffusion baselines on fastMRI, and by 4.6dB and 11.3% over few-step flow- and transition-matching baselines.It attains the highest PSNR and SSIM across R (p < 0.05).Source: A. Comparison Studies · p. 7
  • 7.0dB PSNR and 12.9% SSIM: FLEX-MRI improves over diffusion baselines in cross-dataset testing, and by 3.2dB and 5.3% over few-step baselines.Models trained on M4Raw were tested on fastMRI; transfer reduced FLEX-MRI performance by only 0.8dB PSNR and 0.4% SSIM on average relative to in-domain fastMRI training.Source: A. Comparison Studies · pp. 7–8
  • 4.1dB PSNR and 35.1% SSIM: FLEX-MRI improves over diffusion baselines on M4Raw, and by 2.1dB and 23.3% over few-step baselines.It attains the highest PSNR and SSIM across R (p < 0.05).Source: A. Comparison Studies · p. 7, Table III · p. 8

B. Computational Efficiency

  • Approximately two orders of magnitude faster than 1000-step diffusion baselines, FLEX-MRI runs in the same sub-second regime as few-step methods and requires 1.05 GB peak memory.The diffusion methods require 1.71–3.19 GB peak memory.Source: B. Computational Efficiency · p. 8

C. Ablation Studies

FLEX-MRI’s ablations support inference-time stochastic exploration and support-relaxed learning, with reconstruction quality peaking at J = 10. The discussion and conclusion connect its transition-level variation to unresolved measurement degrees of freedom and report improved quality with few-step computational efficiency.

C. Ablation Studies

  • Removing inference-time stochastic exploration or replacing support-relaxed learning with support-restricted learning degrades reconstruction quality, while perturbing training inputs adds no benefit.The ablations support forward-process states during training, support-relaxed finite-transition learning, and stochastic perturbation confined to inference.Source: C. Ablation Studies · p. 8

VI. DISCUSSION

  • The discussion links k-space perturbation to residual variation within measurement-unresolved degrees of freedom, while ablations support stochastic exploration and support-relaxed learning.Source: §VI · pp. 8–9
  • FLEX-MRI uses a fixed perturbation rate and number of retained transitions; adapting exploration along the recovery trajectory is a proposed extension.Source: §VI · p. 9
  • Experiments focus on supervised training for 2D Cartesian acquisitions; non-Cartesian sampling, volumetric reconstruction, and broader scanner and protocol variation remain to be evaluated.Source: §VI · p. 9
  • Trajectory variation may characterize weakly constrained image content, but radiological evaluation is needed because reference-image distortion metrics cannot establish diagnostic preservation or rule out unsupported content.Source: §VI · p. 9

C. Ablation Studies

  • 1.7 dB PSNR and 1.1% SSIM below J = 10 at J = 2; at J = 1000, PSNR and SSIM are 2.3 dB and 7.8% lower.Quality peaks at J = 10; inference cost scales approximately linearly with J. Smaller J spans larger trajectory intervals, while larger J adds perturbation events and prediction-error accumulation.Source: C. Ablation Studies · p. 8, Figure 5 · p. 9

VII. CONCLUSION

  • Across two multi-coil MRI datasets, FLEX-MRI consistently improves reconstruction quality over competing generative methods while running roughly two orders of magnitude faster than diffusion-based baselines.Source: §VII · p. 9
Loading 2610.13095v1…