Information-Theoretic Optimization for Task-Adapted Compressed Sensing Magnetic Resonance Imaging

Xinyu Peng; Ziyang Zheng; Wenrui Dai; Duoduo Xue; Shaohui Li; Chenglin Li; Junni Zou; Hongkai Xiong

cs.LGApr 14, 2026

Information-Theoretic Optimization for Task-Adapted Compressed Sensing Magnetic Resonance Imaging

Authors: Xinyu Peng, Ziyang Zheng, Wenrui Dai, Duoduo Xue, Shaohui Li, Chenglin Li, Junni Zou, Hongkai Xiong

Organizations: Department of Computer Science and Engineering, Shanghai Jiao Tong University, Shanghai 200240, China. · Department of Electronic Engineering, Shanghai Jiao Tong University, Shanghai 200240, China. · Department of Computer Science, City University of Hong Kong, Hong Kong SAR, China. · College of Information Science and Electronic Engineering, Zhejiang University, Hangzhou 310007, China.

Abstract

Task-adapted compressed sensing magnetic resonance imaging (CS-MRI) is emerging to address the specific demands of downstream clinical tasks with significantly fewer k-space measurements than required by Nyquist sampling. However, existing task-adapted CS-MRI methods suffer from the uncertainty problem for medical diagnosis and cannot achieve adaptive sampling in end-to-end optimization with reconstruction or clinical tasks. To address these limitations, we propose the first task-adapted CS-MRI from the information-theoretic perspective to simultaneously achieve probabilistic inference for uncertainty prediction and adapt to arbitrary sampling ratios and versatile clinical applications. Specifically, we formalize the task-adapted CS-MRI optimization problem by maximizing the mutual information between undersampled k-space measurements and clinical tasks to enable probabilistic inference for addressing the uncertainty problem. We leverage amortized optimization and construct tractable variational bounds for mutual information to jointly optimize sampling, reconstruction, and task-inference models, which enables flexible sampling ratio control using a single end-to-end trained model. Furthermore, the proposed framework addresses two kinds of distinct clinical scenarios within a unified approach, i.e., i) joint task and reconstruction, where reconstruction serves as an auxiliary process to enhance task performance; and ii) task implementation with suppressed reconstruction, applicable for privacy protection. Extensive experiments on large-scale MRI datasets demonstrate that the proposed framework achieves highly competitive performance on standard metrics like Dice compared to deterministic counterpart but provides better distribution matching to the ground-truth posterior distribution as measured by the generalized energy distance (GED).

Figures & tables

Fig. 1 : Illustration of the parametric mapping $r\mapsto\mu_{\theta}(r)$ .

Fig. 2 : Illustration of Algorithm 2 for optimizing ( 20 ). At each training iteration, a pair of data

(\mathbf{x},\mathbf{t})

is sampled from the training set, along with a sampling ratio

r

drawn from

\mathcal{U}[a,b]

. Using

r

as the input, the PGN generates a sampling pattern

\mathbf{m}

, which is then used to create the measurements

\mathbf{y}

. The information in

\mathbf{y}

is encoded into the feature

\hat{\mathbf{y}}

by passing the zero-filled reconstruction

\mathcal{F}^{H}\mathbf{U_{m}}^{T}\mathbf{y}

through the measurement encoder

\mathcal{E}_{\phi^{Y}}

. Subsequently,

\hat{\mathbf{y}}

is fed into the decoders

\mathcal{D}_{\phi^{X}}

\mathcal{D}_{\phi^{T}}

, and

\mathcal{D}_{\phi^{Z}}

to calculate

q(\mathbf{x|y})

q(\mathbf{t|z,y})

q(\mathbf{z|y})

, and

\hat{q}(\mathbf{z|t,y})

. The variational lower bound (loss function) is then evaluated, and the entire MRI model comprising the PGN, encoders

\mathcal{E}_{\phi^{Y}}

\mathcal{E}_{\tilde{\phi}^{Y}}

, and decoders

\mathcal{D}_{\phi^{X}}

