SGD Jittering: A Training Strategy for Robust and Accurate Model-Based Architectures
Authors: Peimeng Guan
Mark A. Davenport
Venue: ICML 2025 poster (2025)
Keywords: Inverse problems
robustness
generalization
robustness-accuracy tradeoff
model-based architecture
noise injection
Primary area: deep_learning->robustness
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.

SGD Jittering: A Training Strategy for Robust and Accurate Model-Based Architectures
Peimeng Guan
Mark A. Davenport!
Abstract
Machine learning algorithms show significantly better re- constructions than classical optimization methods. Most of the work in the field aims to promote reconstruction quality using different deep-learning models, such as transformers (Fabian et al., 2022) and diffusion models (Whang et al., 2022). In particular, model-based architectures (MBAs) such as loop unrolling (LU) networks inspired by classical optimization methods, unroll the objective function into sequence of trainable iterative steps with the forward model A involved in every iteration. MBAs are considered to be more interpretable than neural networks that learn direct inverse without using A (Monga et al., 2021), and achieve state-of-the-ar reconstructions different fields, such as natural image restoration (Gilton al., 2021) and medical image reconstruction (Bian, 2024).
Inverse problems aim to reconstruct unseen data from corrupted or perturbed measurements. While most work focuses on improving recon- struction quality, generalization accuracy and ro- bustness are equally important, especially for safety-critical applications. Model-based architec- tures (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-roul accuracy-robustness robuuttnsss tradcoff in MBAs remains underexplored. In this work, we propose a simple yet effective training scheme for MBAs, called SGD jittering, which injects noise iteration-wis during reconstruction. We theo- retically 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 duitt--distributtin data and demonstrates enhanced robustness against adver- sarial attacks.
While reconstruction quality is crucial, robustness and gen- eralization accuracy are equally important for IPs. Robust- ness refers to the model's ability to handle unwanted noise, while generalization accuracy measures how well the model performs on unseen data. For example, in medical image reconstruction, a model must generalize well to detect subtle details like small bone fractures and be robust enough to eliminate artifacts caused by measurement noise. Despite their importance, robustness and generalization have not been thoroughly explored in the context of IPs, particularly in MBAs. While some studies provide empirical assess- ments of robustness (Gilton al., 2021; Genzel al., 2022), there is noticeable lack of theoretical analysis (Gottschling et al., 2025). In this work, model MBA training as a bilevel optimiza- tion problem without making specific assumptions about neural network architectures. We theoretically analyze its generalization and robustness risks in solving denoising problems and demonstrate that standard mean squared error (MSE) training may perform poorly in both metrics.
1. Introduction
Inverse problems (IPs) aim to recover the underlying signal a E IR" from corrupted observation y E R where m is often less than n, leading to non-unique solutions. IPs are typically modeled as:
Ax + 2,
(1)
where A Rmxn is known forward model for linear IPs, and z represents unknown noise in the observation.
Technologies like adversarial training (AT) (Fawzi et al., 2016; Gilmer et al., 2019; Kannan al., 2018) and input noise injection (Fawzi et al., 2016; Gilmer et al., 2019; Kannan et al., 2018) are commonly used to enhance robust- ness, while approaches such as weight perturbation (Wu et al., 2020), stochastic training (Hardt al., 2016) and sharpness-awe aware minimization (Foret et al., 2021) are de-
Department of Electrical and Computer Engineering, Georgia Institute of Technology, Atlanta, USA. Correspondence to: M. D. <mdav gatech.edu>>
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Stochastic (Proximal) Gradient Descent Jittering
signed to improve generalization. However, there is often tradeoff between robustness and accuracy (Tsipras al., 2019; Zhang et al., 2023; Gottschling et al., 2025; Krainovic et al., 2023). Strategies like careful loss function design (Pang et al., 2022; Kannan et al., 2018; Zhang et al., 2019; Zheng et al., 2021), data augmentation (Zhang et al., 2018), and robust architectural choices such as dropout (Yang et al., 2020) and teacher-studen models (Zhao et al., 2024) have been explored to mitigate this tradeoff in classification tasks, but the challenge of balancing accuracy and robustness in IP solvers still remains an open question (Gottschling et al., 2025). Recent work by Krainovic et al. (2023) introduces jittering in y for black-box IP neural networks, achieving a comparable level of worst-case robustness to AT but with significant accuracy loss.
2. Related Work
Robustness and Generalization for IPs Several strate- gies have been proposed to improve robustness in solving in- verse problems, including training-time exposure to diverse perturbations such noise injection and data augmentation (Krainovic et al., 2023; Zhou et al., 2019), and adversar- ial training (Calivá et al., 2021), enabling models to better handle noise during inference. Other approaches focus on architectural innovations- example, diffusion models have shown inherent robustness to noise in MRI reconstruc- tion (Güngör et al., 2023), while PINN (Peng et al., 2022) embed domain knowledge to enhance stability without com- promising interpretability.. To improve generalization in IPs, methods like data augmentation with synthetic perturbations (Guan et al., 2024) and domain adaptation via techniques such as CycleGAN (Zhu et al., 2017) have been used to expand the training distribution and improve adaptability to unseen data. Additionally, incorporating geometric con- straints, in (Jiang et al., 2022) for electrocardiograpicc image reconstruction, has shown improved generalization by embedding prior knowledge into the learning process. While these approaches typically target either robustness or generalization, achieving both simultaneously remains an open and challenging problem.
To bridge this gap, this work aims to improve robustness through novel training scheme for MBAs, while main- taining generalization accuracy for IPs. Note that while specially designed architectures can also improve robust- ness or generalization, they are not discussed here as this work focuses on investigating different training schemes.
Our Contributions In this work, we propose a stochastic gradient descent (SGD) jittering scheme in training MBAs. This method unrolls the objective function for solving IPs us- ing a gradient descent algorithm, with small, random noises added to the gradient updates at each iteration to improve the training process. During inference, these jittering noises are removed, allowing the model to perform as it would in stan- dard evaluation settings. This technique has demonstrated greater robustness and better generalization than traditional MSE training, while also being more time-efficient than AT. Our contribution can be summarized as follows,
Analysis of Robustness Smoothness is closely related to the robustness of neural networks. Works like Salman et al. (2019); Cohen et al. (2019); Lecuyer et al. (2019); Li et al. (2018) show that regularization in smoothness promotes the robustness of classifiers. On the other hand, the analysis of IP is very different from classifications. Krainovic et al. (2023) analyzes the worst-case robustness for IPs using a simple linear reconstruction model and directly solves for robustness risk. However, optimizing a convex model results in a solution that excels in only one metric (accuracy or robustness). In this work, we explore nonlinear MBAs in their general forms, demonstrating improvements across both metrics.
Network Assumptions: We investigate the performance of MBAs with general non-convex neural networks, where training often results in non-unique parameters. This analysis is crucial for demonstrating improve- ments across various metrics, particularly in robustness and generalization accuracy.
Analysis of Generalization Demonstrating better gener- alization of classifiers often involves showing larger classi- fication margins (Cao et al., 2019; Kawaguchi et al., 2022; Lim et al., 2021), but the proof for IPs is very different. A common but indirect approach shows an algorithm can regularize the Hessian loss landscape with respect to the network parameters. Models that converge to flat minima in the loss landscape tend to generalize better in practice (Jiang* al., 2020; Keskar et al., 2017; Neyshabur et al., 2017; Foret et al., 2021), although the precise theoretical mechanism underlying this phenomenon remains unclear. Many studies have shown that regularizing the Hessian of the loss landscape (Xie et al., 2021; Keskar et al., 2017: Foret et al., 2021) leads to flatter minima and have directly
Robustesss: We theoretically prove that SGD jitter- ing implicitly promotes average-case robustness than regular MSE training. In the experiments, we further demonstrate the effectiveness of SGD jittering under worst-case attacks.
Generalization: To the best of our knowledge, this is the first work to formally analyze generalization accuracy in inverse problems under small perturbations in the test data, by reformulating the generalization risk specific to inverse problems. This goes beyond the commonly observed heuristic that a flatter loss landscape implies better generalization.
2

---

Stochastic (Proximal) Gradient Descent Jittering
linked this to reduced generalization error, but this connec- tion still lacks comprehensive theoretical justification. In this work, rather than focusing on the connection to flat landscapes, we proposed new generalization risk for IPs, and demonstrate that SGD jittering results in a smaller gen- eralization risk than regular MSE training.
the corresponding observation Vy Axg + == y + Ag. Definition 3.1. The average-case robustness risk of Hg measures the distance between the unperturbed to the recovered signal &. Hy) with the presence of g.
Re B.. Py || Ho(v,))
Definition 3.2. The generalization risk of Hp measures the distance between the actual ground-truth Φ to the recovered signal Eg = (by with the presence of g.
Noise Injections Noise injection serves as an implicit regularization technique during training to enhance model robustness and generalization. When applied to input data, it can be viewed as a data augmentation strategy. This ap- proach has been extensively studied classification tasks (Fawzi et al., 2016; Gilmer et al., 2019; Rusak et al., 2020) and has recently demonstrated effectiveness in improving robustness for IPs (Krainovic et al., 2023). Layer-wise noise injection (Lim et al., 2021) in a recurrent network has been shown to generalize better for classifiers. Weight injec- tion perturbs network parameters during training, which is shown to help escape from local minima (Zhu et al., 2019; Nguyen et al., 2019) and promote smooth solutions (Orvieto al., 2023; Liu et al., 2021; Orvieto et al., 2022). Weights can also be perturbed before parameter updates to enhance generalization, as in stochastic gradient descent (SGD). We will explore SGD in detail in the next subsection. In a similar vein, Renaud et al. (2024) adds noise to plug-and and- play algorithms during evaluation to enhance in-distribution performance, though it does not address its impact on ro- bustness and generalization accuracy.
((0) = B. || - Ho (v,)
(3)
Robustness measures how close the reconstruction is com- pared to the data within the dataset, and the generalization accuracy measures how well the reconstruction obeys the physics of the IP even for out-of-distribution (OOD) data. Defining generalization in regression is not straightforward, but for IPs, we can achieve this using the underlying for- ward models. Typically, tradeoff between robustness and accuracy exists for classical feedforward neural networks (Zhang et al., 2019; Yang et al., 2020; Stutz et al., 2019). We aim to illustrate that the SGD jittering technique enhances robustness and accuracy for MBAs without incorporating explicit regularization.
4. Model-based Architecture Overview
MBAs draw inspiration from classical optimization tech- niques, where the underlying signal æ is estimated via,
Connection to SGD Training SGD training is effective in promoting generalization through implicit regularization (Smith et al., 2021). In large overparameterized models with non-unique global solutions, SGD favors networks with flat minima (Hochreiter & Schmidhuber, 1997). The implicit regularization properties of SGD have been explored in con- texts such as linear regression (Zou et al., 2021), logistic regression (Ji & Telgarsky, 2019), and convolutional net- works (Gunasekar et al., 2018). Unlike SGD training, which adds noise to network parameters, the proposed SGD jit- tering method introduces noise at each iterative update of the reconstructions in the lower-level optimization. Despite its name, SGD jittering is structurally similar to layer-wise noise injection within MBAs.
Φ = arg min
1l Ax + r(a),
where r is an arbitrary regularization function to ensure the uniqueness of the reconstruction. When r is differentiable, the optimal solution can be solved via gradient descent (GD) algorithm, iterating through steps k = 1,2, 3... with a constant step-size 17.
(4) In a loop unrolling (LU) architecture using GD, the gradient Vr is replaced by neural network. denoted by fo. The process begins with an initial estimate 2o, which is itera- tively updated using (4) with a fixed number of iterations, denoted by K. The final output æx represents the predicted reconstruction and is compared to the desired ground-trute truth. The weights 0 are then updated through end-to-end back- propagation. Different variants of LU are summarized in Monga et al. (2021).
L 21 (AAs- - y) (rxx).
3. Definition of Robustness and Generalization
We formally define average-case robustness and generaliza- tion risks for learned inverse mapping Ha y + x in IPs. Average-case robustness is also known as the "natu- ral" robustness for average performance (Rice et al., 2021; Hendrycks & Dietterich, 2019), which is equally impor- tant as in corst-case. During evaluation, a small vector g sampled from the distribution P, is added to a data point a E D. The new ground-trute truth becomes F. = x. + g, with
5. MBA Training Scheme Overview
In this section, we summarize some common training schemes for general regression problems, adapted here for MBA training. Following this, we introduce the proposed SGD jittering method.
3

---

Stochastic (Proximal) Gradient Descent Jittering
Regular MSE Training Mean-squared error (MSE) is the most commonly used training loss for regression and IPs. In particular, the goal of training a MBA is to minimize the following risk w.r.t. 8,
the input jittering approach minimizes,
Sow (0) E (w. -D - 2112 where,
(7)
# arg min
1 + w
A + ra(e).
H(0)= = Em |0 2 where,
This method learns an inverse mapping from y + w to æ, which might suffer from generalization issues in the same way as AT. In fact, Krainovic et al. (2023) demonstrates a significant accuracy loss using input jittering for black-box neural networks.
1 2 ll Acll2 + ra(w))
(5)
ē = arg min
Notice that MBAs do not directly learn r, but instead learn its gradient Vr using a neural network fo. We denote rg to highlight its dependency of 0. The lower-level optimization implemented using iterative steps, where each iteration corresponding to a single module in MBA. For LU archi- tecture, with finite and fixed number of iterations K, MSE training minimizes the squared distance between the output æ and the ground-truth æ. However, this training scheme can overfit to patterns in training data, potentially making the model more vulnerable to variations in y.
6. Proposed Training Scheme
Another line of work injects noises layer-wise into the neu- ral network, where Hodgkinson et al. (2021); Jim et al. (1996); Liu et al. (2020) model this layer-wise injection as a stochastic differential equation, demonstrating its role in im- plicit regularization for robustness. Additionally, Lim et al. (2021) introduces noise into recurrent neural networks and shows improved generalization for classification tasks. In this work, we propose adding random zero-mean Gaussian noises, Wy. to the gradient updates fo(#)). with the noise re-sampled independently for each iteration k = The noisy gradients are unbiased, i.e., fo + wsg [[]((a))] for all æ. Let = (w), .... wxx J. The estimated reconstruction is solved iteratively using this noisy gradi- ent descent, which is essentially stochastic gradient descent (SGD). Thus, with step size 17, the learning objective to minimize the following risk,
Ad Adversarial Training AT is widely recognized as one of the most effective methods for training robust neural net- works (Wang et al., 2019). AT finds the worst-case attack for each training instance and adds it to the input data designed to deceive the model. By learning from these challenging examples, AT improves the model's robustness. For MBA training, AT aims to minimize the worst-case robustness risk as follows,
|| #112
max. e:||ellase
R. (0) Ea
I 1タ2 ly + Aml
à arg min
(=))
Jows (0) Es (m..)) (m..)) >D)))wll 2 where,
(6)
a = lim w2 (1A (Axt - 4) + {o(ws) + wwx)
Fast gradient sign method and projected gradient descent (ProjGD) are common methods to find the attack vector e, where Athalye et al. (2018) demonstrates that ProjGD adversarial attacks are more robust. However, AT is slow as it requires iterative solvers for the attack vector (Shafahi et al., 2019).
8
(8)
The jittering noise satisfy the following assumption. Assumption 6.1. (SGD noises) Assume that at each gra- 3.2,3...., noise Wn.
dient descent iteration (O, L. /nI) independently sampled from an i..d zero- mean Gaussian distribution, where EE| w o
AT might suffer from poor generalizations. Imagine in noiseless IPs where y = Ax, and assume the inverse A-1 exists. The goal of AT is to learn an inverse mapping He
In AT and input jittering training, the noise vectors are added to y throughout the entire inversion, introducing a bias the reconstruction compared to the actual ground- truth. In contrast, the lower-level objective for SGD jittering is the same in regular MSE training. Extensive research has explored the convergence properties of SGD in solving nonconvex optimization problems (Ghadimi & Lan, 2013; Li et al., 2024; Madden et al., 2024; Lei et al., 2019), which ensures a correct inverse mapping even with the presence of noise, thus preserves high accuracy in reconstruction.
such that Ho + e) is close to x, for non-zero vector e. However, the true estimated is 41 (y + e) = x + A -ee. AT learns to ignore the latter term, thus introducing errors that can degrade generalization accuracy.
Input Jittering Noise injection to the input is widely used in general ML tasks to promote robustness (Salman al., 2019; Cohen et al., 2019; Lecuyer et al., 2019). In IPs that aim to map from y x, Krainovic et al. (2023) proposes adding zero-mean random Gaussian vector w to y as input to a neural network that learns a direct reconstruction. They demonstrate that, with careful selection of the jittering variance, this training scheme can attain a comparable level of worst-case robustness to that achieved by AT. In MBAs,
We adapted the SGD convergence result in (Garrigos & Gower, 2023) for SGD jittering as follows, where the proof of Corollary 6.2 can be found in the Appendix.
Corollary 6.2. (Convergence of MBAs-SGD SGD for IPs, modi- fied from Theorem 12 in (Garrigos & Gower, 2023)) Let
4

