Anatomical Token Uncertainty for Transformer-Guided Active MRI Acquisition

TL;DR

This paper introduces a novel active sampling framework for accelerated MRI that uses a pretrained tokenizer and transformer to guide sampling based on uncertainty, improving image reconstruction quality.

Contribution

It proposes a new active sampling method leveraging latent token uncertainty for MRI acceleration, with two strategies: LES and GEO, outperforming existing methods.

Findings

Outperforms state-of-the-art baselines in perceptual metrics.

Effective in both knee and brain MRI datasets at high acceleration factors.

Code available at https://github.com/levayz/TRUST-MRI.

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 $k$-space domain to identify informative sampling lines, and (2) Gradient-based Entropy Optimization (GEO), which identifies regions of maximum uncertainty reduction via the $k$-space gradient of a total latent entropy loss. We evaluate our framework on the fastMRI singlecoil Knee and Brain datasets at $\times 8$ and $\times 16$ 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.