\mathcal{D}_{\phi^{T}}

\mathcal{D}_{\phi^{Z}}

is optimized through backpropagation.

Fig. 3 : The rate-distortion curves of various methods across different sampling ratios in terms of PSNR and SSIM on the fastMRI database [ 59 ] . The curves for LOUPE and SeqMRI are plotted at four fixed sampling ratios of 6.25%, 12.5%, 18.75%, and 25%. In contrast, the curves for other methods are plotted using 64 points sampled equidistantly from the sampling rate range of [1%, 30%] for 2D unconstrained sampling and [6.25%, 25%] for 1D Cartesian sampling .

Fig. 4 : Visualizations of MR reconstruction images and the corresponding sampling patterns at the 25% sampling rate on the fastMRI database. The rightmost column shows the ground truth MR images (top), the amplified regions (middle), and the k-space measurement (bottom). The other columns display the reconstruction images, the amplified areas, and the generated sampling patterns from top to bottom. The PSNR and SSIM values are labeled in the top right corner of the reconstruction images.

Method	Sampling Pattern	$\frac{100}{15}\times$ Acceleration		$\frac{100}{10}\times$ Acceleration		$\frac{100}{5}\times$ Acceleration
Method	Sampling Pattern	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
MD-Recon-Net [ 67 ]	Radial	37.25	0.9360	35.10	0.9120	31.05	0.8250
PnP-Parallel-MRI [ 66 ]	Radial	37.33	0.9388	35.14	0.9111	31.01	0.8221
ADMM-Net [ 64 ]	Radial	36.83	0.9306	34.46	0.8972	30.13	0.7958
Self-Supervised MRI [ 68 ]	Radial	36.58	0.9218	33.42	0.8924	30.18	0.8068
ISTA-Net+ [ 2 ]	Radial	37.07	0.9343	34.73	0.9052	30.64	0.8176

TABLE I : Results on the Brain MRI dataset from ADMM-Net. InfoMRI-Unroll-MSE means remove the uncertainty weight and train the model solely on MSE.

Fig. 5 : Reconstruction performance comparison of models trained with

p(r)=\mathcal{U}[0,0.3]

versus models trained individually at specific sampling ratios with

p(r)=\delta(r-r_{0})

. The performance is evaluated on the fastMRI database.

Fig. 6 : Sampling patterns produced by traditional sampling strategies and our PGN as $r$ increases from 0.05 to 0.3.

Methods	$r=0.05$	$r=0.15$	$r=0.25$
Uniform random	0.11	0.20	0.28
Variable density	0.51	0.52	0.51
Spectrum	1.0	1.0	1.0
InfoMRI (Ours)	0.07	0.09	0.14

TABLE II : Redundancy ratios for the sampling masks shown in Fig. 6 . Lower values denote higher sampling efficiency.

Fig. 7 : Visualization of InfoMRI at $r=0.15$ for uncertainty quantification on the Brain MRI dataset . InfoMRI-Unroll reconstructs better than InfoMRI-UNet and both variants accurately predict the reconstruction error with learned uncertainty.

Fig. 8 : Comparison of segmentation predictions on QUBIQ 2021 dataset . Pixel-wise probability methods exhibit spatially incoherent segmentation maps. LI-Net produces spatially coherent segmentation map but no variation. InfoMRI successfully recovers the spatial covariance of the segmentations, and produces diverse and spatially coherent samples.

Method	Sampling Patterns	$8\times$ Acceleration		$16\times$ Acceleration		$24\times$ Acceleration		$32\times$ Acceleration
Method	Sampling Patterns	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$
TradPat-UNet-RS	Uniform random	0.2396	0.9211	0.4265	0.8566	0.7296	0.7275	0.7685	0.7039
	Poisson disk	0.2319	0.9147	0.3124	0.8903	0.4570	0.8343	0.5226	0.8129
	Variable density	0.1573	0.9474	0.2324	0.9171	0.2801	0.9047	0.3008	0.8979
