Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Overcoming Distribution Shifts in Plug-and-Play Methods with Test-Time Training

Edward P. Chandler, Shirin Shoushtari, Jiaming Liu, M. Salman Asif, Ulugbek S. Kamilov

TL;DR

This paper proposes PnP-Ttt as a new method for overcoming distribution shifts in PnP and shows that given a sufficient number of measurements, PnP-Ttt enables the use of image priors trained on natural images for image reconstruction in magnetic resonance imaging (MRI).

Abstract

Plug-and-Play Priors (PnP) is a well-known class of methods for solving inverse problems in computational imaging. PnP methods combine physical forward models with learned prior models specified as image denoisers. A common issue with the learned models is that of a performance drop when there is a distribution shift between the training and testing data. Test-time training (TTT) was recently proposed as a general strategy for improving the performance of learned models when training and testing data come from different distributions. In this paper, we propose PnP-TTT as a new method for overcoming distribution shifts in PnP. PnP-TTT uses deep equilibrium learning (DEQ) for optimizing a self-supervised loss at the fixed points of PnP iterations. PnP-TTT can be directly applied on a single test sample to improve the generalization of PnP. We show through simulations that given a sufficient number of measurements, PnP-TTT enables the use of image priors trained on natural images for image reconstruction in magnetic resonance imaging (MRI).

Overcoming Distribution Shifts in Plug-and-Play Methods with Test-Time Training

TL;DR

This paper proposes PnP-Ttt as a new method for overcoming distribution shifts in PnP and shows that given a sufficient number of measurements, PnP-Ttt enables the use of image priors trained on natural images for image reconstruction in magnetic resonance imaging (MRI).

Abstract

Plug-and-Play Priors (PnP) is a well-known class of methods for solving inverse problems in computational imaging. PnP methods combine physical forward models with learned prior models specified as image denoisers. A common issue with the learned models is that of a performance drop when there is a distribution shift between the training and testing data. Test-time training (TTT) was recently proposed as a general strategy for improving the performance of learned models when training and testing data come from different distributions. In this paper, we propose PnP-TTT as a new method for overcoming distribution shifts in PnP. PnP-TTT uses deep equilibrium learning (DEQ) for optimizing a self-supervised loss at the fixed points of PnP iterations. PnP-TTT can be directly applied on a single test sample to improve the generalization of PnP. We show through simulations that given a sufficient number of measurements, PnP-TTT enables the use of image priors trained on natural images for image reconstruction in magnetic resonance imaging (MRI).
Paper Structure (10 sections, 8 equations, 2 figures, 1 table, 1 algorithm)

This paper contains 10 sections, 8 equations, 2 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Inverse Problems
  4. Plug-and-Play Priors
  5. Deep Equilibrium Models
  6. Test-Time Training
  7. Our contribution
  8. Method
  9. Results
  10. Conclusion

Figures (2)

  • Figure 1: Evaluation of PnP-TTT for different sampling ratios in accelerated MRI. The left-most chart displays the best PSNR performance achieved by PnP-TTT at different sampling ratios. The remaining five charts show PSNR over all examples at each TTT iteration. Note that the best performance is above the lower baseline for all the sampling ratios; however, TTT eventually overfits to the test-time measurement, reducing performance. Additionally, note that at larger sampling ratios, the performance of PnP-TTT prior can surpass that of the matched prior due to the DEQ training.
  • Figure 2: Visual evaluation of various priors at different CS ratios for CS-MRI problem with corresponding PSNR (upper left) and SSIM (upper right). Note that natural priors (top row) performs suboptimally compared to matched MRI prior (bottom row). Additionally, note the improvement due to the usage of PnP-TTT framework for all CS ratios (middle row).