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

TL;DR

This paper introduces SGD jittering, a training strategy for model-based architectures that improves robustness and generalization in inverse problems, demonstrated across various imaging tasks and adversarial scenarios.

Contribution

The paper proposes SGD jittering, a novel noise injection training scheme for MBAs, enhancing robustness and generalization compared to standard methods.

Findings

SGD jittering improves robustness against perturbations and attacks.

It yields cleaner reconstructions for out-of-distribution data.

Theoretical analysis confirms better generalization and robustness.

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.