LI-Net [ 15 ]	Uniform random	0.2516	0.9104	0.3601	0.8737	0.4750	0.8311	0.4872	0.8263
	Poisson disk	0.2089	0.9213	0.2443	0.9109	0.3732	0.8707	0.3871	0.8590

TABLE III : Segmentation performance comparisons on QUBIQ 2021 dataset . Lower GED and Higher DSC indicate better segmentation performance. We use bold and underline to highlight the best and the second best, respectively. Note that

8\times

16\times

24\times

, and

32\times

accelerations correspond to sampling ratios of 1/8, 1/16, 1/24 and 1/32, respectively.

Fig. 9 : Visualization of posterior samples from $q(\mathbf{t|y})$ under different acceleration factors for our InfoMRI on QUBIQ 2021 dataset . As the acceleration factor increases, $q(\mathbf{t|y})$ generates more diverse samples to address the increasing uncertainty caused by the reduced amount of information.

Methods	$32\times$ Acceleration		$48\times$ Acceleration
Methods	PSNR $\downarrow$	Accuracy $\uparrow$	PSNR $\downarrow$	Accuracy $\uparrow$
Uniform random	11.33	0.8236	11.20	0.8061
Variable density	14.05	0.9410	12.74	0.8838
LOUPE	15.61	0.9692	13.73	0.9253
A-DPS	15.10	0.9773	13.02	0.9651
InfoMRI ( $\beta=0$ )	14.84	0.9776	12.19	0.9587

TABLE IV : Privacy preserving classification on MNIST. Higher accuracy and lower PSNR indicate better classification and privacy preserving performance. We use bold to highlight the best and underline for the second best.

Methods	$16\times$ Acceleration		$32\times$ Acceleration
Methods	PSNR $\downarrow$	AUROC $\uparrow$	PSNR $\downarrow$	AUROC $\uparrow$
Uniform random	22.58	0.8968	19.63	0.8848
Variable density	25.33	0.9079	23.81	0.8989
LOUPE	26.20	0.9129	24.85	0.9064
InfoMRI ( $\beta=0$ )	25.42	0.9206	24.00	0.9157
InfoMRI ( $\beta>0$ )	24.41	0.9152	23.56	0.9096

TABLE V : Privacy preserving diagnosis on fastMRI. Higher AUROC and lower PSNR indicate better classification and privacy preserving performance. We use bold and underline to highlight the best and the second best, respectively.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Fig. 10 : Visualizations of MR reconstruction images at the 5% and 15% sampling rates on the Brain MRI dataset . The leftmost column shows the ground truth MR images (top) and the amplified regions (bottom). The other columns display the reconstruction images and the amplified areas. The PSNR and SSIM values are labeled in the reconstruction images.

Method	2D Unconstrained		1D Cartesian
Method	BD-PSNR	BD-Rate	BD-PSNR	BD-Rate
Uniform random	-4.22	38.97%	-3.96	32.64%
Variable density	-1.72	26.45%	-	-
Spectrum based	-2.10	30.17%	-	-
Equispace	-	-	-3.06	26.06%
LOUPE	0.03	-0.017%	-0.89	14.22%

Appendix

TABLE VI : BD-PSNR gains (dB) and BD-Rate changes (%) for the baselines compared with InfoMRI on the fastMRI dataset . The results of LOUPE and SeqMRI are averaged over the four sampling ratios of 6.25%, 12.5%, 18.75%, and 25%, whereas those of other methods are calculated using all sampling ratios within the range [6.25%, 25%] (note that there are a total of 100 equidistantly sampled points within the sampling rate range of [1%, 30%]). A negative BD-PSNR indicates that the baseline provides lower quality than InfoMRI at the same sampling rate, while a positive BD-Rate means the baseline requires higher sampling rate than InfoMRI to achieve the same quality.

Method	Sampling Pattern	$2\times$ Acceleration		$4\times$ Acceleration
Method	Sampling Pattern	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$
TradPat-UNet-RS	Uniform random	0.1372	0.9548	0.1666	0.9431
	Poisson disk	0.1190	0.9572	0.1266	0.9563
	Variable density	0.1300	0.9529	0.2029	0.9266
