Revisiting Global Token Mixing in Task-Dependent MRI Restoration: Insights from Minimal Gated CNN Baselines
Authors: Xiangjian Hou, Chao Qin, Chang Ni, Xin Wang, Chun Yuan, Xiaodong Ma
Organizations: Dept. of Radiology & Imaging Sciences, University of Utah, Salt Lake City, UT, USA · Dept. of Electrical & Computer Engineering, University of Utah, Salt Lake City, UT, USA · Dept. of Computer Vision, Mohamed bin Zayed University of Artificial Intelligence, Abu Dhabi, UAE · Dept. of Biomedical Engineering, University of Utah, Salt Lake City, UT, USA · Dept. of Electrical & Computer Engineering, University of Washington, Seattle, WA, USA
Global token mixing, implemented via self-attention or state-space sequence models, has become a popular model design choice for MRI restoration. However, MRI restoration tasks differ substantially in how their degradations vary over image and k-space domains, and in the degree to which global coupling is already imposed by physics-driven data consistency terms. In this work, we ask the question whether global token mixing is actually beneficial in each individual task across three representative settings: accelerated MRI reconstruction with explicit data consistency, MRI super-resolution with k-space center cropping, and denoising of clinical carotid MRI data with spatially heteroscedastic noise. To reduce confounding factors, we establish a controlled testbed comparing a minimal local gated CNN and its large-field variant, benchmarking them directly against state-of-the-art global models under aligned training and evaluation protocols. For accelerated MRI reconstruction, the minimal unrolled gated-CNN baseline is already highly competitive compared to recent token-mixing approaches in public reconstruction benchmarks, suggesting limited additional benefits when the forward model and data-consistency steps provide strong global constraints. For super-resolution, where low-frequency k-space data are largely preserved by the controlled low-pass degradation, local gated models remain competitive, and a lightweight large-field variant yields only modest improvements. In contrast, for denoising with pronounced spatially heteroscedastic noise, token-mixing models achieve the strongest overall performance, consistent with the need to estimate spatially varying reliability. In conclusion, our results demonstrate that the utility of global token mixing in MRI restoration is task-dependent, and it should be tailored to the underlying imaging physics and degradation structure.
Figures & tables
Figure 1: Task-dependent MRI restoration framework. (a-b) Degradations and restoration pipelines across three tasks. (c-d) Shared CNN backbone comparing local and LSConv token mixing.
Method
4 × acceleration
8 × acceleration
PSNR ↑
SSIM ↑ a
NMSE ↓
PSNR ↑
SSIM ↑ a
NMSE ↓
Zero-filled
29.49
0.6541 (0.7504)
0.0532
26.84
0.5500 (0.6392)
0.0893
CS [ 21 ]
29.54
0.5736 (-)
0.0583
26.99
0.4870 (-)
0.0903
U-Net [ 17 ]
31.88
0.7142 (-)
0.0357
29.78
0.6424 (-)
0.0511
KIKI-Net [ 5 ]
31.87
0.7172 (-)
0.0353
29.27
0.6355 (-)
0.0546
Kiu-net [ 13 ]
32.06
0.7228 (-)
0.0342
29.86
0.6456 (-)
0.0497
Table 1: Quantitative comparison on the fastMRI single coil knee dataset. PSNR and NMSE are computed under a consistent protocol for all methods, while SSIM is reported under two evaluation protocols to ensure fair comparison. SSIM values are formatted as Slice-wise (Volumetric) . Methods with missing values under a specific protocol are denoted by a dash (-).
Method
fastMRI knee
Stanford2D
PSNR ↑
SSIM ↑
NMSE ↓
PSNR ↑
SSIM ↑
NMSE ↓
Zero-filled
27.42
0.7000
0.0704
29.24
0.7175
0.0948
U-Net [ 17 ]
34.18
0.8601
0.0151
33.77
0.8953
0.0333
E2E-VarNet [ 19 ]
36.81
0.8873
0.0090
36.48
0.9220
0.0172
HUMUSNet [ 6 ]
36.84
0.8879
0.0092
35.43
0.9134
0.0219
PDAC [ 23 ]
37.12
0.8905
0.0085
36.77
0.9247
0.0166
Table 2: The quantitative results of 8× accelerated multi-coil MRI reconstruction using our proposed model and recent methods on the fastMRI knee and Stanford 2D datasets.
We present In-Token Learning, an image restoration framework that adapts a pretrained diffusion transformer using conditional rectified flow matching. Clean targets paired with degraded inputs supervise transport from Gaussian noise to restored images. Spatially aligned degraded-image tokens are fused with evolving latent tokens along the channel dimension, preserving the image-token count at a given resolution. Direct Low-Quality Guidance (DLG) combines frozen degraded-image embeddings with a fixed task prompt through the native conditioning pathway, without a trainable ControlNet-style branch or image captioning. We evaluate super-resolution and denoising on DIV2K, LSDIR, FFHQ, RealLQ250, and RealPhoto60, and automatic colorization on DIV2K and LSDIR. The tasks use separately trained checkpoints under the same framework. Results show competitive fidelity and perceptual quality under the evaluated protocols, with weaker generalization on RealLQ250. We report full-image QHD (2560×1440) inference and a tiled 12K restoration demonstration of Along the River During the Qingming Festival. Attention cost still increases with resolution. This technical report preserves the early broader study underlying Fill2SR, which subsequently developed the real-world super-resolution direction.
Xingfu Yi, Xiaoxue Yu
Independent Researcher, Hangzhou, China · Zhejiang University, Hangzhou, China
While traditional image restoration focuses on perceptual quality, Task-Driven Image Restoration (TDIR) aims to maximize the performance of downstream high-level vision tasks. Recent approaches leveraging generative priors have shown promise for TDIR; however, they typically suffer from computational inefficiency and potential semantic alteration by indiscriminately updating all latent tokens. In this paper, we posit that not all visual information is equally important for machine perception. Through an analysis of the latent token space, we observe that task-relevant cues are unevenly distributed across the token sequence, exhibiting index-wise specialization. This suggests that selectively refining a subset of tokens can be sufficient for task-driven objectives. Leveraging this insight, we propose TaskTok, a novel framework that selectively restores only task-relevant tokens via a learnable token switch and a lightweight token refinement module. Extensive experiments across image classification, semantic segmentation, and object detection demonstrate that TaskTok significantly enhances task performance with high computational efficiency. The source code is available at https://github.com/jimmy9704/TaskTok
Efficient visual tokenization requires balancing representation compactness with the preservation of fine-grained image content. Conventional grid-based tokenizers allocate tokens uniformly across spatial locations despite non-uniform visual complexity, whereas compact 1D tokenizers aggregate global information but can struggle to preserve local details. We introduce MixTok, a hybrid visual tokenizer that combines a compact 1D global representation with sparse 2D tokens for local refinement. A content-adaptive router ranks candidate locations by refinement priority and determines the number of 2D tokens for each image under a target average token budget. At the highest-scoring locations, 2D token features replace the spatial features recovered from the 1D representation, and a shared decoder reconstructs the image from the resulting mixed features. On ImageNet 256 x 256, MixTok achieves highly competitive reconstruction and generation performance, with an rFID of 1.195 and a gFID of 1.95 at respective average budgets of approximately 128 and 131 content tokens per image.
Chenge Li, Yefei He, Xinyu Liu +3
Zhejiang University · The Hong Kong University of Science and Technology · University of Adelaide