---

Stochastic (Proximal) Gradient Descent Jittering
F(x)= - Axll + ra) denote the lower-level ob- jective function in (8). Assume fo Vro is T-I.ipschitz continuous, satisfying
Under assumptions 6.1, 7.1 and 7.2, SGD jilterinn general- izes better than regular MSE training,
G( S G(
llfo - A) - fo(æ))l2 S 21 Vx,
The detailed proof is provided in Appendix C. SGD jittering can be interpreted as a noisy version of the regular MSE training. The idea is decompose the generalization ac- curacy in the presence of g into two terms: the MSE term corresponding to the reconstruction from the clean trajectory Ho(y), and penalty term depends on the noisy trajectory Ho(y - g). Then we show that the training loss of SGD jit- tering includes an implicit regularization term that penalizes deviations in the intermediate reconstructions caused by the perturbation. Minimizing the SGD jittering training loss also minimizes the penalty term in (()), leading to better generalization accuracy compared to regular MSE training.
Consider a sequence ( generated by SGD in (8) with a constant step-size 17 = V ILmack for K 2 1, the
2
following holds:
2T. - F))))))
min V 11 S
OSKCK
K
where, x² arg min F(x)), and
A3 = F(x") inl( (211v lB AxlB3 in)),
denote X as the eigenvalues of A A and Lmax max L, max; )=}-
On the other hand, SGD jittering noise also implicitly en- hances robustness. While we will analyze this relationship theoretically in the following section, intuitively, exposing neural network to noisier inputs during training encourages it to converge to smoother regions of the loss landscape, which are less sensitive to input variations. This results in more stable outputs that are less affected by perturbations during evaluation.
Furthermore, we rewrite the regularization term derived in the proof as follows in SGD training via Taylor expansion.
K-1
regularization = E Σ n(1 ((l-n)) ((o(())((( (((
(10)
Let ( K 0==0 and {=/ denote the reconstruction
trajectory from SGD and GD, respectively. As derived through iterative expansions in the proof, at iteration k,
7. Analysis for Denoising Problem
In this section, we analyze the generalization and robustness of SGD jittering in addressing the denoising problem. We begin by outlining the relevant assumptions.
7(1 - (m())"))) - Σt (1 = As- sumption 7.2 ensures the existence of the first and second derivatives of fo. Let fg R" -> R be the function defined by (a() [fo())) representing the ith component of fo(a)) for i E (1....,m.. Then, we have,
we have 2x sgd = 0'g òR. where òz = E k-1 )))
Assumption 7.1. (Denoising) The forward model in denois- ing is identity mapping, so that y = x + z. Signals æ E IR" are sampled from distribution X, and noises are sampled i..dd from z (0, o22n/ /m.l).
Assumption 7.2. (Neural Network) Assume the neural net- work architecture is twice differentiable with respect to the input, so the Hessian exists.
fa(21")) = fo(x)) + o% fa(a)) + 6% H2
where, V and Hj denote the gradient and the Hessian of fgg As & increases, the regularization weight m/1- (111))x11- K-1-i in (10) becomes more significant. Thus minimizing the SGD jittering risk implicitly promotes smoother and wider land- scapes of the lower-level objective F with respect to the iterative inputs z for larger values of K. It is important to note that this differs from the concept of flat minima in a general feedforward neural network, which flatness refers to regions of low training loss with respect to the network parameters, often associated with improved gener- alization. Our findings align with the observations in (Lim et al., 2021), which demonstrate that layer-wise noise in- jections to recurrent neural networks promotes a smaller Hessian with respect to the hidden inputs to the final layer. While it is widely observed that flat minima in the loss landscape (with respect to network parameters) imply better
Definition 7.3. Let Omeee gat and (ayd denote the optimal neural network parameters of f obtained from regular MSE training in (5), AT in (6) and SGD jittering training in (8), respectively.
7.1. Main Theoretical Results
Given the assumptions, we state the first theorem regarding generalization accuracy. Let Ha denote the inverse process learned using MBAs for the rest of the analysis.
Theorem 7.4. Assuming a small zero-meam Gaussian ran- dom vector g with TE] 1g 112 = 2 is added to data x E DD during evaluation. The generalization risk for denoising problem is defined as,
(0)) = FT R...yDD.al | g Holy + +9
5

---

Stochastic (Proximal) Gradient Descent Jittering
didiisrribution
attack
biex
120-
MSE Training
Input
Adversarial Training
-0
SGD littering
Avg case Robustness
Worst-case Rebustness
Coneralization
Figure 2D point denoising example using MSE training, AT, and SGD jittering. Each method is evaluated using regular in diistributiinn data, under average- and worst-case attacks to and with additional bias (from left to right columns respectively). Blue dots represents each test data instance, and the red dots indicate the ground-truth The dashed circle highlights cluster under attack, the arrows represent the distance from the center of the data cluster to the biased ground truth.
generalization summarized in Section 2, Lim et al. (2021) does not clarify why flat minima with respect to hidden- layer inputs improve generalization in regression settings. Our work fills this gap in the context of IPs.
Toy Problem: 2-dimensional Point Denoising We begin with 2D point denoising problem to validate our theorems. The task involves recovering the point noisy observations generated by adding zero-meat Gaus- sian noise with a variance of 0.01. We use 200 samples for training and 50 for evaluation. Figure I visualizes the denoising results under ID data, average-case attack, worst- case attack, and biased ground-truth ID data are sampled from the same distribution as the training data, centered at (0, O), The worst-case attack finds vector within 0.01 dg-distance around y that maximizes the squared error, while the average-case attack uniformly samples attack vec- tors within the same ball and adds to y. Finally, we evaluate generalization accuracy by reconstructing the slightly biased ground truth (0.005, 0) from its corresponding measure- ments. The results show that both SGD jittering and AT produce tighter clusters of reconstructed points compared to MSE training under adversarial attacks. Moreover, when reconstructing a slightly biased ground truth, SGD jittering achieves the most accurate reconstruction, showcasing its better generalization capability to small variations in data. Notice that MSE training underperforms SGD jittering on ID data, likely because the noise injection in SGD jittering helps the model escape from local minima in training.
(0,0) from
Furthermore, we present the robustness analysis results.
Theorem 7.5. Let e be the perturbation vector sampled i.i.d. from the zero-mean Gaussian distribution, Pe. such that Bllell2 o². For an inverse mapping He, the average- case robustness riskfor denoising is as follows,
(11) SGD jittering is more robust than MSE training against counded-varianee perturbations around the measurements,
Re (0) = E( -. [|[ - Ha(y + e]]
R.(009) S &R
The proof is similar to the one in proving generalization, we include the full proof in the Appendix D. The implicit regularization induced by SGD jittering help characterize both generalization and robustness of MBAs.
8. Experiments and Discussions
In our experiments, we use a toy denoising example to val- idate our theorems and demonstrate the proposed training scheme using seismic deconvolution and magnetic reso- nance imaging reconstruction. For each task, train the same LU architecture using different training schemes. We then evaluate the models using in-distribution (ID) testing data, adversarial attacks with various strengths, and task- dependent OOD data. Detailed training procedures are pro- vided the Appendix.
Seismic Deconvolution Seismic deconvolution a cru- cial problem in geophysics. It aims reverse the convo- lution effects in the received signal, which occurs when artificial source waves travel through the Earth's 's layers. The goal is to extract sparse reflectivity series from recorded seismic data. The training data is generated following the same procedure as in (Iqbal et al., 2019; Guan et al., 2024).
6

---

Stochastic (Proximal) Gradient Descent Jittering
MSE Training
jittering
Ground-truth
MSE Training
AT
SGD jittering
Ground-truth
Largo magnitude
Ex.
Small magnitude
(a) Adversarial Attack
(b) COD with additional Layer
Figure 2. Visualization of seismic deconvolution using (stochastic) gradient descent variants within loop unrolling architectures. Figure (a) illustrates the reconstruction performance under adversarial attack with 1. Figure (b) evaluates the model's performance on OOD seismic data, where the ground-trunt reflectivity includes an additional small horizontal layer (with magnitude of 0.3 and 0.1), highlighted in the red boxes.
PSNR/ / SSIM
ID Test Data
Adv. Attack
OOD Data
PSNR / SSIM fastMRI
Adv. Attack
Tumor Cell
MSE training
34.64/ 0.921
28.85 / 00.829
34.57 0.918
MSE training
28.21 0.603
25.68 0.382
29.92 0.779
AT
30.27 0.849
33.45 / 0.902
AT
27.68 / 0.564
27.17 0.549
27.74/ 0.597
Input Jittering
32.92 / 0.882
30.10/ / 0.882
32.91 / 0.882
Input Jittering
28.18 / 0.595
25.05
0.420
29.97 0.740
SGD jittering
35.10/ 0.928
29.89 / 0.842
34.93 / 0.927
SGD jittering
28.22 / 0.607
26.77 / 0.552
30.36/ / 0.788
Table 1. Scismic deconvolution evaluation for in-distribution data (column 2), under adversarial attack with = (column 3), and for OOD data with additional horizontal layer (columns 4). The best and second-boe -bstt performances are in bold and underlined respectively.
Table 2. MRI evaluation for in-distribution data (column 2), ad- versarial attack (column 3), and OOD data (column 4). The best and second best performances are in bold and underlined respec- tively.
used for reconstruction. To assess generalization, we use different knee dataset from Bickle & Jin (2021), which includes giant tumor cells absent from the training data. The results in Table 2 show that, on ID data, MSE training and SGD jittering perform similarly, while AT underperforms due to its inherent tradeoff in resolution. On OOD tumor data, SGD jittering outperforms all other methods, demon- strating better generalization to unseen features. Visual comparisons in Figure 3 further support these findings. Re- constructions produced by AT are smoother but often miss fine details, consistent with prior observations in black- box IP solvers (Krainovic et al., 2023). In contrast, SGD jittering produces high-resolution reconstructions, preserving critical anatomical structures, partieularly in OOD cases.
To test the generalization accuracy, we generate OOD data by introducing an arbitrary horizontal layer with magni- tudes of 0.1 to the ground truth, reflected accordingly in the measurements. The numerical results Table show that SGD jittering achieves the highest peak signal-to-noise ratio (PSNR) and structural similarity index (SSIM) on ID and OOD data, while maintaining competitive performance under adversarial attacks. In contrast, AT demonstrates the best robustness under adversarial attacks but suffers signif- icant performance degradation for both ID and OOD data. Visual reconstructions in Figure further highlight the dif- ferences between these methods. While MSE training often introduces artifacts, especially under adversarial attacks, AT fails to recover the small horizontal layer in OOD data with magnitude of 1, likely perceiving it as noise. In compar- ison, SGD jittering accurately recovers the additional layer, demonstrating its ability to generalize to subtle features not present in the training data.
8.1. Robustness vs. Generalization
Figure 4 presents comparison of various training schemes in terms of in-distribution accuracy, robustness and general- ization measured by average test PSNR on seismic decon- volution (top row) and MRI reconstruction (bottom row) tasks. The left panels show that the proposed SGD jitter- ing approach achieves better generalization to OOD data. Mcanwhile, the right panels depict the trade-off between robustness and accuracy for baseline methods, where our
Accelerated MRI Reconstruction Accelerated MRI aims to recover human- inttannitterpreted body structures from partial k-space measurements. We train models with single-coil knee MRI from the fastMRI dataset (Knoll et al., 2020) with 4 acceleration, or 1/4 of the measurements in k-space is

---

Stochastic (Proximal) Gradient Descent Jittering
MSE Training PENR-300767
jittering
Ground-truth
MSE Training
S@D jittering
Ground-truth
PSMR-30.899
PSNRx31.25
PSNR- 29, 20.188
PSN 25,981
PSNR- 29.431
PSNR-30.0o sb1
PaNR-28 1bb
PSNR-30.
(a) adv. attack
(b) OOD tumor cell
Figure 3. Gradient descent LU for MRI evaluation (a) under adversarial attack obtained from projected gradient descent with more than = 1. Each row represents an MRI reconstruction of sample from fastMRI (Knoll et al., 2020). (b) MRI evaluation using OOD tumor knee data from (Bickle & Jin, 2021).
Seismic deconvolution
8.3. Stochastic Proximal Gradient Descent Jittering
35.0
We further extend the proposed stochastic training idea to proximal gradient descent (PGD) variant of model-basee architecture training, and denote the algorithm SPGD jittering, where the learning objective is,
34.5
Robustress PSNR
30.0
OOD PSNR
34.0
33.5
JSPGD (8) (m.y) =112
where,
29.0
Φ lim prozs (wE (we (1AA' (Aæx 3) wx))
34.0
34.5
33.5
34.0
34.5
35.0
ID PSNR
ID PSNR
(12)
MRI reconstruction
Here, prox represents a proximal operator, which is re- placed by a neural network parameterized by 6, and wy denotes the stochastic noises injected independently the input of the proximal network at each iteration. Similar SGD) jittering, SPGD) jittering ensures that the noise we is re-sampled at every iteration, maintaining unbiased gradient updates while promoting variations in training data.
27.0
Robustness PSNR
OOD PSNR
26.5
26.0
28.5
28.0
27.8
28.0
28.2
27.8
26.0
28.2
Figure compares the same proximal gradient descent framework with different training schemes for accelerated MRI reconstruction. MSE training denotes the classical proximal GD variant of LU trained with MSE loss, AT, and SPGD jittering for accelerated MRI reconstruction. MSE training demonstrates reasonable generalization to OOD tumor data but remains highly sensitive to adversarial at- tacks. contrast, AT exhibits robustness against attacks but produces reconstructions with lower resolution, which can obscure fine details in critical medical imaging tasks. SPGD jittering effectively enhances both robustness and accuracy, achieving cleaner reconstructions under adversarial attacks and superior generalization to OOD tumor data.
Training
Training
Input jittering
jittering
Figure Comparing robustness generalization of different training schemes for GD L.U for seismic deconvolution problem (top row), and MRI reconstruction (hottom row).
approach provides a more effective balance between the two metries.
8.2. How to Choose SGD Noise Level?
Figure examines the impact of SGD jittering noise levels o²x in the context the denoising problem. The jittering noise level determines the robustness and in- and out-of- distribution accuracies. This parameter search allows the identification of an optimal o that balances robustness against adversarial attacks while maintaining high accuracy.
Numerical evaluations on SPGD jittering for accelerated MRI and further experiments on seismic deconvolution are presented in Appendix E. We also compare the training speed of AT and the proposed stochastic noise injection
8

---

Stochastic (Proximal) Gradient Descent Jittering
Training
SPGD jittering
Ground-truth
MSE Training
jittering
Ground-truth
PSNR-28.27
PSNR=29.352
PSNR-31.1377
PSNR=28,
PSNN-288 445
PSNIR=26.7 751
PSNR=18.666
PSNR-20.566
PSNA-31..0
PSNR-26.60
PENR=27.18
PSNR=29.70 731
(a) adv. attack
(b) OOD tumor cell
Figure Proximal gradient descent LU MRI evaluation (a) under adversarial attack obtained from projected gradient descent with norm more than 1. Each row represents MRI reconstruction of sample from fastMRI (Knoll et al., 2020). (b) MRI evaluation using OOD tumor knee data from (Bickle & Jin, 2021).
Robusemess
SGD jittering training scheme to address the tradeoff be- tween robustness and accuracy solving IPs. Our method proves effective for high-quality model-based solvers. Through mathematical analysis, we demonstrate that this approach implicitly regularizes the gradient and Hessian of the neural network with respect to the input, resulting in improved generalization and enhanced robustness compared to traditional MSE training. Experimental results show that while MSE training is vulnerable to adversarial attacks and AT recovers smooth estimates, SGD and SPGD jittering consistently produces robust, high-quality outcomes.
0.0000
0.0004
00007
0.00
0.01
0.02
0.03
0.04
0.03
Accurecy
evel
000 eval
6.0-
5.5
Software and Data
8.00
6.01
+ 0.04
0.03
0.03
The code for the proposed SGD and SPGD jittering methods are provided here:
com/
LwVProbs/SD-jjttter:
Figure 6. The effect of SGD jittering noise level 2 model performance for denoising problem. (Top): Worst-case robustness R. recorded each noise level. Each line represents different value of e, with the red dot indicating the ows that achieves the best robustness
Acknowledgements
This work was supported by the Center for Energy and Geo Processing (CeGP) at Georgia Institute of Technology.
corresponding e. (Bottom): In-distribution (ID) and out-of-distribation (OOD) accuracies in MSE are recorded for each owko with best noise levels highlighted.
Impact Statement
This work highlights the critical need for both robustness and generalization accuracy in solving inverse problems. Model-based architectures are high- perffoerrormacce and data- efficient machine learning architecture for IPs. The pro- posed SGD/SPGD jittering strategy mitigates the robustness- accuracy tradeoff for this specific architecture, demonstrat- ing both efficiency and effectiveness across diverse applica- tions. However, for inverse problems with unknown forward models where model-based architectures are unsuitable, fu- ture research may require alternative techniques to overcome the robustness-accuracy tradeoff.
techniques in Appendix H, demonstrating that the proposed method is time-efficient strategy promote robustness.
9. Conclusion
Robustness and generalization accuracy both crucial in solving IPs, yet current analyses often focus on simple black-box solvers that enhance robustness at the cost of accuracy. While empirical studies have examined the ro- bustness of MBAs, there a lack of theoretical analysis. In this work, we introduce a novel and easily implemented

---

Stochastic (Proximal) Gradient Descent Jittering
References
Gilmer, J., Ford, N., Carlini, N., and Cubuk, E. Adversarial examples are natural consequence of test error in noise. In International Conference on Machine Learning, pp. PMLR, 2019.
Athalye, A., Carlini, N., and Wagner, D. Obfuscated gra- dients give false sense of security: Circumventing de- fenses to adversarial examples. In International confer- ence on machine learning, pp. 274-283. PMLR, 2018.
Gilton, D., Ongie, G., and Willett, R. Deep equilibrium architectures for inverse problems in imaging. IEEE Transactions on Computational Imaging, 7:1 0.:123-111 1133, 2021.
Bian, W. A brief overview of optimization- based al algorithms for mri reconstruction using deep learning. arXiv preprint arXiv:2406.02626, 2024.
Bickle, I. and Jin, T. Giant cell tumor knee. case study, radiopaedia.org, 2021. URL doi org/10.533477 rTD- 93992.
Gottschlingg N. M., Antun, V., Hansen, A. C., and Ad- cock, B. The troublesome kernel: On hallucinations, no free lunches, and the accuracy-stability tradeoff in in- problems. SIAM Review, 67(1):73-104, 2025. doi: 10.1 10.11377233156687399 URL https:: //to:///oi///// org/10.
Calivá, F. Cheng, K., Shah, R., and Pedoia, V. Adver- sarial robust training of deep learning mri reconstruc- tion models. Machine Learning for Biomedical Imag-
ing, 1:1-32, 2021. ISSN 2766-905X. doi: https:
Guan, P. Iqbal, N.. Davenport, M. A., and Masood, M. Solving inverse problems with model mismatch using untrained neural networks within model-based architectures. Transactions on Machine Learning Re- search, 2024. ISSN 2835-8856. URL https:/ openreview net www.fforrmmii/rr fo
URL https:
//melba- www.eejoonnnn..ggo..
Cao, K., Wei, C., Gaidon, A., Arechiga, N., and Ma, T. Learning imbalanced datasets with ditri-distriuutionn aware margin loss. Advances in neural information pro- cessing systems, 32, 2019.
Gunasekar, S., Lee, J. D., Soudry, D., and Srebro, N. Im- plicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018.
Cohen, Rosenfeld, E., and Kolter, Z. Certified adversarial robustness via randomized smoothing. In international conference machine learning, pp. 1310-1320. PMLR, 2019.
Güngöö, A., Dar, S. U., Öztürk, $., Korkmaz, Y., Bedel, H. A., Elmas, G., Ozbey, M., and Çukur, T. Adaptive dif- fusion priors for accelerated mri reconstruction. Medical image analysis, 88:102872, 2023.
Fabian, Z., Tinaz, B., and Soltanolkotabi, M. Humus-net: Hybrid unrolled multi-scale network architecture for celerated mri reconstruction. Advances in Neural Infor- mation Processing Systems, 35:25306-25319, 2022.
Hardt, M., Recht, B., and Singer. Y. Train faster, generalize better: Stability of stochastic gradient descent. In Balcan, M. F. and Weinberger, K. Q. (eds.), Proceedings of The 33rd International Conference on Machine Learning, vol- ume 48 of Proceedings of Machine Learning Research, pp. 1225-1234 New York, New York, USA, 20-22 Jun 2016. PMLR. URL ht attpps: : http////pceeinggeeeeeee mmlr.
Fawzi, A., Moosavi-Dezfooli, S.-M., and Frossard, P. Ro- bustness of classifiers: from adversarial to random noise. Advances in neural information processing systems, 29, 2016.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021. URL https: openreview. net. forum? d= 6m mpos rM.
press/ 8 ha rdt ht.l.
Hendrycks, D. and Dietterich, T. Benchmarking neural network robustness common corruptions and perturba- tions. In International Conference Learning Represen- tations, 2019. URL https: wpwpr/vieneeeeeeeeee/ee/ee/oe..e nnt/ forum? 1==
Garrigos, G. and Gower, R. M. Handbook of convergence theorems for (stochastic) gradient methods, 2023.
Genzel, M., Macdonald, J., and März, M. Solving in- verse problems with deep neural networks-robustner included? IEEE transactions on pattern analysis and machine intelligence, 45(1):1119-111 2022.
Hochreiter, S. and Schmidhuber, J. Flat minima. Neural computation, 9(1):1-42, 1997.
Hodgkinson, L., van der Heide, C., Roosta, F,, and Mahoney, M. W. Stochastic continuous normalizing flows: training sdes as odes. In Uncertainty in Artificial Intelligence, pp. 1 1130-1140, PMLR, 2021.
Ghadimi, S. and Lan, G. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM journal on optimization, 23(4):2341-2366 2013.
10

---

Stochastic (Proximal) Gradient Descent Jittering
Iqbal, N., Liu, E., Mcclellan, H., and Al-Shuhail, A. A. Sparse multichannel blind deconvolution of seismic data via spectral projected-gradien IEEE Access, 7:23740- 23751, 2019. doi: 10. WWW.ACCCSSSSS.2..999999
Lecuyer, M., Atlidakis, V., Geambasu, R., Hsu, D., and Jana, S. Certified robustness adversarial examples with differential privacy, In 2019 IEEE symposium security and privacy (SP), pp. 656-672 IEEE, 2019.
Ji, Z. and Telgarsky, M. The implicit bias of gradient descent on nonseparable data. In Beygelzimer, A. and Hsu, D. (eds.). Proceedings of the Thirty-Second Conference Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp- 1772-1798. PMLR, 25-28 Jun mlr press//
Lei, Y., Hu, T., Li, G., and Tang, K. Stochastic gradient descent for nonconvex learning without bounded gradient assumptions. IEEE transactions neural networks and learning systems, 31 (1100))434444400 2019.
2019. URL https: v
Li, B., Chen, C.,, Wang, W., and Carin, Second order adversarial attack and certifiable robustness, 2019. In URL https://openrevien. net/forum, 2018.
Jiang, X., Toloubidokhti, M., Bergquist, J., Zenger, B., Good, W. W., MacLeod, R. S., and Wang, L. Improv- and
Li, H., Qian, J., Tian, Y., Rakhlin, A., and Jadbabaie, A. Convex and non-convex optimization under generalized smoothness. Advances in Neural Information Processing Systems, 36, 2024.
ing generalization by learning geometry-dependee physics-based based reconstruction of image sequences. IEEE fransactions on medical imaging, 42(2):403-415, 2022.
Jangg* R Y., Neyshabur*, B., Mobahi, H., Krishnan, D., and Bengio, S. Fantastic generalization measures and where to find them. In International Conference Learning Representations, 2020. URL https: //wp/epeeeee.e.eee////e net/ net/forum?? forum? (ttppmmmiidS000PPBBFVPV
Lim, S. H., Erichson, N. B., Hodgkinson, I.., and Ma- honey, M. W. Noisy recurrent neural networks. In Beygelzimer, A., Dauphin, Y.. Liang, P., and Vaughan, J. W. (eds.), Advances Neural Information Processing Systems, 2021. URL https: / openreview. .nt// forum? indiiddfffii 9xiRCE
Jim, K.-C., Giles, C. L., and Horne, B. G. analysis of noise in recurrent neural networks: convergence and generalization. IEEE Transactions on neural networks, 7 (6):1424-144 1438, 1996.
Liu, T., Li, Y., Wei, S., Zhou, E., and Zhao, T. Noisy gradient descent converges to flat minima for nonconvex matrix factorization. In International Conference on Artificial Intelligence and Statistics, pp. 1891-1899. PMLR, 2021.
Kannan, H., Kurakin, A., and Goodfellow, Adversarial logit pairing. arXiv preprint arXiv: 1803.06333. 2018.
Liu, X., Xiao, T., Si, S., Cao, Q., Kumar, S., and Hsich, C.-J. How does noise help robustness? explanation and exploration under the neural sde framework. In Proceed- ings of the IEEE/CVF conference computer vision and pallern recognition, pp. 282-290, 2020.
Kawaguchi, K., Kaelbling, I.. P., and Bengio, Y. Gener- alization in deep learning. In Mathematical Aspects of Deep Learning. Cambridge University Press, 2022. doi: 10.1017/978100000550666
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M.. and Tang. P. T. P. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017. URL httph:t////// http:::://pppereviieee www.rreeieenneeeeeeeeeee net/ forum? id=HloyRlygg.
Madden, L., Dall' Anese, E., and Becker, S. High probabil- ity convergence bounds for non-convex stochastic gradi- ent descent with sub-weibull noise. Journal of Machine Learning Research, 25(241):1-36, 2024.
Monga, V., Li, Y., and Eldar, Y. C. Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing. IEEE Signal Processing Magazine, 38(2): 18 44, 2021.
Knoll, F., Zbontar, J., Sriram, A., Muckley, M. J., Bruno, M.. Defazio, A., Parente, M., Geras, K. J., Katsnelson, J.. Chandarana, H., et al. fastmri: A publicly available raw k-space and dicom dataset of knee images for accel- erated mr image reconstruction using machine learning. Radiology: Artificial Intelligence, 2(1):e 2020.
Neyshabur, B., Bhojanapalli, S., McAllester, D., and Srebro, N. Exploring generalization deep learning. Advances in neural information processing systems, 30, 2017.
Krainovic, A., Soltanolkotabi, M., and Heckel, R. Learn- ing provably robust estimators for inverse problems via jittering. In Thirty-seventh Conference on Neural In- formation Processing Systems, 2023. URL ht https: //openreview nwt. forum? wi=p:iddgggglgidddde
Nguyen, T. H., Simsekli. U., Gurbuzbalaban, M., and
Richard, G. First exit time analysis of stochastic gradient descent under heavy-tailed gradient noise. Advances in neural information processing systems, 32, 2019.

---

Stochastic (Proximal) Gradient Descent Jittering
Orvieto, A., Kersting, H., Proske, F., Bach, F., and Lucchi, A. Anticorrelated noise injection for improved generaliza- tion. In International Conference on Machine Learning, pp. 77094-17116, PMLR, 2022.
Wang, Y., Ma. X., Bailey, J., Yi, Zhou, B., and Gu, Q On the convergence and robustness of adversarial training. In Chaudhuri, K. and Salakhutdinov, R. (eds.), Proceed- ings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learn- ing Research, pp. 6586-6595, PMLR, 09-15 Jun 2019. URL https: /ptpooeeeeeingss.ceeeeeep.e...se mwl..press/vel wang19i.. html.
Orvieto, A., Raj, A., Kersting, H., and Bach, F. Explicit regu- larization in overparametrized models via noise injection. In International Conference on Artificial Intelligence and Statistics, pp. 7265-7287. PMLR, 2023.
Whang. J., Delbracio, M., Talebi, H., Saharia, C., Dimakis, A. G., and Milanfar, P. Deblurring via stochastic refine- ment. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16293- 16303, 2022.
Pang. T., Lin, M., Yang, X., Zhu, J., and Yan, S. Robustness and accuracy could be reconcilable by (proper) definition. In International Conference on Machine Learning, pp- 17258-17277. PMLR, 2022.
Peng, W., Yao, W., Zhou, W. Zhang. X., and Yao, W. Robust regression with highly corrupted data via physics informed neural networks, October 2022.
Wu, D., Xia, S.-T., and Wang, Y. Adversarial weight pertur- bation helps robust generalization. Advances in neural information processing systems, 33:2958-2969, 2020.
arXiv:2210.10646.
Renaud, M., Prost, J., Leclaire, A., and Papadakis, N. Plug- and-play image restoration with stochastic denoising reg- ularization. In Forty-first International Conference on Machine Learning, 2024.
Xie, Z., Sato, I., and Sugiyama, M. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Confer- ence on Learning Representations, 2021. URL https: / / www.opeeeee.eee..e...........eee
www.eeeweww...... net/ forum?
Rice, L., Bair, A., Zhang, H., and Kolter, J. Z. Robustness between the worst and average case. Advances in Neural Information Processing Systems, 34:27840- -2851, 2021.
Yang. Y.-Y., Rashtchian, C., Zhang, H., Salakhutdinov, R. R., and Chaudhuri, K. A closer look at accuracy vs. robustness. Advances in neural information processing systems, 33:8588-8600 2020.
Rusak, E., Schott, L., Zimmermann, R. S., Bitterwolf, J., Bringmann, O., Bethge, M., and Brendel, W. A simple way to make neural networks robust against diverse im- age corruptions. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part III 16, pp. 53-69. Springer, 2020.
Zhang, H., Cisse, M., Dauphin, Y. N., and Lopez-Paz. D. Mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018. URL httpt:://peeneee////// /ttp///pppeeeieeeeeeeeeeeee net forum? id=r1Dép1-Rb
Salman, H., Li, J., Razenshteyn, I., Zhang, P., Zhang, H., Bubeck, S., and Yang, G. Provably robust deep learning via adversarially trained smoothed classifiers. Advances in neural information processing systems, 32, 2019.
Zhang. H., Yu, Y., Jiao, J., Xing, E., El Ghaoui, L., and Jordan, M. Theoretically principled trade-off between robustness and accuracy. In International conference on machine learning, pp. 7472-7482. PMLR, 2019.
Shafahi, A., Najibi, M., Ghiasi, M. A., Xu, Z., Dickerson, J., Studer, C., Davis, I.. S., Taylor, G., and Goldstein, T. Adversarial training for free! Advances in neural information processing systems, 32, 2019.
Zhang, T. T., Lee, B. D., Hassani, H., and Matni, N. Ad- versarial tradeoffs in robust state estimation. In 2023 American Control Conference (ACC), pp. 4083-4089, IEEE, 2023.
Smith, L., Dherin, B., Barrett, D. G. T. and De, S. On the origin of implicit regularization in stochastic gradient de- scent. In ICLR, 2021. URL https: /tttt::://ppeeeeiiieee net/forum? httridrrrrrrreeeeeeeeee
Zhao, S., Wang, X., and Wei, X. Mitigating Accuracy- Robustness Trade-Off via Balanced Multi-Teacher Adversarial Distillation on Pattern Analysis & Machine Intelligence, 46 (12):9338-9351 December 2024. ISSN 1939-
IEEE Transactions
Stutz, D., Hein, M., and Schiele, B. Disentangling adversar- ial robustness and generalization, 2019.
3539. doi: 10. 10.1009/TAAM.22044445555
URL
Tsipras, D., Santurkar, S., Engstrom, L., Turner, A., and Madry, A. Robustness may be at odds with accuracy. In International Conference on Learning Representations, 2019. URL ptttps://openreeiieee wtw.://eeet/rrrrteeeeee dd=SyxAb30cY7.
https::/doi.:ft http:///bb...eeccccutttsreeettte orrg
10. N0.1109/TPAMT. 2024 341 43080
Zheng, Y., Zhang, R., and Mao, Y. Regularizing neural net- works via adversarial model perturbation. In Proceedings
12

---

Stochastic (Proximal) Gradient Descent Jittering
of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8156-8165, 2021.
Zhou, M., Liu, T., Li, Y. Lin, D., Zhou, E., and Zhao, T. Toward understanding the importance noise in training neural networks. In International Conference on Machine Learning, pp. 7594-7602. PMLR, 2019.
Zhu, J.-Y., Park, T., Isola, P., and Efros, A. Unpaired image-to-image translation using cycle-consistent adver- sarial networks. In Proceedings the IEEE international conference on compuler vision, pp. 2223-2232, 2017.
Zhu, Z., Wu, J., Yu, B., Wu, L., and Ma, J. The anisotropic noise in stochastic gradient descent: Its behavior es- caping from sharp minima and regularization effects. In Chaudhuri, K. and Salakhutdinov, (eds.), Proceed- ings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learn- ing Research, pp. 7654-7663. PMLR, 09-15 Jun 2019. URL https::pp....m....tt. httpt:://pppceeediienpeeo http///poeedggg....ce mlr wtt..//ggg.....r.e//////11 press/v97/ zhul html.
Zou, D., Wu, J., Braverman, V., Gu, Q., and Kakade, S. Be- nign overfitting of constant-stepsize sgd for linear regres- sion. In Conference on Learning Theory, pp. 4633 4635.
PMLR. 2021.
13

---

Stochastic (Proximal) Gradient Descent Jittering
A. Notations
clarify the notations used in the paper.
Hg: learned inverse mapping parameterized by 8, i.e., black-box neural networks, MBAs. For MBAs in particular,
Hp = arg min, ly Aa (a)).
r: regularization function from predefined functions, i.e., & and &g norm.
TØ: regularization function whose gradient is learned from a neural network fa.
fe: learned gradient update of fo(x) (o)
F(x) = b1 Axl12 rr)) denote the value of the lower-level objective function.
B. Proof of SGD Jittering Convergence (Corollary 6.2)
Corollary B.1. (Restate Corollary 6.2. Convergence of MBAs-SGD for IPs) Let F(x)) = - Axl + ro(a) denote the lower-level objective function. function, Assume fo = Vrg is L-Lipschitz continuous, satisfying
[Ifo + fo(x))l2 LI| Vx, A,
let μ max;==,.... Xo where d are the eigenvalues of AT A, and define Lmaz = max L}. Consider sequence
fax 2 1k=0 generated by SGD in (8) with a constant step-size 17 Elimau K- For 2 1, the the following holds:
2LLmoz (2(F(( - inf F) + A3)),
min B OSECK
( (22
K
where, x² =arg min F, and
inf (21ee~ Axl inf ra()..
Proof. The corollary follows directly from Theorem 5.12 in (Garrigos & Gower, 2023), which assumes the objective function is Lipschitz. In the case of the SGD jittering training scheme, the lower-level objective function is given by F(x)= = 1 Axl 2 + (xx). We show that F is (L + p1) -)pisscchi,, allowing the result from Theorem 5.12 to be directly applied.
Assume So is T.-Lipschitz, expanding the gradient of F, we have,
| F(x + 4) - [=F(==]
(A(æ (A + A) y) + fa + A) - A (Ax - y) - -f(2)]l
||A AA fo(w + A) fa(x)]1
||A
All | + ||ʃa( A) - fa(2)]12
max i
A A [] A) - fo(x)]122
where A be the eigenvalues of AT A, and let max: As Thus is (L + p)-Lipshitz and we conclude the proof.
C. Proof of Generalization (Theorem 7.4)
Theorem C.1. (Restate Theorem 7.4) Assuming a small zero-meam Gaussian random vector g with | o is added to data x during evaluation. The generalization risk denwising problem defined as,
((0) = Em.yyDDa a| + g Ha(y + g)ll2
(13)
Under ussumptions 6.1, 7.1 and 7.2, SGD jittering generalizes better than regular MSE training.
g( S G(
14

---

Stochastic (Proximal) Gradient Descent Jittering
Proof. For any trained model 0, have generalization risk in (13). We can express Ho + g) in terms of Ho(y) for the same set of parameters via iterative expansions. For iterations k: = 0, 1, ...; K, let the sequence (00, 21,..., FK denote the denoising trajectory from observation y + g. where Hø(y + 9) æx. and let { x1, ....... T'y ) denote the denoising trajectory from observation where Hy(y) = z'x- Then we can rewrite the reconstruction trajectory with the appearance of g as follows,
2a 2 g = y + g
Φ = (1 - 1) + nyy + g) - nfo(2zo)
((1 r) + my ()o(())) + (f -
ω g (()( fo(
22 (1 1))=1 nly 9) nfo((1)
=/m g 11 (fa(b)) fa(2o)) + (Tf((()))))) fo(x)))
#++1 = 2 + g Σ (1 1)) (Jo(a²) So(wi)
Let K = k + 1 and plug in the above expansion to (9). For brevity, we omit the subscripts of expectations in the proof.
K-1
((0) F|| xxx
K-1
K 711 - 17) - (fo(m)) - (a.))) [ rk. Σ 111 1)
= F|| 11
B
Σ t=
()>---((fo(at) (fo(=)) - fo(a:)))
The first term is the testing accuracy when y are in the data distribution, while the second and third terms are affected by g. MSE training minimizes the first term solely, but no extra regularization on the other terms. On the other hand, we will show that SGD jittering training also penalizes the magnitude the other terms, thus obtaining a smaller generalization risk than MSE training.
We view the SGD jittering process as a noisy version of regular MSE training, the goal is to write the SGD jittering risk in terms of the noiseless updates with some implicit regularization. Let (#6, #1 ...,, 2xx denote the noiseless GD trajectory of model O. We write the iterative updates of SGD jittering training {x x "..,,KK xxd) for the same set of parameters 0 in terms of s's,, which tells the deviation of each intermediate reconstruction between the SGD jittering and its noiseless counterpart,
ao we + wi
(1 n) agd + 1/ (x 1/41
æ + v(fo((t))) fo(=x)")) - (11(1 w)) + syw1)
(1 naardd + "1 (a 1/222
æ + v((o(())) fa( :n(1= (So((s)) fo( (11(1 - -1) + 11(1 m) + syw2))
sgd m k--
(1 - -7)a
+ my - mffo 2nd) mpwx+1
k+1
Σ " )o(
Σ n11 1)) -
0=0
1=0
Let K = k + 1. Since the jittering noise is zero-mean, the SGD jittering training risk becomes,
K
K-1
B|ix
sod1 1
x'k
11
+ E.
Σ
(11
n)) K
m. +
TR.
Σ
(11 x( - fo(a
(=0
1=0
K K-1
2F
(x
K
(11 7) (fo-)a( (a
The first component corresponds to the MSE loss when the trajectory is not perturbed by jittering noise. The second term is weighted noise variance independent of the network parameters. The third term adds extra regularization to the difference in
15

---

Stochastic (Proximal) Gradient Descent Jittering
outputs of Se between noisy and clean trajectories. It penalizes more the iterations approach the final output. The last term is the cross product between the reconstruction error under noiseless GD trajectory and the perturbed reconstructions from f. which is lower bounded since the overall SGD jittering risk is non-negative.
When choosing ow at iteration i such that Σ 11 - 1)) -))((=)0 in SGD jittering training where W2 are independently sampled at each iteration. Thus, minimizing the regularization term fo," )>l2 also reduces the penalty term in computing generalization accuracy Ex > - 711 (fo(aa) - fo (..))). Minimizing the last term also minimizes -21 [ fo (xx))) in the generalization accuracy. Therefore, (( G(0mse) due to the implicit regularization term in SGD jittering training.
og. the reconstruction from + + is special case
E (m. vk
- 771 (ffo
a'k. Ei=o
K-1 7(1 -)
D. Proof of Robustness (Theorem 7.5)
Theorem D.1. (Restate Theorem 7.5) Let be the perturbation vector and let Pe be iid zero-mean Gaussian distribution such that E Bl|e|l = o.. The average-case robustness risk for denoising problem is follows,
Re (0) = B( [| 1oa-Hool Holy + e)lla
(14)
For the denoising problem, SGD jittering is more robust than MSE training against bounded-variance perturbations around the measurement, or
Re (0=gd) S Re (fmse).
Proof. The proof is similar to proving generalization accuracy. To evaluate the average-case robustness for any trained we first write the iterative updates under attack { क. ....x%k } in terms of the noiseless trajectory {x/) "''....''k
x x/ + =
21 (1 n)=² n e) - n
v(fo(ah) fo(x)')
x2 et (1 - n) + nly e) (2(()
= x + e + 11 (f((()))) fo( + n(fo(x()) - fo(x("))
# t e + Σ (1 (ff(x)) fo(x:"))
K- + we rewrite the robustness risk. Notice that the generalization loss in (9) includes the additional vector whereas the robustness risk in (11) does not.
K-1
(2( B||x x'x -e
Σ km0
n11 11) (fo((() fo(a)"))
K
K-1
=E
El|
7111 1) K -1-- (fo(x))) fo(x fo( -Blee1 2E
(x
FK Σ n(1 x0 ({-o<a(( fola) ())
U-U
2ma
K
E|i - x
112
E||
Σ
1) K 10--1-((o(()) - (x(""))) o² 2E (x
T'K
(11 1) K-1 (fo(x))) - fo
The first term the mean-squared testing loss for in-distribution data. The second term is an additional penalty that appears in SGD jittering training. When picking o²_ for all & and o2 such that Σ 17 1 - 17) 2kk o2 = o² this term in robustness risk is a special case in SGD jittering risk. The last term also appears in the SGD jittering risk; therefore, minimizing the SGD jittering risk implicitly reduces this term as well. Therefore, for models that perform equally well in MSE for in-distribution testing data, minimizing SGD jittering risk improves the robustness.
16

---

Stochastic (Proximal) Gradient Descent Jittering
E. Extra Stochastic Proximal Gradient Descent Experiments
Similar to gradient descent variants loop unrolling architectures, for proximal gradient descent (PGD) variants, we train a PGD LU with regular MSE loss. Then finding the worst-case attack to y over all PGD updates.
E.1. Seismic Deconvolution
In the same seismic deconvolution problem setup discussed in the main manuscript, but now solved using proximal gradient variants, Figure 7(a) shows the performance of the trained models under an adversarial attack with c 1. Figure 7(b) focuses on reconstructing the additional layer highlighted in red boxes. Table 3 provided numerical evaluations. The proposed SPGD jittering demonstrates greater robustness compared to MSE training while better preserving detailed information than adversarial training (AT).
MSE Training
AT
SPOD jittering
Ground-truth
MSE Training
AT
SPGD jittering
Ground-truth
Large magnitude
Smell magnitude
Adversarial Attack
OOD with additional Layer
Figure Visualization seismic deconvolution using (stochastic) proximal gradient descent variants within loop unrolling architectures. Figure (a) illustrates the reconstruction performance under adversarial attack with 1. Figure (b) evaluates the model's performance on out-of-diittrbuution (OOD) seismic data, where the ground-truth reflectivity includes an additional small horizontal layer, highlighted the red boxes.
PSNR SSIM
in-distribution
Adv. Attack
OOD
MSE training
34.91/ /0.23
28.78 0.842
34.810/ 0.920
AT
32.99/ 0.899
31.44 / 0.875
32.966 / 0.899
Input Jittering
33.02 0.899
30.56 0.861
32.910 / 0.895
SPGD Jittering
35.05 / 0.931
29.93 / 0.870
34.952 / 0.929
Table PGD variants of L.U evaluated on seismic deconvolution across three scenarios: in-distribution data (column 2), adversarial attacks (column 3), and out-of iitto---iitrrbution (OOD) data (column 4). The best performance in each case highlighted in bold, while the second- bost performance is underlined.
E.2. Accelerated MRI Reconstruction
We complete the missing numerical results discussed in the main manuscript for 4x accelerated MRI reconstruction using SPGD jittering in Table 4.
PSNR SSIM
fastMRI
Adv. Attack
Tumor Cell
MSE training
27.93/ 0.582
22.82 / 0.179
29.39 0.771
AT
27.20 0.533
26.78 0.518
27.13 0.560
Input Jittering
24.88 / 0.299
24.44 / 0.297
25.85 0.478
SPGD Jittering
28.15 / 0.597
27.56 0.577
29.762 / 0.777
Table PGD variants of LU evaluated MRI data across three scenarios: in-distribution data (column 2), adversarial attacks (column 3), and out-of-distribution (OOD) data (column 4). The best performance in each case is highlighted in bold, while the second-best performance is underlined.
17

---

Stochastic (Proximal) Gradient Descent Jittering
F. Compared to Other Baselines for MRI reconstruction
We evaluate our method against a widely used robustness technique Lipschitz regularization using spectral normalization (SN). SN constrains the Lipschitz constant of the network by bounding the spectral norm (i.e., largest singular value) of each weight matrix, which helps stabilize training and improve robustness. As shown in the second last row in Table for 4xMRI reconstruction. SN promotes adv. robustness, but the cost of reduced accuracy some extend. This is likely due its restrictive nature of SN limiting the model's expressive power.
Additionally, we also compares the performance to a diffusion models (DMs), as they also demonstrates impressive results in image generation. Since our work proposes a general framework for IPs, we chose to compare against the standard DDPM rather than specialized task-specific DMs, consistent with prior work (Güngör et al., 2023). To ensure a fair comparison under similar computational constraints, we adapt the denoising U-Net fit within the same GPU memory as other methods.
Table 5 compares DDPM with other methods. DDPM achieves comparable in-distribution (ID) performance to MBAs trained with both MSE loss and the proposed SGD jittering, and shows stronger robustness than standard MSE-trained MBAs. However, it underperforms in OOD generalization. While DDPM serves as strong baseline for robustness and ID accuracy, SGD jittering achieves better generalization under distribution shifts. We will add the comparison to DM in the main mamuscript as an interesting baseline.
is also worth noting that DDPM requires 10x more parameters and is significantly more data- data-intensive, whereas MBAs are more data-efficient (Monga al.. 2021) due to their optimization-inspired iterative structure. We also refer to prior work (Güngör et al., 2023), which compares DDPM, DiffRecon, AdaDiff to MSE-trained MBAs for MRI reconstruction. Their results show that while DM can generalize well in some cases, MBA methods consistently perform better on ID data. Results in (Güngör et al., 2023) shows that MSE-trained MBA is a strong baseline, and our proposed SGD jittering further improves their robustness and generalization.
Model
Training Scheme/Design
fastMRI
Adv. Attack
Tumor Cell
MSE training
28.21 0.603
25.68 / 0.332
29.92 / 0.779
AT
27.68 / 0.564
27.17 / 005449
27.74 / 0.597
GD LU
Input Jittering
28.18/ / 0.595
25.05 / 0,422
29.97 / 0.740
SGD jittering (Ours)
28.22 0.607
26.77 / 0.552
30.36 / 0.788
Lipschitz regularization (SN)
27.71 0,576
27.09 0.542
27.92 / 0.594
DDPM
Default training
28.17 / 0.611
27.13 / 0.536
29.72 / 0.782
Table 5. MRI reconstruction evaluation compared to the Lipschitzzeegularzzatiii regularization method (last row). The best and second best performances are in bold and underlined respectively.
G. Training Details
We use 10-iteration gradient descent loop unrolling architectures for all tasks. For the toy example, we use a 3-layer MLP with hidden dimension of 32 as the learned gradient network. For seismic deconvolution and MRI reconstruction, a 5-layer and 8-layer DnCNN with 64 hidden channels are used for the learned gradient network, respectively. All models are trained using Nvidia RTX3080, using Adam optimizer with a learning rate of 1e 4.
The jittering noise variance for each method is selected based on performances of both robustness and accuracy. In input jittering training, the three methods in order of 2D denoising, seismic deconvolution and MRI reconstruction use the noise variance of 0.01, 0.05 and 0.05. In SGD jittering training, we choose the SGD jittering variance for each task with 0.01, 0.1 and 0.01 respectively.
H. Training Speed
We also compare the training speed among different training schemes for GD and PGD LU architectures. All methods are trained on the same device. Training batch sizes for the 2D denoising problem, seismic deconvolution and MRI are 256, 16, and 4 respectively. AT is trained by projected gradient descent. For the 2D denoising problem, = 0.01 with 50 projected GD iterations and a stepsize of 0.05. For the seismic deconvolution and MRI reconstruction problem, c c= 1, with 20 projected GD iterations and a stepsize of 0.1. AT is significantly slower than other methods. Notice that the jittering in SGD and SPGD does introduce very minimal time overhead compared to standard MSE training, even when accounting for
18

---

Stochastic (Proximal) Gradient Descent Jittering
the additional noise sampling steps.
2D demoising
Seis. Deconv.
MRI
MSE training
5360.84 it/s
23.69 it/s
5.23 it/s
AT
562.75 it/s
1.53 it/s
0.25 it/s
GD LU
Input Jittering
5256.77 it/s
22.80 it/s
5.20 it/s
SGD jittering
4337.84 it/s
22.72 it'
5.27 it/s
MSE training
5136.34 it/s
16.78 it/s
6.23 it/s
AT
486.68 it/s
1.12 it/s
0.17 it/s
PGD LU
Input Jittering
4965.22 2.55.25 it/s
16.44 it/s
6.17 it/s
SPGD jittering
4669.76 it/s
14.52 it's
6.07 it/s
Table 6. Training speed measured in items per second are recorded.
19


Reviews:
[{"review_id": "IvypFm61bO", "reviewer": "Reviewer_KSe9", "summary": "The paper investigates robustness-accuracy tradeoffs, where the authors focus on unrolling-based methods. The authors consider different training strategies for increasing the robustness to average-case perturbations or distribution-shifts. As a specific solution for unrolling-based methods, the authors propose to add jittering noise in each step of unrolling, and demonstrate a good robustness-accuracy tradeoff for practical setups.\n\n## update after rebuttal\nThanks again to the authors for their rebuttal. With the exception of my comment regarding the choice of noise levels, the authors addressed my concerns well. I think this is an interesting paper and keep my score as (weak) accept.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful feedback and helpful suggestions. We address the reviewer’s comments and questions below.\n\n>Jittering outperforms MSE training in in-distribution results\n\nWe acknowledge the reviewer’s observation and appreciate the opportunity to elaborate further. As noted in L.310–313 (prior to the Seismic Deconvolution section), training with standard MSE loss may lead to suboptimal solutions due to the highly non-convex loss surface with respect to network parameters. Our proposed SGD jittering acts similarly to layer-wise noise injection, which is known to help escape local minima by promoting exploration of the loss landscape [1]. Consequently, jittering can achieve better in-distribution and generalization performance than MSE training in some cases, as supported by both our empirical results and prior work [2].\n\n> L.167: number of steps used in PGD for AT\n\nAs detailed in Appendix G (L.836), we used 20 PGD steps with a step size of 0.1 for both seismic deconvolution and MRI reconstruction tasks. While the runtime efficiency is a useful byproduct, our main point is that methods like AT and input noise injection introduce perturbations directly to the input, often break forward model consistency and result in overly smoothed reconstructions. Our contribution is to introduce SGD jittering, which promotes robustness in a more principled manner while preserving high accuracy and physical model fidelity.\n\n> How jittering noise levels were chosen\n\nAs shown in Figure 5 (L.420), we compared robustness and accuracy across a range of jittering noise levels and selected the setting with the best overall trade-off. A similar hyperparameter search was conducted for the input injection baseline. Notably, we observed that input injection exhibited greater sensitivity to noise level variations, whereas our proposed SGD jittering showed more stable performance across a range of noise levels. While some tuning is still required, this suggests that SGD jittering is more robust to hyperparameter selection in practice. We will include a short discussion of this in the main paper.\n\n> Notion of Generalization\n\nWe agree that in classical learning theory, generalization typically refers to a model's ability to perform well on unseen in-distribution data, often analyzed in the finite-sample setting. However, in this work, we adopt a broader and increasingly common notion of generalization under distribution shift, which aligns with recent literature on robustness and transferability. To avoid confusion, we will revise the manuscript to clarify this broader usage of \"generalization\" and explicitly distinguish it from classical in-distribution generalization. \n\n> Typos and clarification\n\nWe thank the reviewer to pointing it out, we will fix the typo and make clarification of the terms in the manuscript.\n\n[1] Orvieto et al. Explicit Regularization in Over-parametrized Models via Noise Injection\n\n[2] Lim et al. Noisy Recurrent Neural Networks\n\nWe sincerely appreciate the reviewer’s time for reviewing and the constructive suggestions, and we believe that the additional clarifications improve the quality of the submission. We hope this addresses your concerns.", "claims_and_evidence": "- The main claim of the paper is that SGD jittering (the proposed training strategy) yields better generalization and higher average-case robustness compared to standard MSE training.  This is not surprising, since the same holds true for other robustness-enhancing methods (input jittering or adversarial training). The authors support this by proving it for denoising setup (and technical assumptions), and provide sufficiently convincing experiments.\n- Another claim of the paper is that SGD jittering improves accuracy at the same time as robustness, and thereby overcomes or mitigates the robustness-accuracy-tradeoff (see e.g. L061-066). This is very interesting, but not supported by theory, and I have some concerns and questions regarding the experimental setup.", "methods_and_evaluation_criteria": "The authors train unrolled networks with the four different training strategies (MSE, adversarial training, input jittering, SGD jittering) and evaluate on in-distribution data, adversarial attacks and out-of-distribution examples (Tables 1 and 2), which is appropriate to investigate the claims. See also experimental design and analyses.", "theoretical_claims": "I checked the proof of Theorem 7.5 (the average-case robustness result) and don't see any major issues.", "experimental_designs_or_analyses": "I find the visualization of the toy problem results interesting, but think that real-world dataset examples would be more beneficial for the reader. The MRI results are relatively convincing, the seismic deconvolution problem is more niche. For the MRI results (and also the toy problem), I have some concerns regarding the in-distribution performance of standard MSE training. Specifically, it is very surprising to me that jittering improves the in-distribution accuracy compared to MSE training (which optimizes for it) (see also questions).\nMoreover, the authors write that the jittering levels (of SGD and input methods) are chosen based on 'robustness and accuracy' (L.831) and only states the concrete values, which are hard to interpret. Since the main conclusion is that the method performs better in robustness and accuracy than competing methods, I think there should be a more principled approach to choosing it (or description of it if applicable).", "relation_to_broader_scientific_literature": "The authors sufficiently describe the related work, including existing demonstration of robustness-enhancing methods via noise injection.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "The investigated problem of addressing the robustness-accuracy tradeoff is important, and proposing methods which are less expensive than adversarial training is valuable. The idea of injecting noise for increasing robustness to the input, or the intermediate layers, has already been presented in the literature. However, the analysis and theoretical arguments of this approach with respect to average-case robustness is interesting and novel.", "other_comments_or_suggestions": "- I find the title a bit miss-leading, as \"unrolling methods\" (considered in the paper) is only a small subset of \"model-based architectures\". Moreover, the strengths of the paper are with respect to robustness, but there is noticeably less support for the accuracy part (minor comment).\n- I am a bit confused by the notion of generalization, and I think the papers actually discusses robustness to distribution-shifts. While there are varying notions in the literature, generalization results are often associated to a finite sample case and in-distribution data, but here the authors effectively investigate a distribution shift by adding noise to the training distribution.\n- L. 758,759: A term is missing in the inner product.\n- Notation of the robustness risk: The authors denote the average-case robustness risk by $R_{\\epsilon}$ and the worst-case robustness risk by $R_{e}$, which leads to confusion at first I think.\n- L: 173,174: The authors write that AT learns to 'ignore' $A^{-1} e$, but this is not true in general (e.g. only when the perturbations are very large). AT finds a tradeoff between reconstructing the signal and minimizing the error.\n- L. 355,356: \"tradeoff in resolution\"", "questions_for_authors": "- In L. 167, the authors write that AT is slow and requires iterative solvers for the attack vector. How many steps were used in PGD for adversarial training / testing?\n- Regarding L. 831, could you give more details on how the jittering noise levels were chosen?\n- In Figure 1 I would expect that the MSE training yields a better in-distribution results (particularly compared to adversarial training)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xEYtAw8jOz", "reviewer": "Reviewer_v9Xr", "summary": "The authors study the robustness and generalization properties of model-based architectures. The goal is to solve inverse problems with interpretable algorithms, such as loop-unrolling networks, and maintain two desirable properties: i) robustness to adversarial attacks, ii) generalization to small natural shifts in test-time data. The authors propose an algorithm called SGD jittering that makes progress in achieving these properties without sacrificing performance. The algorithm is based on a noisy version of Gradient Descent in the loop-unrolling architecture.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful suggestion, and please find the point-to-point response below.\n\n> Stronger baselines such as diffusion models (DM)\n\nWe agree that diffusion models (DMs) have demonstrated impressive results in image generation. In response, we added DDPM-based experiments for MRI reconstruction. To ensure a fair comparison under similar computational constraints, we adapted the denoising U-Net to fit within the same GPU memory as other methods. Since our work proposes a general framework for IPs, we chose to compare against the standard DDPM rather than specialized task-specific DMs, consistent with prior work [1].\n\nTable 1 compares DDPM with other methods. DDPM achieves comparable in-distribution (ID) performance to MBAs trained with both MSE loss and the proposed SGD jittering, and shows stronger robustness than standard MSE-trained MBAs. However, it underperforms in OOD generalization. While DDPM serves as a strong baseline for robustness and ID accuracy, SGD jittering achieves better generalization under distribution shifts. We will add the comparison to DM in the main manuscript as an interesting baseline.\n    \nIt is also worth noting that DDPM requires ~10× more parameters and is significantly more data-intensive, whereas MBAs are more data-efficient [2] due to their optimization-inspired iterative structure. We also refer to prior work [1], which compares DDPM, DiffRecon, AdaDiff to MSE-trained MBAs for MRI reconstruction. Their results show that while DM can generalize well in some cases, MBA methods consistently perform better on ID data. Results in [1] shows that **MSE-trained MBA is a strong baseline**, and our proposed SGD jittering further improves their robustness and generalization.\n\n|    PSNR/SSIM     | ID | Adv. Attack | OOD |\n|:--------------------:|:-------------:|:-------------:|:-------------:|\n|     MSE training     | <28.21> / 0.603 | 25.68 / 0.382 | 29.92 / 0.779 |\n|          AT          | 27.68 / 0.564 | **27.17** / <0.549> |  7.74 / 0.597 |\n|    Input Jittering   | 28.18 / 0.595 | 25.05 / 0.420 | <29.97> / 0.740 |\n| SGD jittering (Ours) | **28.22** / <0.607> | 26.77 / **0.552** | **30.36 / 0.788** |\n|  DDPM (new baseline) | 28.17 / **0.611** | <27.13> / 0.536 | 29.72 / <0.782> |\nTable 1: MRI reconstruction. Best performances in **bold**, second best in <...>.\n\n> Robustness and generalization tradeoff\n\nWe thank the reviewer for raising this important question regarding a fundamental challenge. The tradeoff has been studied from various perspectives (i.e, distributional and optimization) [3,4], with evidence that no single training objective can simultaneously optimize both. In our work, we contextualize this tradeoff within IPs by explicitly defining robustness and generalization accuracy in relation to the forward model. Our theoretical and empirical results show that AT and standard MSE training prioritize robustness and accuracy, respectively, but fail to address both objectives effectively. In contrast, the proposed SGD jittering implicitly regularizes the model, enabling simultaneous improvement in both metrics.\n\n> How to interpret $x_g$\n\nFor robustness (Eq.2), $g$ is interpreted as artifacts due to noise to measurement $y$. Eq.2 measures how reconstruction $H_\\theta(y_g)$ deviates from the clean $x$.\n    \nFor generalization (Eq.3), $x_g$ is considered as natural shift of $x$. Eq.3 measures how well the model reconstruct $x_g$ from its corresponding measurement $y_g$, or maintaining consistency with the physics model.\n\n[1] Dar et al. Adaptive diffusion priors for accelerated MRI reconstruction\n\n[2] Monga et al. Algorithm Unrolling: Interpretable, Efficient Deep Learning for Signal and Image Processing\n\n[3] Zhang et al. Theoretically Principled Trade-off between Robustness and Accuracy\n\n[4] Krainovic et al. Learning Provably Robust Estimators for Inverse Problems via Jittering\n\nWe sincerely appreciate the reviewer’s constructive feedback. We believe the added experiments and clarification strengthened the manuscript and hope they satisfactorily address your concerns.", "claims_and_evidence": "The claims are supported by the evidence.", "methods_and_evaluation_criteria": "The authors only include comparisons to other model-based algorithms. I think it would be better if stronger baselines were included, such as solving these problems with diffusion models. The latter lack the interpretability of model-based methods, but it would be nice to see what price we pay in performance to gain this interpretability. It would also be nice to include examples of where this interpretability is important.", "theoretical_claims": "I did not rigorously check all the proofs, but I read the theoretical results and the assumptions and they make sense.", "experimental_designs_or_analyses": "I checked the experimental results (seismic deconvolution and MRI). As I mentioned above, I believe the biggest weakness is the lack of stronger baselines.", "relation_to_broader_scientific_literature": "The paper improves the robustness and generalization of model-based algorithms that offer interpretable solutions to inverse problems.", "essential_references_not_discussed": "N/A.", "other_strengths_and_weaknesses": "Some weaknesses are mentioned above. I would further add that it is unclear what is the fundamental limit on the trade-off between generalization and robustness. How should the reader think of $x_g$ in Section 3? As an adversarial input or as a natural shift of $x$? \nIt would also be nice for the paper to include some convincing evidence on why model-based algorithms are interpretable in the context of MRI/Seismic Deconvolution and compare the performance with stronger baselines.\n\nIn terms of strengths, the paper has interesting theoretical results, it is well-presented and it has a clear motivation.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "Could you include comparisons with stronger baselines for the problems you want to address? MBAs do not need to necessarily improve upon these baselines, but it needs to be more clear what's the trade-off between interpretability, performance, robustness and generalization.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FBloO1tyfe", "reviewer": "Reviewer_mUA8", "summary": "The paper introduces \"SGD Jittering,\" a new  training strategy designed to enhance the robustness and generalization of Model-Based Architectures (MBAs) for image inverse problems. Specifically, the authors propose to inject random zero-mean Gaussian noises into gradient updates at each iteration within deep unrolling networks during training. Theoretically, they demonstrate that this simple noise injection improves average-case robustness and generalization accuracy compared to standard mean-squared-error and adversarial training, respectively. Empirically, they validate their method across several inverse problem tasks, including a toy denoising example, seismic deconvolution, and single-coil MRI reconstruction.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments and helpful questions. Please find detailed responses below.\n\n> Generality of Theoretical Results\n\nWe agree that Theorems 7.4 and 7.5 were established specifically in the denoising setting. Extending the theoretical analysis to more complex inverse problems would require additional assumptions and broader treatment of the forward model. We will explicitly acknowledge this limitation in the revised manuscript and outline it as a promising direction for future work.\n\n> Extension to Proximal-Based MBAs\n\nThe analysis of robustness and generalization for SPGD are not a direct extension of the current SGD theoretical results. \nWhile GD-LU and PGD-LU differ structurally—based on whether the consistency update and neural network modules (which learns $\\nabla r$ in GD-LU or acts as a proximal operator in PGD-LU) are applied in parallel or sequentially—both architectures allow for noise injection into hidden states during training. This perturbs intermediate representations, thus promoting robustness.\n\nWhile our theoretical analysis focuses on SGD, prior work has established convergence for both SGD and SPGD in related setting. These results ensure that the reconstructed outputs remain consistent with the forward model $A$, thus preserving reconstruction accuracy. Motivated by this, we extend our approach empirically to the proximal setting. Our experiments confirm that jittering improves both robustness and generalization in proximal settings.\n\n> Generalization Risk Definition Eq.3\n\nOur current definition aims at capturing small shifts in test data, but assuming same and known forward model is used at inference. We agree that explicitly addressing forward-model mismatches between training and inference would be an important direction. This work serves as an initial discussion in addressing generalization and robustness issue only due to data-shift.\n\n> Clarification of Regularization Functional $r$\n\nIn model-based architectures, the regularization function $r$ is implicitly defined through its gradient learned by a neural network $NN_\\theta = \\nabla r$. Thus, we do not explicitly specify the functional form of $r$. Please refer to L:117 under Eq.4 for a detailed discussion.\n\n> Clarifying Connections to RED frameworks\n\nTheir motivation and implementation are fundamentally different. MBA is a supervised learning strategy whose training is formulated as a bilevel optimization problem (Eq.5-Eq.8), where the network is trained end-to-end and allows for noise injection during training. In contrast, RED uses a pre-trained denoiser as an explicit regularization function during inference, without end-to-end training. The denoiser in RED is fixed and not learned as part of the reconstruction process, so SGD jittering is not applicable to RED.\n\n> Review of robustness, generalization for IPs\n\nThank the reviewers suggestion, and we will include the following paragraph in the manuscript:\n\n\"To improve robustness, several strategies have been proposed, including training-time exposure to diverse perturbations such as noise injection and data augmentation [1-2], and adversarial training, enabling models to better handle noise during inference. Other approaches focus on architectural innovations—for example, diffusion models have shown inherent robustness to noise in MRI reconstruction [4], while PINN [5] embed domain knowledge to enhance stability without compromising interpretability. To improve generalization in IPs, methods like data augmentation with synthetic perturbations [6] and domain adaptation via techniques such as CycleGAN [7] have been used to expand the training distribution and improve adaptability to unseen data. Additionally, incorporating geometric constraints, as in [8] for electrocardiographic image reconstruction, has shown improved generalization by embedding prior knowledge into the learning process. While these approaches typically target either robustness or generalization, achieving both simultaneously remains an open and challenging problem.\"\n\n[1] Krainovic et al. Learning Provably Robust Estimators for Inverse Problems via Jittering\n\n[2] Zhou et al. Towards Understanding the Importance of Noise in Training Neural Networks\n\n[4] Dar et al. Adaptive diffusion priors for accelerated MRI reconstruction\n\n[5] Peng et al. Robust Regression with Highly Corrupted Data via Physics Informed Neural Networks\n\n[6] Guan et al. Solving Inverse Problems with Model Mismatch using Untrained Neural Networks within Model-based Architectures\n\n[7] Zhu et al. Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks\n\n[8] Jiang et al. Improving Generalization by Learning Geometry-Dependent and Physics-Based Reconstruction of Image Sequences\n\nWe sincerely appreciate the reviewer’s constructive feedback. We believe the added literature review and clarification will strengthened the manuscript and hope they satisfactorily address your concerns.", "claims_and_evidence": "Yes. The paper provides both empirical and theoretical results on non-convex deep neural networks, systematically comparing SGD Jittering with standard MSE training and worst-case adversarial attacks. Additionally, to the best of the reviewer's knowledge, this work is the first to offer a theoretical analysis of generalization accuracy in inverse problems, particularly in the presence of small perturbations in test data.", "methods_and_evaluation_criteria": "Yes, particularly in the empirical evaluation, where the model is trained on the fastMRI dataset and tested on a different knee dataset from Bickle & Jin (2021). Note that the test set contains giant tumor cells absent from the training data, ensuring a robust out-of-distribution (OOD) evaluation.", "theoretical_claims": "The assumptions used for the theorems align well with existing literature. The training convergence analysis under the SGD optimizer is a direct extension of Garrigos & Gower (2023). Additionally, I do not see major flaws in Theorems 7.2 and 7.4.", "experimental_designs_or_analyses": "Yes, the empirical design is methodologically sound and provides informative insights.", "relation_to_broader_scientific_literature": "Prior works, Fawzi et al., 2016, showed that small perturbations in data space can significantly degrade performance, particularly for deep networks trained on ill-posed inverse problems. This paper’s approach of introducing SGD perturbations at the optimization level offers a new perspective.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths\n\n1, This paper presents an interesting and easily implementable training strategy to improve robustness and generalization in model-based deep learning.\n\n2, The paper is well-written and easy to follow, with clear mathematical formulations and consistent notations.\n\n3, The theoretical analysis effectively supports the intuition behind the proposed method, explicitly linking SGD jittering to improved robustness and generalization.\n\nWeaknesses\n\n1, The main theorems, 7.4 and 7.5, are primarily derived for image denoising, leaving their applicability to more general inverse problems (e.g., MRI or seismic reconstruction) unclear. Extending the theory may require additional assumptions on the forward model.\n\n2, The results on proximal-based MBAs appear to be an ad-hoc empirical extension, rather than a direct consequence of the theoretical analysis. It is unclear whether the current theoretical framework applies directly or if additional assumptions are needed for deep neural network priors.\n\n3, The paper lacks a more comprehensive review of generalization and robustness in inverse problems. While Definition 3.2 introduces a generalization risk formulation, it does not fully account for forward model mismatches between training and inference, which are critical in practical applications.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1, Can the authors clarify what regularization functionals are used for constructing $r_\\theta$ ?\n\n2, Can similar approach and theoretical analysis to deep priors without explicit potentials, such as those based on regularization by denoising (RED) frameworks ?\n\n3, Can the authors further clarify why the chosen definition of generalization risk (Eq. (3)) is suitable for general inverse problems? Would other definitions (such as those involving a distributional shift of the forward operator $A$ or the noise $z$) be potentially more relevant?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BpT6UK8hew", "reviewer": "Reviewer_x3Tj", "summary": "The paper introduces SGD Jittering, a training method for model-based architectures (MBAs) solving inverse problems. By adding small, random noise to gradient updates during training, SGD Jittering improves robustness and generalization accuracy without modifying input data or increasing computational cost like adversarial training (AT). Theoretical analysis proves its advantages over standard mean squared error (MSE) training. Experiments on denoising, seismic deconvolution, and MRI reconstruction show superior performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments and helpful questions. Please find detailed responses below.\n\n> Comparing to SGLD\n\nWe thank R-x3Tj for mentioning SGLD, but we clarify that our SGD jittering is fundamentally different from SGLD in both goal and mechanisms.\n    \nSGLD adds noise directly to network parameter updates to approximate sampling from the Bayesian posterior. It aims to capture model uncertainty, not necessarily to improve robustness or generalization. This process requires careful scheduling of the noise and often results in slower training.\n\nIn contrast, SGD jittering injects noise into the hidden inputs (intermediate activations) during training—not into the parameters—with the explicit goal of enhancing robustness and generalization in IPs. Importantly, no noise is used at inference time, so reconstructions remain deterministic. Moreover, our method integrates naturally with MBAs and does not require any modification to the optimizer.\n\n> Theoretical guarantees are limited to denoising tasks; provide broader analysis for general IPs\n\nWe agree that extending the theoretical guarantees to more general IPs, including nonlinear and ill-posed settings with data/model mismatch, remains a rich direction for future work. We see our current analysis as an important first step toward building a theoretical foundation for robust training schemes in MBAs, and we hope it will inspire further progress in this area.\n\n> Testing on more diverse tasks (e.g., super-resolution) could better validate generality\n\nIn addition to the main experiments, we also evaluated our method on a natural image deblurring task using the CelebA dataset. Due to space constraints, these results were not included in the main text. Below, we provide test performance on in-distribution (CelebA), out-of-distribution (FairFace), and adversarial settings. Table 2 shows improved ID/OOD accuracies and robustness over MSE training, with robustness nearly matching AT. We will include these results and discussion in the revised manuscript to further support the generality of our approach.\n\n| PSNR/SSIM            | ID-CelebA         | Adv. Attack       | OOD-FairFace      |\n|----------------------|-------------------|-------------------|-------------------|\n|     MSE training     |   <34.14 / 0.954>   |   29.81 / 0.812   |   <32.94 / 0.940>   |\n|          AT          |   32.10 / 0.928   | **31.83 / 0.902** |   31.26 / 0.918   |\n|    Input Jittering   |   34.06 / 0.942   |   31.28 / 0.857   |   31.04 / 0.912   |\n| SGD jittering (Ours) | **35.12 / 0.960** |   <31.46 / 0.884>   | **33.23 / 0.945** |\nTable 2: Image deblurring. Best performances in **bold**, second best in <...>.\n\n> Other robustness methods\n\nThank the reviewer for the suggestion. Prior work [1] has compared a MBAs trained with AT to end-to-end randomized smoothing (RS) for MRI reconstruction, and showed that AT improves robustness significantly than RS. It supports our choice of **AT as a strong baseline for evaluating robustness in IPs**. \n\nAs suggested, we evaluated our method against a widely used robustness technique Lipschitz regularization using spectral normalization (SN). SN constrains the Lipschitz constant of the network by bounding the spectral norm (i.e., largest singular value) of each weight matrix, which helps stabilize training and improve robustness. As shown in Table 3 result for 4xMRI reconstruction. SN promotes adv. robustness, but at the cost of reduced accuracy to some extend. This is likely due to its restrictive nature of SN limiting the model’s expressive power.\n\n|    PSNR/SSIM     | ID | Adv. Attack | OOD |\n|:--------------------:|:-------------:|:-------------:|:-------------:|\n|     MSE training     | <28.21 / 0.603> | 25.68 / 0.382 | 29.92 / <0.779> |\n|          AT          | 27.68 / 0.564 | **27.17** / <0.549> |  27.74 / 0.597 |\n|    Input Jittering   | 28.18 / 0.595 | 25.05 / 0.420 | <29.97> / 0.740 |\n| SGD jittering (Ours) | **28.22 / 0.607** | 26.77 / **0.552** | **30.36 / 0.788** |\n|  SN (new baseline) | 27.71 / 0.576  | <27.09> / 0.542 | 27.92 / 0.594 |\nTable 3: MRI reconstruction. Best performances in **bold**, second best in <...>.\n\n[1] Alkhouri et al. Robust physics-based deep MRI reconstruction via diffusion purification\n\nWe sincerely appreciate the reviewer’s constructive feedback. We believe the added experiments and clarification will strengthened the manuscript and hope they satisfactorily address your concerns.", "claims_and_evidence": "The claims are supported by clear evidence and proofs.", "methods_and_evaluation_criteria": "The method and evaluations make sense, but the method is similar to previous works like SGLD.  Therefore, the novelty is not very good.", "theoretical_claims": "I checked some proofs but not all proofs.", "experimental_designs_or_analyses": "Experiments are ok but can be improved when more inverse tasks like superresolution are added.", "relation_to_broader_scientific_literature": "The key contributions are similar to Stochastic gradient Langevin dynamics.", "essential_references_not_discussed": "The key contribution is similar to stochastic gradient Langevin dynamics (SGLD), the the comparison between the proposed method and SGLD is not discussed.", "other_strengths_and_weaknesses": "Strengths:\n1. Provides analysis linking noise injection to implicit regularization of gradient/Hessian smoothness, advancing understanding of robustness in iterative inverse solvers.\n2. Demonstrates effectiveness on real-world tasks (seismic deconvolution, MRI) where robustness and generalization are critical. The extension to proximal gradient (SPGD Jittering) highlights broader applicability.\n3. Well-structured presentation of theory, experiments, and ablation studies.\n\nWeaknesses:\n\n1. SGD jittering is similar to stochastic gradient Langevin dynamics, which makes the method less novel. In addition, the theoretical or empirical differences are not discussed.\n2. Theoretical guarantees are limited to denoising tasks; broader analysis for general inverse problems (e.g., ill-posedness, nonlinearity) is not provided.\n3. Experiments focus on a few inverse problems. Testing on more diverse tasks (e.g., super-resolution) could better validate generality.\n4. While outperforming MSE and adversarial training, comparisons to other robustness methods (e.g., gradient penalty, Lipschitz regularization) are missing.", "other_comments_or_suggestions": "The section structure should be reorganized. e.g., Sec. 3 only has one paragraph.", "questions_for_authors": "Please see the weakness part.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper considers model based architectures for solving inverse problems, and proposes a training strategy which inects noise at each iteration during SGD, called SGD jittering. The paper shows that jittering can generalize better and can be more robust to average case perturbations. 

The paper received four reviews and after rebuttal, all are supportive of publication. 

The reviewers find the paper's theoretical and empirical results interesting. 

The reviewers identified some minior weaknesses, including:
- Reviewer v9Xr finds the main weakness is the lack of baselines beyond unrolled networks; the authors added a Diffusion model as baseline. 
- Reviewer mUA8 notes that is would be interesting to add results beyond denoising, e.g., for compressive sensing. I agree with the reviewer that this would be interesting and with the authors that this is beyond of the scope of the paper.
- Reviewer KSe9 had some concerns regarding the statement that SGD jittering improves accuracy at the same time as robustness.

I also find the paper interesting and recommend it for acceptance. 

Besides the feedback from the reviewers, in order to understand the robustness benefits of the method better I think it would be interesting (in the final version or follow-up work) to evaluate out-of-distribution performance as a function of in-distribution performance (see e.g., Fig. 2 in Darestani et al. ``Measuring Robustness in Deep Learning Based Compressive Sensing’’, ICML 2021). By evaluating out-of-distribution performance like this it becomes clear whether the improved performance to out-of-distribution data (like in Figure 4b in the paper under review) is due to a robustness benefit or is because the in- and out-of-distribution performance improves (i.e., the algorithm still lies on a line as in Fig 2 in Darestani et al).
Decision: Accept (poster)