LI-Net	Uniform random	0.1688	0.9393	0.1868	0.9328
	Poisson disk	0.1341	0.9520	0.1347	0.9520

Appendix

TABLE VII : Segmentation performance comparisons on QUBIQ 2021 dataset . Lower GED and Higher DSC indicate better segmentation performance. We use bold and underline to highlight the best and the second best, respectively.

Method	Sampling Pattern	$2\times$ Acceleration		$4\times$ Acceleration		$8\times$ Acceleration		$16\times$ Acceleration
Method	Sampling Pattern	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$
TradPat-UNet-RS	Uniform random	0.3965	0.4935	0.4028	0.4887	0.4368	0.4742	0.5058	0.4416
	Poisson disk	0.3835	0.4996	0.3859	0.4966	0.4126	0.4863	0.4503	0.4679
	Variable density	0.3913	0.4944	0.3932	0.4927	0.4236	0.4803	0.4768	0.4571
LI-Net	Uniform random	0.5038	0.4529	0.5096	0.4482	0.5433	0.4331	0.6283	0.3927
	Poisson disk	0.5008	0.4524	0.5074	0.4489	0.5115	0.4478	0.5429	0.4332

Appendix

TABLE VIII : Segmentation performance comparisons on SKM-TEA dataset . Lower GED and Higher DSC indicate better segmentation performance. We use bold and underline to highlight the best and the second best, respectively.

Method	Sampling Patterns	$8\times$ Acceleration		$16\times$ Acceleration		$24\times$ Acceleration		$32\times$ Acceleration
Method	Sampling Patterns	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$
TradPat-UNet-RS	Uniform random	0.9109	0.6423	1.1624	0.5132	1.4424	0.3612	1.5017	0.3263
	Poisson disk	0.7598	0.7099	0.9390	0.6237	1.1742	0.5018	1.1890	0.4956
	Variable density	0.6085	0.7816	0.7307	0.7299	0.8449	0.6770	0.9354	0.6300
LI-Net	Uniform random	1.0508	0.5793	1.2417	0.4783	1.4856	0.3409	1.5214	0.3204
	Poisson disk	0.9039	0.6405	1.0580	0.5644	1.2976	0.4347	1.3174	0.4247

Appendix

TABLE IX : Segmentation performance comparisons on BRISC 2025 dataset . Lower GED and higher DSC indicate better performance. We use bold and underline to highlight the best and the second best, respectively.

Fig. 11 : Demonstration of practical utility evaluated on the challenging real-world clinical BRISC 2025 dataset . Row 1: A standard clinical case where InfoMRI posterior samples exhibit high consistency, indicating reliable segmentation. Rows 2 and 3: A shared challenging case featuring a small and ambiguous lesion (sharing the same Ground Truth). Row 2: Under high ambiguity, InfoMRI generates diverse, structurally coherent posterior samples. The high inter-sample variance explicitly indicates low confidence, serving as a reliable built-in flag to prompt human expert intervention. Row 3: In contrast, deterministic pixel-wise probability maps fail to capture this diagnostic uncertainty, yielding consistently incorrect segmentation boundaries regardless of the chosen threshold.

Fig. 12 : Visual examples of fine-structure tumor segmentation at $16\times$ acceleration on the clinical BRISC 2025 dataset. The proposed method demonstrates a superior capability to accurately delineate intricate lesion boundaries and capture diverse, irregular tumor shapes, thereby maintaining strict structural integrity under significant undersampling.

