Anatomical Token Uncertainty for Transformer-Guided Active MRI Acquisition
Organizations: Tel Aviv University, Faculty of Engineering, Tel Aviv, Israel · University of Cambridge, Department of Engineering, Cambridge, UK
Abstract
Full data acquisition in MRI is inherently slow, which limits clinical throughput and increases patient discomfort. Compressed Sensing MRI (CS-MRI) seeks to accelerate acquisition by reconstructing images from under-sampled k-space data, requiring both an optimal sampling trajectory and a high-fidelity reconstruction model. In this work, we propose a novel active sampling framework that leverages the inherent discrete structure of a pretrained medical image tokenizer and a latent transformer. By representing anatomy through a dictionary of quantized visual tokens, the model provides a well-defined probability distribution over the latent space. We utilize this distribution to derive a principled uncertainty measure via token entropy, which guides the active sampling process. We introduce two strategies to exploit this latent uncertainty: (1) Latent Entropy Selection (LES), projecting patch-wise token entropy into the -space domain to identify informative sampling lines, and (2) Gradient-based Entropy Optimization (GEO), which identifies regions of maximum uncertainty reduction via the -space gradient of a total latent entropy loss. We evaluate our framework on the fastMRI singlecoil Knee and Brain datasets at and acceleration. Our results demonstrate that our active policies outperform state-of-the-art baselines in perceptual metrics, and feature-based distances. Our code is available at https://github.com/levayz/TRUST-MRI.
Figures & tables
| Dataset | R | Method | Model | PSNR | SSIM | NMSE | LPIPS | DISTS | SSFD |
|---|---|---|---|---|---|---|---|---|---|
| Knee | Random | U-Net | 31.28 | 0.7250 | 0.0368 | 4.47 | 0.24 | 14.35 | |
| LOUPE [ 3 ] | U-Net | 32.21 | 0.7459 | 0.0317 | 5.26 | 0.31 | 9.41 | ||
| PUERT [ 20 ] | ISTA-Unfold [ 22 ] | 33.63 | 0.7963 | 0.0232 | 5.25 | 0.32 | 9.16 | ||
| Ada-sel [ 11 ] | VarNet [ 17 ] | 31.23 | 0.7396 | 0.0385 | 4.29 | 0.21 | 11.58 | ||
| AdaSense [ 8 ] | DDRM [ 13 ] | 31.19 | 0.6552 | 0.0420 | 4.68 | 0.24 | 11.63 | ||
| LES (Ours) | Transformer | 30.29 | 0.6490 | 0.0478 | 3.70 | 0.15 | 7.35 |
| Sampling Policy | PSNR | SSIM | LPIPS | DISTS | SSFD | |
|---|---|---|---|---|---|---|
| Random | 0 | 26.00 | 0.6387 | 3.2587 | 0.1421 | 7.4903 |
| LES | 1 | 27.07 | 0.6697 | 3.0733 | 0.1380 | 6.8490 |
| GEO | 1 | 26.95 | 0.6676 | 3.0656 | 0.1381 | 6.9482 |
| LES | 22 | 27.11 | 0.6702 | 3.0730 | 0.1379 | 6.7607 |
| GEO | 22 | 27.26 | 0.6783 | 2.9630 | 0.1315 | 6.6656 |
| Method | Step Latency | Total Time | Throughput (fps) | |
|---|---|---|---|---|
| AdaSense [ 8 ] | 8 | 8.49 0.01 s | 76.63 0.12 s | 0.01 |
| Ada-Sel [ 11 ] | 1 | 5.43 1.32 s | 5.45 0.14 s | 0.18 |
| Random (ours) | 0 | — | 112.8 2.0 ms | 8.86 |
| LES (Ours) | 8 | 229.1 17.0 ms | 1.04 0.08 s | 0.97 |
| GEO (Ours) | 8 | 350.3 6.1 ms | 2.02 0.03 s | 0.50 |