SGD Jittering: A Training Strategy for Robust and Accurate Model-Based Architectures

Peimeng Guan, Mark A. Davenport

TL;DR: This work proposes a novel stochastic (proximal) gradient descent training strategy for model-based architectures in solving inverse problems, offering improved robustness and better generalization compared to standard MSE training.

Abstract: Inverse problems aim to reconstruct unseen data from corrupted or perturbed measurements. While most work focuses on improving reconstruction quality, generalization accuracy and robustness are equally important, especially for safety-critical applications. Model-based architectures (MBAs), such as loop unrolling methods, are considered more interpretable and achieve better reconstructions. Empirical evidence suggests that MBAs are more robust to perturbations than black-box solvers, but the accuracy-robustness tradeoff in MBAs remains underexplored. In this work, we propose a simple yet effective training scheme for MBAs, called SGD jittering, which injects noise iteration-wise during reconstruction. We theoretically demonstrate that SGD jittering not only generalizes better than the standard mean squared error training but is also more robust to average-case attacks. We validate SGD jittering using denoising toy examples, seismic deconvolution, and single-coil MRI reconstruction. Both SGD jittering and its SPGD extension yield cleaner reconstructions for out-of-distribution data and demonstrates enhanced robustness against adversarial attacks.

Lay Summary: Machine learning methods for image and signal reconstruction problems (also known as inverse problems) should be both accurate and robust to noise. However, conventional training often fails to achieve both due to the robustness–accuracy tradeoff. We introduce SGD jittering, a simple yet effective noise injection training strategy for model-based architectures. Our theoretical and empirical results show that SGD jittering significantly improves robustness and generalization over conventional mean squared error (MSE) training.

Link To Code: https://github.com/InvProbs/SGD-jittering

Primary Area: Deep Learning->Robustness

Keywords: Inverse problems, robustness, generalization, robustness-accuracy tradeoff, model-based architecture, noise injection

Submission Number: 5218