Method	Sampling patterns	$8\times$ Acceleration		$16\times$ Acceleration		$24\times$ Acceleration		$32\times$ Acceleration
Method	Sampling patterns	ECE $\downarrow$	BS $\downarrow$	ECE $\downarrow$	BS $\downarrow$	ECE $\downarrow$	BS $\downarrow$	ECE $\downarrow$	BS $\downarrow$
TradPat-UNet-RS	Variable density	0.0121	0.0307	0.0204	0.0425	0.0191	0.0505	0.0318	0.0597
LI-Net	Variable density	0.0228	0.0266	0.0298	0.0337	0.0462	0.0499	0.0415	0.0459
LOUPE-UNet-RS	Learned	0.0083	0.0223	0.0129	0.0344	0.0156	0.0397	0.0364	0.0681
LOUPE-LI-Net		0.0299	0.0299	0.0319	0.0319	0.0407	0.0407	0.0532	0.0532
SemuNet		0.0141	0.0205	0.0178	0.0237	0.0155	0.0302	0.0223	0.0368

Appendix

TABLE X : Calibration performance comparisons on QUBIQ 2021 dataset . Lower BS and ECE indicate better calibration performance. We use bold and underline to highlight the best and the second best, respectively. Note that

8\times

16\times

24\times

, and

32\times

accelerations correspond to sampling ratios of 1/8, 1/16, 1/24 and 1/32, respectively.

$w_{1}$	$w_{2}$	$w_{3}$	$4\times$ Acceleration		$8\times$ Acceleration		$16\times$ Acceleration		$24\times$ Acceleration
$w_{1}$	$w_{2}$	$w_{3}$	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$	GED $\downarrow$	DSC $\uparrow$
1.0	1.0	1.0	0.7087	0.3466	0.6938	0.3053	0.6372	0.4541	0.6489	0.4356
1.0	5.0	1.0	0.5884	0.4633	0.7075	0.4463	0.6224	0.4880	0.5750	0.5480
1.0	10.0	1.0	0.2388	0.8673	0.2235	0.8826	0.2589	0.8674	0.2352	0.8855
1.0	50.0	1.0	0.0943	0.9546	0.0955	0.9542	0.1054	0.9472	0.1208	0.9380
1.0	100.0	1.0	0.0953	0.9558	0.0900	0.9574	0.1397	0.9199	0.2032	0.8981

Appendix

TABLE XI : Sensitivity analysis on hyperparameters of the loss function evaluated on QUBIQ 2021 dataset .

Fig. 13 : Visualization of reconstruction examples under different $\beta$ values on MNIST. As observed, reconstruction becomes progressively more challenging as $\beta$ increases. Top: Reconstruction; Bottom: sampling patterns.

Method	Number of Samples	GED	DSC	Infer. Time (s)
InfoMRI (Ours)	1	0.1369	0.9547	0.006
	2	0.1146	0.9535	0.011
	4	0.0936	0.9576	0.019
	6	0.0955	0.9542	0.029
	8	0.0869	0.9562	0.038
	16	0.0872	0.9563	0.075

Appendix

TABLE XII : Performance and inference time trade-offs evaluated on QUBIQ 2021 dataset . Performance results are based on subsampling with

16\times

Acceleration. Inference times are reported based on one NVIDIA GeForce RTX 4090.

Fig. 14 : Training dynamics of InfoMRI. We report the evolution of three sub-losses (Task NLL, KL Divergence, and Reconstruction NLL) and the performance metrics on the validation dataset against the number of training steps. The smooth convergence demonstrates the stability of the optimization process.

Method	Trainable Modules
SemuNet	Rec-UNet (31.0 M) + Seg-UNet (31.0 M) + LOUPE (16.4 K) = 62.1 M
TradPat-UNet-RS	Rec-UNet (31.0 M) + Seg-UNet (31.0 M) = 62.0 M
LOUPE-UNet-RS	Rec-UNet (31.0 M) + Seg-UNet (31.0 M) + LOUPE (16.4 K) = 62.1 M
LI-Net	Seg-Enc (34.1 M) + Seg-Dec (49.5 M) + Measure-Enc (34.1 M) = 117 M
LOUPE-LI-Net	Seg-Enc (34.1 M) + Seg-Dec (49.5 M) + Measure-Enc (34.1 M) + LOUPE (16.4 K) = 117 M
Tackle	Rec-UNet (31.0 M) + Seg-UNet (31.0 M) + LOUPE (16.4 K) = 62.1 M

