Overcoming Distribution Shifts in Plug-and-Play Methods with Test-Time Training (2403.10374v1)
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).
- “Convolutional neural networks for inverse problems in imaging: A review,” IEEE Signal Process. Mag., vol. 34, no. 6, pp. 85–95, 2017.
- “Using deep neural networks for inverse problems in imaging: Beyond analytical methods,” IEEE Signal Process. Mag., vol. 35, no. 1, pp. 20–36, Jan. 2018.
- “Plug-and-play methods for integrating physical and learned models in computational imaging,” IEEE Signal Process. Mag., vol. 40, no. 1, pp. 85–97, Jan. 2023.
- “Plug-and-play priors for model based reconstruction,” IEEE Glob. Conf. Signal Inf. Process., pp. 945–948, 2013.
- “Plug-and-play priors for bright field electron tomography and sparse interpolation,” IEEE Trans. Comput. Imag., vol. 2, no. 4, pp. 408–423, December 2016.
- “Dense recurrent neural networks for accelerated MRI: History-cognizant unrolling of optimization algorithms,” IEEE J. Sel. Topics Signal Process., vol. 14, no. 6, pp. 1280–1291, Oct. 2020.
- “Deep learning techniques for inverse problems in imaging,” IEEE J. Sel. Areas Inf. Theory, vol. 1, no. 1, pp. 39–56, May 2020.
- “Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing,” IEEE Signal Process. Mag., vol. 38, no. 2, pp. 18–44, Mar. 2021.
- “Deep equilibrium architectures for inverse problems in imaging,” IEEE Trans. Comput. Imag., vol. 7, pp. 1123–1133, 2021.
- “Online deep equilibrium learning for regularization by denoising,” in Proc. Advances Neural Inf. Process. Syst., New Orleans, LA, 2022.
- “Deep model-based architectures for inverse problems under mismatched priors,” IEEE J. Sel. Areas Inf. Theory, 2022.
- “Deep-learning systems for domain adaptation in computer vision: Learning transferable feature representations,” IEEE Signal Process. Mag., vol. 34, no. 6, pp. 117–129, Nov. 2017.
- “A brief review of domain adaptation,” in Advances in Data Science and Inf. Eng., Cham, 2021, pp. 877–894, Springer International Publishing.
- “Test-time training with self-supervision for generalization under distribution shifts,” in Proc. 37th Int. Conf. on Mach. Learn., 13–18 Jul 2020, vol. 119, pp. 9229–9248.
- “Test-time training can close the natural distribution shift performance gap in deep learning based compressed sensing,” in Proc. 39th Int. Conf. on Mach. Learn., 17–23 Jul 2022, vol. 162, pp. 4754–4776.
- A. Beck and M. Teboulle, “Fast gradient-based algorithm for constrained total variation image denoising and deblurring problems,” IEEE Trans. Image Process., vol. 18, no. 11, pp. 2419–2434, November 2009.
- “Beyond a Gaussian denoiser: Residual learning of deep CNN for image denoising,” IEEE Trans. Image Process., vol. 26, no. 7, pp. 3142–3155, 2017.
- “Spectral normalization for generative adversarial networks,” in Int. Conf. on Learn. Represent., 2018.
- “A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics,” in Proc. 8th Int. Conf. Comput. Vis., July 2001, vol. 2, pp. 416–423.
- J. Zhang and B. Ghanem, “Ista-net: Interpretable optimization-inspired deep network for image compressive sensing,” IEEE Conf. Comput. Vis. Pattern Recognit., pp. 1828–1837, 2018.
- Y. Nesterov, Introductory lectures on convex optimization: A basic course, vol. 87, Springer Science & Business Media, 2003.
- D. G. Anderson, “Iterative procedures for nonlinear integral equations,” J. ACM, vol. 12, no. 4, pp. 547–560, 1965.