Overcoming Distribution Shifts in Plug-and-Play Methods with Test-Time Training
This addresses distribution shift issues in computational imaging for applications like MRI, but it is incremental as it adapts existing test-time training to PnP methods.
The paper tackles the problem of performance drop in Plug-and-Play Priors (PnP) methods due to distribution shifts between training and testing data, proposing PnP-TTT which uses test-time training to improve generalization, showing it enables priors trained on natural images to work in MRI reconstruction with sufficient measurements.
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).