Appendix

TABLE XIII : Training-time trainable parameters by method (task pipeline components). Counts are shown in millions (M) or thousands (K).

Method	Total Inference Params
SemuNet	Rec-UNet (31.0 M) + Seg-UNet (31.0 M) + LOUPE (16.4 K) = 62.1 M
TradPat-UNet-RS	Rec-UNet (31.0 M) + Seg-UNet (31.0 M) = 62.0 M
LOUPE-UNet-RS	Rec-UNet (31.0 M) + Seg-UNet (31.0 M) + LOUPE (16.4 K) = 62.1 M
LI-Net	Seg-Dec (49.5 M) + Measure-Enc (34.1 M) = 83.6 M
LOUPE-LI-Net	Seg-Dec (49.5 M) + Measure-Enc (34.1 M) + LOUPE (16.4 K) = 83.6 M
Tackle	Rec-UNet (31.0 M) + Seg-UNet (31.0 M) + LOUPE (16.4 K) = 62.1 M

Appendix

TABLE XIV : Inference-time total parameters by method (modules executed at test time). Counts are shown in millions (M) or thousands (K).

Fig. 15 : The GED and Dice performance on the QUBIQ 2021 dataset for InfoMRI trained with the sampling ratios sampled from the full range [0,1]. The experimental configurations are the same as Table II in the manuscript.

Explore similar work

Bayesian Uncertainty-Aware MRI Reconstruction
Mar 13, 2026Ahmed Karam Eldaly, Matteo Figini, Daniel C. AlexanderAccelerated MRI Reconstruction Uncertainty Quantification
Department of Computer Science, University of Exeter, UK · UCL Hawkes Institute, Department of Computer Science, University College London, UK
Trustworthy MRI Reconstruction via Bayesian Uncertainty Quantification with Sparsity Prior Models
Jun 15, 2026Ahmed Karam Eldaly, Matteo Figini, Daniel C. AlexanderSparse Recovery Uncertainty Quantification
Department of Computer Science, University of Exeter, Exeter EX4 4QF, U.K. · UCL Hawkes Institute, Department of Computer Science, University College London, London, U.K.
Anatomical Token Uncertainty for Transformer-Guided Active MRI Acquisition
Mar 23, 2026Lev Ayzenberg, Shady Abu-Hussein, Raja Giryes +1Magnetic Resonance Imaging Accelerated MRI Reconstruction
Tel Aviv University, Faculty of Engineering, Tel Aviv, Israel · University of Cambridge, Department of Engineering, Cambridge, UK

Mar 13, 2026J/K move · Enter open · S save

Bayesian Uncertainty-Aware MRI Reconstruction

Ahmed Karam Eldaly, Matteo Figini, Daniel C. Alexander

Department of Computer Science, University of Exeter, UK · UCL Hawkes Institute, Department of Computer Science, University College London, UK

Accelerated MRI Reconstruction Uncertainty Quantification Bayesian Inverse Problems

Open paper AskPDF arXiv

We propose a novel framework for joint magnetic resonance image reconstruction and uncertainty quantification using under-sampled k-space measurements. The problem is formulated as a Bayesian linear inverse problem, where prior distributions are assigned to the unknown model parameters. Specifically, we assume the target image is sparse in its spatial gradient and impose a total variation prior model. A Markov chain Monte Carlo (MCMC) method, based on a split-and-augmented Gibbs sampler, is then used to sample from the resulting joint posterior distribution of the unknown parameters. Experiments conducted using single- and multi-coil datasets demonstrate the superior performance of the proposed framework over optimisation-based compressed sensing algorithms. Additionally, our framework effectively quantifies uncertainty, showing strong correlation with error maps computed from reconstructed and ground-truth images.

Abstract

Explore similar work

Bayesian Uncertainty-Aware MRI Reconstruction

Trustworthy MRI Reconstruction via Bayesian Uncertainty Quantification with Sparsity Prior Models

Anatomical Token Uncertainty for Transformer-Guided Active MRI Acquisition