# Test-Time Training Can Close the Natural Distribution Shift Performance Gap in Deep Learning Based Compressed Sensing

arXiv: https://arxiv.org/abs/2204.07204 · HTML: https://mlsanity.com/paper/2204.07204/

## Key facts

- Cited by 37 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 14 Apr 2022.

## Built on by (most cited first, 30 of 37)

- [A Comprehensive Survey on Test-Time Adaptation under Distribution Shifts](https://mlsanity.com/paper/2303.15361/) (2023)
- [Test-Time Adaptation via Conjugate Pseudo-labels](https://mlsanity.com/paper/2207.09640/) (2022)
- [Score-Based Generative Models for PET Image Reconstruction](https://mlsanity.com/paper/2308.14190/) (2023)
- [SLOT: Sample-specific Language Model Optimization at Test-time](https://mlsanity.com/paper/2505.12392/) (2025)
- [SMRD: SURE-based Robust MRI Reconstruction with Diffusion Models](https://mlsanity.com/paper/2310.01799/) (2023)
- [Steerable Conditional Diffusion for Out-of-Distribution Adaptation in Medical Image Reconstruction](https://mlsanity.com/paper/2308.14409/) (2023)
- [Scaling Laws For Deep Learning Based Image Reconstruction](https://mlsanity.com/paper/2209.13435/) (2022)
- [Prior Mismatch and Adaptation in PnP-ADMM with a Nonconvex Convergence Analysis](https://mlsanity.com/paper/2310.00133/) (2023)
- [Robustness of Deep Learning for Accelerated MRI: Benefits of Diverse Training Data](https://mlsanity.com/paper/2312.10271/) (2023)
- [Deep Model-Based Architectures for Inverse Problems under Mismatched Priors](https://mlsanity.com/paper/2207.13200/) (2022)
- [Stability and Generalizability in SDE Diffusion Models with Measure-Preserving Dynamics](https://mlsanity.com/paper/2406.13652/) (2024)
- [Test-time Adaptation for Foundation Medical Segmentation Model without Parametric Updates](https://mlsanity.com/paper/2504.02008/) (2025)
- [Test-Time Training for Speech](https://mlsanity.com/paper/2309.10930/) (2023)
- [Joint Supervised and Self-supervised Learning for MRI Reconstruction](https://mlsanity.com/paper/2311.15856/) (2023)
- [Benchmarking Self-Supervised Learning Methods for Accelerated MRI Reconstruction](https://mlsanity.com/paper/2502.14009/) (2025)
- [Unsupervised Detection of Distribution Shift in Inverse Problems using Diffusion Models](https://mlsanity.com/paper/2505.11482/) (2025)
- [Towards Architecture-Agnostic Untrained Network Priors for Image Reconstruction with Frequency Regularization](https://mlsanity.com/paper/2312.09988/) (2023)
- [Overcoming Distribution Shifts in Plug-and-Play Methods with Test-Time Training](https://mlsanity.com/paper/2403.10374/) (2024)
- [Generalized Semi-Supervised Learning via Self-Supervised Feature Adaptation](https://mlsanity.com/paper/2405.20596/) (2024)
- [Test-time adaptation for image compression with distribution regularization](https://mlsanity.com/paper/2410.12191/) (2024)
- [Pruning Unrolled Networks (PUN) at Initialization for MRI Reconstruction Improves Generalization](https://mlsanity.com/paper/2412.18668/) (2024)
- [UGoDIT: Unsupervised Group Deep Image Prior Via Transferable Weights](https://mlsanity.com/paper/2505.11720/) (2025)
- [Bridging Generalization Gaps in High Content Imaging Through Online Self-Supervised Domain Adaptation](https://mlsanity.com/paper/2311.12623/) (2023)
- [MotionTTT: 2D Test-Time-Training Motion Estimation for 3D Motion Corrected MRI](https://mlsanity.com/paper/2409.09370/) (2024)
- [Efficient Test-Time Adaptation through Latent Subspace Coefficients Search](https://mlsanity.com/paper/2510.11068/) (2025)
- [Image Reconstruction via Deep Image Prior Subspaces](https://mlsanity.com/paper/2302.10279/) (2023)
- [Robustness and Exploration of Variational and Machine Learning Approaches to Inverse Problems: An Overview](https://mlsanity.com/paper/2402.12072/) (2024)
- [Resolution-Robust 3D MRI Reconstruction with 2D Diffusion Priors: Diverse-Resolution Training Outperforms Interpolation](https://mlsanity.com/paper/2412.18584/) (2024)
- [TAMER: A Test-Time Adaptive MoE-Driven Framework for EHR Representation Learning](https://mlsanity.com/paper/2501.05661/) (2025)
- [Equivariant Test-Time Training with Operator Sketching for Imaging Inverse Problems](https://mlsanity.com/paper/2411.05771/) (2024)

## Abstract (arXiv)

Deep learning based image reconstruction methods outperform traditional methods. However, neural networks suffer from a performance drop when applied to images from a different distribution than the training images. For example, a model trained for reconstructing knees in accelerated magnetic resonance imaging (MRI) does not reconstruct brains well, even though the same network trained on brains reconstructs brains perfectly well. Thus there is a distribution shift performance gap for a given neural network, defined as the difference in performance when training on a distribution P and training on another distribution Q, and evaluating both models on Q. In this work, we propose a domain adaptation method for deep learning based compressive sensing that relies on self-supervision during training paired with test-time training at inference. We show that for four natural distribution shifts, this method essentially closes the distribution shift performance gap for state-of-the-art architectures for accelerated MRI.