Title: Learning Provably Robust Estimators for Inverse Problems via Jittering
Venue: NeurIPS 2023
Accepted: yes
Keywords: Adversarial Training

Jittering

Denoising

Deconvolution

Compressive Sensing

Inverse Problems

Robustness
Abstract: Deep neural networks provide excellent performance for inverse problems such as denoising. However, neural networks can be sensitive to adversarial or worst-case perturbations. This raises the question of whether such networks can be trained efficiently to be worst-case robust. In this paper, we investigate whether jittering, a simple regularization technique that adds isotropic Gaussian noise during training, is effective for learning worst-case robust estimators for inverse problems. While well studied for prediction in classification tasks, the effectiveness of jittering for inverse problems has not been systematically investigated. In this paper, we present a novel analytical characterization of the optimal $\ell_2$-worst-case robust estimator for linear denoising and show that jittering yields optimal robust denoisers. Furthermore, we examine jittering empirically via training deep neural networks (U-nets) for natural image denoising, deconvolution, and accelerated magnetic resonance imaging (MRI). The results show that jittering significantly enhances the worst-case robustness, but can be suboptimal for inverse problems beyond denoising. Moreover, our results imply that training on real data which often contains slight noise is somewhat robustness enhancing.

Learning Provably Robust Estimators for Inverse Problems via Jittering

Abstract

Deep neural networks provide excellent performance for inverse problems such as denoising. However, neural networks can be sensitive to adversarial or worstcase perturbations. This raises the question of whether such networks can be trained efficiently to be worst-case robust. In this paper, we investigate whether jittering, a simple regularization technique that adds isotropic Gaussian noise during training, is effective for learning worst-case robust estimators for inverse problems. While well studied for prediction in classification tasks, the effectiveness of jittering for inverse problems has not been systematically investigated. In this paper, we present a novel analytical characterization of the optimal ℓ 2 -worst-case robust estimator for linear denoising and show that jittering yields optimal robust denoisers. Furthermore, we examine jittering empirically via training deep neural networks (U-nets) for natural image denoising, deconvolution, and accelerated magnetic resonance imaging (MRI). The results show that jittering significantly enhances the worst-case robustness, but can be suboptimal for inverse problems beyond denoising. Moreover, our results imply that training on real data which often contains slight noise is somewhat robustness enhancing.

1 Introduction

Deep neural networks achieve state-of-the-art performance for image reconstruction tasks including compressive sensing, super-resolution, and denoising. Due to their excellent performance, deep networks are now used in a variety of imaging technologies, for example in MRI and CT. However, concerns have been voiced that neural networks can be sensitive to worst-case or adversarial perturbations. Those concerns are fuelled by neural networks being sensitive to small, adversarially selected perturbations for prediction tasks such as image classification (Szegedy et al., 2014).

Recent empirical work for image reconstruction tasks (Huang et al., 2018; Antun et al., 2020; Genzel et al., 2022; Darestani et al., 2021) found that worst-case perturbations can have a significantly larger effect on the image quality than random perturbations. This sensitivity to worst-case perturbations is not unique to neural networks, classical imaging methods are similarly sensitive (Darestani et al., 2021).

This raises the question of whether networks can be designed or trained to be worst-case robust. A successful method proposed in the context of classification is adversarial training, which optimizes a robust or adversarial loss during training (Madry et al., 2018). However, the robust loss requires finding worst-case perturbations during training which is difficult and computationally expensive.

0

0

.

1

0

.

2

0

.

3

0

0

.

2

0

.

4

0

.

6

0

.

8

ϵ

2

/

E

[

∥

x

∥

2

2

]

Robust risk

ˆ

R

ϵ

/n

Linear denoising

0

0

.

1

0

.

2

0

.

3

ϵ

2

/

E

[

∥

x

∥

2

2

]

U-net denoising

0.00

0.01

0.02

0.03

0

.

5

1

1

.

5

ϵ

2

/

E

[

∥

Ax

∥

2

2

]

U-net MRI reconstruction

0.0

0.1

0.2

0.3

0.1

1.0

10

ϵ

2

/

E

[

∥

Ax

∥

2

2

]

U-net deconvolution

Standard

Jittering

Adversarial

In this work, we study jittering, a simple regularization technique that adds noise during training as a robustness-enhancing technique for inverse problems. It is long known that adding noise during training has regularizing effect and can be beneficial for generalization (Bishop, 1995; Holmstrom and Koistinen, 1992; Reed and Marks II, 1999). Prior work also studied adding noise for enhancing adversarial robustness for classification (Zantedeschi et al., 2017; Kannan et al., 2018; Gilmer et al., 2019). However, jittering has not been systematically studied as a robustness enhancing technique for training robust networks for inverse problems.

We consider the following signal reconstruction problem. Let f : R m → R n be an estimator (a neural network in practice) for a signal (often an image) x ∈ R n based on the measurement y = Ax + z ∈ R m , where A is a measurement matrix and z is random noise. We want to learn an estimator that has small robust risk defined as

The robust risk is the expected worst-case error with respect to a ℓ 2 -perturbation of norm at most ϵ of f measured with the mean-squared error.

Theoretical results. We start with Gaussian denoising of a signal lying in a subspace, and first characterize the optimal linear robust denoiser, i.e., the estimator that minimizes the robust risk R ϵ ( f ) . While the resulting estimator is quite intuitive, proving optimality is fairly involved and relies on interesting applications of Jensen's inequality.

Second, we show that the optimal linear robust estimator minimizes the Jittering-risk

where w ∼ N (0 , σ 2 w I ) is Gaussian jittering noise with noise variance σ 2 w that depends on the desired robustness level ϵ . This finding implies that instead of performing robust training via minimizing an empirical version of robust risk, we can train a denoiser via jittering, i.e., injecting Gaussian noise during training, at least for denoising a signal lying in a subspace. Figure 1, left panel, demonstrates the equivalence of training via minimizing a jittering risk and robust training numerically for the subspace model. It is evident that both methods of training yield an equally robust estimator.

Moreover, we discuss extensions of our theory for linear inverse problems y = Ax + z and find that jittering can result in suboptimal worst-case robust estimators for some classes of forward operators.

Empirical results for real-world denoising, deconvolution and compressive sensing. Jittering is also effective for learning robust neural network estimators for solving inverse problems in practice. Figure 1, second from left to right, depicts the worst-case risk achieved by training a U-net model for denoising, compressive sensing, and deconvolution, with standard training (blue), with jittering (purple), and with adversarial training (orange). For denoising, we see that jittering is as effective for obtaining a worst-case robust estimator as adversarial training, as suggested by theory. For compressive sensing and deconvolution, we find that jittering can be suboptimal beyond denoising, but is still effective for enhancing robustness.

Those findings make jittering a potentially attractive method for learning robust estimators in the context of inverse problems, since jittering can also be implemented easily and needs far less computational resources than adversarial training. Moreover, those findings imply that training on real data which often contains slight noise is somewhat robustness enhancing.

2 Related work

Empirical investigation of worst-case robustness for imaging. Several works investigated the sensitivity of neural networks for image reconstruction tasks to adversarial perturbations, for limited angle tomography (Huang et al., 2018), MRI and CT (Antun et al., 2020; Genzel et al., 2022; Darestani et al., 2021), and image-to-image tasks (colorization, deblurring, denoising, and superresolution) (Choi et al., 2022; Yan et al., 2022; Choi et al., 2019). Collectively, those works show that neural networks for imaging problems are significantly more sensitive to adversarial perturbations than to random perturbations, as expected. The effect of adversarial ℓ 2 -perturbations measured in mean-squared-error is roughly proportional to the energy of the perturbations in most of those problems, demonstrating that up to a constant (that might be large) neural networks can be relatively stable for imaging tasks. Classical reconstruction methods, in particular ℓ 1 -regularized least-squares, are similarly sensitive to adversarial perturbations (Darestani et al., 2021).

Learning robust methods with robust optimization. To learn robust classifiers, Madry et al. (2018) proposed to minimize a robust loss and to find worst-case perturbations during training with projected gradient descent. Adversarial training can be effective for learning robust methods, but is computationally expensive due to the cost of finding adversarial perturbations. A variety of heuristics exist to lower the computational cost of robust training for neural networks. For example, Raj et al. (2020) consider compressive sensing and CT reconstruction problems and propose to generate adversarial perturbations for training with an auxiliary network instead of solving a maximization problem. As another example, Wong et al. (2020) considers adversarial training of classifiers and propose to calculate adversarial perturbations during training by first randomly perturbing the initial point and then applying a single step of projected gradient descent.

Jittering for enhancing robustness in inverse problems. The literature is somewhat split on whether jittering is effective for enhancing worst-case robustness for imaging. Genzel et al. (2022) suggested that jittering can enhance worst-case robustness of networks trained on compressive sensing and CT tasks. Contrary, Gandikota et al. (2022) consider the robustness to ℓ ∞ -perturbations for neural-network based deblurring and observed that the DeepWiener architecture, trained via Jittering with noise levels sampled from a fixed interval, is sensitive to adversarial perturbations. In this work, we present the first systematic study on the effectiveness of jittering as a robustness-enhancing technique for inverse problems.

Robustness for inverse problems versus robustness for classification problems. Robustness in general and adding noise during training in particular, has been intensively studied in the classification setting. However, inverse problems and classification/prediction problems are very different. Adversarial robustness for classifiers is defined as the (average) minimal distance to the decision boundary, and random noise robustness as the minimal noise strength (for example the radius of Gaussian noise sphere) such that one likely crosses the decision boundary (Fawzi et al., 2018). For inverse problems, there is no notion of a decision boundary. Therefore, results and intuitions from classification, which are often based on geometric insights on distances to surfaces (see for example Fawzi et al. (2016) and Shafahi et al. (2019)) do not apply to inverse problems.

Jittering for enhancing robustness in classification. Prior work in classification considered Gaussian data augmentation or adding Gaussian noise during training (which is conceptually very similar to jittering) as an robustness-enhancing technique and found that adding noise enhances adversarial robustness, but reported mixed results on its effectiveness. Fawzi et al. (2018) proved for linear classifiers that adding Gaussian noise during training increases adversarial robustness, and Gilmer et al. (2019) demonstrated that empirically adding Gaussian noise during training also increases adversarial robustness for neural networks in the context of classification. Rusak et al. (2020) also found Gaussian noise addition beneficial for corruption robustness (including noise, compression and weather artifacts). Furthermore, Kannan et al. (2018) and Zantedeschi et al. (2017) considered adding Gaussian noise during training together with other regularization methods and report that adding noise at a fixed noise level alone yields a noticeable increase in robustness. Contrary, Carlini and Wagner (2017) reevaluated the methods proposed by Zantedeschi et al. (2017) and reported that the robustness gains are small compared to adversarial training.

Randomized smoothing. Randomized smoothing is a very successful technique for obtaining robust classifiers (Cohen et al., 2019; Carlini et al., 2023), and is based on constructing a smoothed classifier from a base classifier by averaging the base classifier's outputs under Gaussian noise perturbation. The smoothed classifier is provably robust within a specified radii, without making any restrictions on the base classifier. Despite similarities at first sight, randomized smoothing considers surrogate smoothed models, whereas jittering is a training technique (see the appendix on a detailed discussion).

3 Theory for robust reconstruction of a signal lying in a subspace

In this section, we characterize the optimal robust estimator for denoising a signal in a subspace. While the resulting estimator is quite intuitive, proving optimality is fairly involved and relies on interesting applications of Jensen's inequality. We then show that the optimal robust estimator is also the unique minimizer of the jittering loss. Finally, we conjecture a precise characterization of optimal estimators for linear inverse problems beyond denoising, and show that jittering can result in suboptimal estimators for linear inverse problems beyond denoising.

3.1 Problem setup

We consider a signal reconstruction problem, where the goal is to estimate a signal x based on a noisy measurement y = Ax + z , where z ∼ N (0 , σ 2 z 1 /m I ) is Gaussian noise and A ∈ R m × n a measurement or forward operator. The random noise is scaled so that the expected noise energy is E [ ∥ z ∥ 2 2 ] = σ 2 z . The random noise is denoted by z , to distinguish it from the adversarial noise or worst-case error, denoted by e . We assume that the signal is (approximately) chosen uniformly from the intersection of a sphere and a subspace. Specifically, the signal is generated as x = Uc , where c ∼ N (0 , σ 2 c 1 /d I ) is Gaussian and U ∈ R n × d is an orthonormal basis for a d -dimensional subspace of R n . The expected signal energy is E [ ∥ x ∥ 2 ] = σ 2 .

We consider a linear estimator of the form f ( y ) = Hy for estimating the signal from the measurement. For the standard reconstruction problem of estimating the signal x from the measurement y , performance is often measured in terms of the expected least-squared error. We are interested in robust reconstruction and consider the expected worst-case reconstruction error with respect to an ℓ 2 -perturbation, defined in equation (1), and given by

For ϵ = 0 , the robust risk reduces to the standard expected mean-squared error.

3.2 Denoising

We start with denoising where the forward map is the identity, i.e., A = I . The following theorem characterizes the optimal worst-case robust denoiser.

α

Scaling factor

0

Optimal robust denoising

1

5

0

.

0

0

.

5

ϵ

0.0

0.05

0.1

0.15

0

.

5

1

1

.

5

2

ϵ

2

/σ

2

c

Robust risk

R

ϵ

/n

Standard

Adversarial

Jittering

1

/σ

1

2

c

Theorem 3.1. For d → ∞ , the optimal worst-case estimator, i.e., the estimator minimizing the worst-case risk R ϵ ( f ) amongst all estimators of the form f ( y ) = Hy , is H = α UU T , where

The worst-case optimal estimator projects onto the signal subspace, and then shrinks towards zero, by a factor determined by the noise variance σ 2 z and the worst-case noise energy ϵ 2 . We consider the asymptotic setup where d → ∞ only for expositional convenience; our proof shows that the estimator in the theorem is also near optimal for finite d .

To understand the implications of the theorem, let us first consider the case where the worst-case perturbation is zero. Then, the optimal estimator simply projects on the signal-subspace and shrinks towards zero, by a factor of α = σ 2 c σ 2 c + σ 2 z d n . The larger the noise, the more shrinkage.

Next, consider the most interesting regime, where non-zero adversarial noise is present. If the adversarial noise energy is larger than the signal energy, the estimator projects onto zero. However, this is an extreme regime since the adversarial noise can cancel the signal, and no good estimate of the signal can be achieved.

For the more practical regime where the adversarial noise energy is smaller than the signal energy, the theorem states that the optimal estimator projects onto the signal-subspace and shrinks towards zerojust like the optimal estimator for the noise-free case-but this time by a factor α , that decreases in the adversarial noise energy ϵ 2 .

The proof of Theorem 3.1 is in the appendix. Note that for estimators H = α UU T , a worst-case perturbation can be computed in closed form for a fixed y and x : a worst-case perturbation is the vector that points into the direction of the signal plus noise lying in the signal subspace, i.e., e = U ϵ (1 -α ) c + α U T z ∥ (1 -α ) c + α U T z ∥ 2 . However, for a general estimator H , the perturbation can not be written in closed form, which makes proving optimality quite challenging. Our proof relies on a characterization of the inner maximization problem as the solution to an optimization problem in one variable, and several unusual applications of Jensen's inequality.

3.2.1 Robust denoisers via jittering

An important consequence of the characterization of the worst-case optimal estimator in Theorem 3.1 is that, at least for the linear setup considered in this section, a worst-case optimal estimator can be obtained by regularization with jittering.

2

.

5

2

No noise

Medium noise

High noise

Robustness beyond denoising

Recall from the introduction that regularization via jittering simply adds Gaussian noise w ∼ N (0 , σ 2 w I ) to the measurement during training. The jittering risk (2) of the estimator f ( y ) = Hy for denoising is J σ w ( f ) = E ( x , y ) , w [ ∥ f ( y + w ) -x ∥ 2 2 ] . Choosing the variance of the jittering noise level accordingly as a function of the desired robustness level ϵ yields an optimal worst-case robust estimator by minimizing the jittering loss, as formalized by the following corollary of Theorem 3.1. Corollary 3.2. For ϵ 2 < σ 2 c , the linear estimator f ( y ) = Hy that minimizes the jittering risk J σ w with noise level chosen as a function of the desired noise level ϵ as σ w ( ϵ ) = √ ϵ 2 σ 2 z d n + σ z √ d n σ c ϵ √ σ 2 c -ϵ 2 + σ 2 z d n d ( σ 2 c -ϵ 2 ) also minimizes the worst-case risk R ϵ .

Hence, if we aim for a robustness level ϵ < σ c , we can simply apply training via Jittering instead of adversarial training by choosing the Jittering noise level using the explicit formula for the jittering noise level σ w ( ϵ ) in corollary 3.2.

Figure 1, left panel, shows the results of numerical simulation for adversarial training and jittering. In the implementation we treat the linear reconstructions as neural networks with a single layer without bias and perform adversarial training and jittering for each perturbation level. The simulations show that the robust risk performance of the models are identical, as predicted by the theory. Details on how adversarial training is performed are in Section 4.

3.2.2 Robustness accuracy trade-off

Another consequence of Theorem 3.1 is an explicit robustness-accuracy trade-off: increased worstcase robustness comes at a loss of accuracy. In the practically relevant regime of 0 ≤ ϵ 2 < σ 2 c the standard risk of the optimal worst-case estimator f α ( ϵ ) is R 0 ( f α ( ϵ ) ) = σ 2 c · σ 2 z d n σ 2 c + σ 2 z d n -ϵ 2 . This expression yields the optimal standard error for ϵ = 0 and is strictly monotonically increasing in ϵ , hence showing the loss of accuracy when increasing robustness. Robustness-accuracy tradeoffs can also be observed in other machine learning settings, for classification and regression settings, see for example (Tsipras et al., 2019). For linear inverse problems with applications in control, robustness accuracy-tradeoffs were recently characterized by Lee et al. (2021) and Javanmard et al. (2020).

3.3 General linear inverse problems

In the previous section, we characterized the optimal worst-case robust estimator and found that jittering yields optimal robust denoisers. In this section, we derive a conjecture for the worst-case optimal robust estimator for more general linear inverse problems of reconstructing a signal x from a measurement y = Ax + z , with a forward operator A (with A = I in general), and show that this estimator is in general not equal to the estimator obtained with jittering, thus jittering is in general suboptimal.

̸

Optimal robust estimator. Let AU = WΛV T be the singular value decomposition of the matrix AU with singular values λ i . As formalized by Lemma A.1 of the appendix, the robust-risk (1) of the estimator f can be written as an expectation involving a minimization problem over a single variable (instead of a maximization over an n -dimensional variable, as in the original definition):

̸

Here, σ max is the largest singular value of the matrix H . In order to find the optimal robust estimator we wish to solve the optimization problem arg min H R ϵ ( H ) . The difficulty in solving this optimization problem is that we can't solve the minimization problem within the expectation (3) in closed form. In order to prove Theorem 3.1 for denoising (i.e., for A = I ) we derived an upper and a matching lower bound of the risks using several unusual applications of Jensen's inequality. The proof does not generalize in a straightforward manner to the more general case where A = I . However, for large d , the random variable v T ( I -1 λ HH T ) -1 v concentrates around it's expectation, and thus we conjecture that for large d , we can exchange expectation and minimization, which yields:

The expectation in the risk expression (4) can be explicitly computed, which yields the following conjecture for the worst-case optimal estimator:

Conjecture 3.3. For d →∞ the optimal worst-case estimator, i.e., the estimator minimizing the worst-case risk R ϵ ( f ) amongst all estimators of the form f ( y ) = Hy , is H = UV diag ( σ i ) W T with

̸

if λ i = 0 and σ i = 0 otherwise. Here, the parameter λ is a solution of:

The optimization problem involved is convex and box-constrained and can thus be solved numerically. Besides the argument above, we confirmed our conjecture with numerical simulations.

Optimal jittering estimator. Unlike for denoising, for general inverse problems, the jittering-risk minimizing estimator is in general not equal to the optimal worst-case estimator, but the two estimators are often close. The optimal estimator minimizing the jittering risk is given as (see Appendix C):

where as before AU = WΛV T is the singular value decomposition with singular values λ i . While the estimator (5) has the same form as the worst-case optimal estimator in Conjecture 3.3, the diagonal matrix in the two estimators is in general slightly different.

Numerical Simulation. The worst-case suboptimality of the jittering-risk optimal estimator (5) depends on the singular values of the matrix AU ; if they are equal the jittering-risk estimator is optimal, and if they are not equal there is typically a small gap. To illustrate the gap, we consider a forward operator A with linearly decaying singular values i n , for 1 ≤ i ≤ n with signal energy σ 2 c = 1 and noise level σ z = 0 . 2 . We compare the (conjectured) optimal robust estimator specified by Conjecture 3.3 with the optimal Jittering estimator (5) at noise level σ w , where σ w is optimized such that one obtains minimal robust risk R ϵ at a given perturbation level ϵ . The results in Figure 2, right panel, show a small gap in robust risk, which implies that Jittering is suboptimal for this case. However, simulations with varying forward operators and noise levels indicate that the gap is small relative to the robust risk of the standard estimator. Experiments on image deconvolution using U-Net presented in Section 4 show similar results.

4 Experiments

In this section, we train standard convolutional neural-networks with standard training, adversarial training, and jittering for three inverse problems: denoising images, image deconvolution, and compressive sensing, and study their robustness. We find that Jittering yields well-performing robust denoisers at a computational cost similar to standard training, which is significantly cheaper than adversarial training. We also find that jittering yields robust estimators for deconvolution and compressive sensing. This indicates that training on real data which often contains slight measurement noise is robustness enhancing.

4.1 Problem setup

We start by describing the datasets, networks, and methodology.

Natural images. We consider denoising and deconvolution of natural images, where our goal is to reconstruct an image x ∈ R n from a noisy measurement y = Ax + z , where z ∼ N (0 , σ 2 z 1 /n I ) is Gaussian noise and A a measurement matrix, which is equal to identity for denoising, and implements

a convolution for deconvolution. For deconvolution we use a 8 × 8 -sized discretization of the 2 -dimensional Gaussian normal distribution with standard deviation 2 . The kernel is visualized in Figure 3 in the appendix. We obtain train and validation datasets { ( x 1 , y 1 ) , . . . , ( x N , y N ) } of sizes 34 k and 4 k, respectively, from colorized images of size n = 128 · 128 · 3 generated by randomly cropping and flipping ImageNet images. The methods are tested on 2 k original-sized images.

Medical data. We also perform experiments on accelerated singlecoil magnetic resonance imaging (MRI) data, where the goal is to reconstruct an image x ∈ R n from a noisy and subsampled measurement in the frequency domain y = MFx + z ∈ R 2 m . We use the fastMRI singlecoil knee dataset (Zbontar et al., 2018), which contains the images x and fully sampled measurements ( M = I ). We process it by random subsampling at acceleration factor 4 and obtain train, validation and test datasets with approximately 31 k, 3 . 5 k and 7 k slices, respectively. While perturbations are sought in frequency domain, the inverse Fourier transform is applied to the measurements before feeding the 320 × 320 cropped and normalized images into the network.

Network architecture. We use the U-net architecture (Ronneberger et al., 2015) since it gives excellent performance for denoising (Brooks et al., 2019) and medical image reconstruction tasks, such as computed tomography (Jin et al., 2017) and is used as a building block for state-of-the-art methods for magnetic resonance imaging (Zbontar et al., 2018; Sriram et al., 2020). For natural images, we use a U-net with 3 × 3 padded convolutions with ReLU activation functions, 2 × 2 max-pooling layers for downscaling and transposed convolutions for upscaling. The network has 120 k parameters. For MRI reconstruction we use a U-Net architecture similar to Zbontar et al. (2018) with 3 × 3 padded convolutions, leaky ReLU activation function, 2 × 2 average pooling and transposed convolutions ( 480 k learnable parameters). We denote the U-Net by the parameterized mapping f θ : R m → R n in the following.

Evaluation. We evaluate networks by measuring its robustness via the empirical robust risk defined as ˆ R ϵ ( θ ) = 1 N ∑ N i =1 max ∥ e i ∥ 2 ≤ ϵ ∥ f θ ( y i + e i ) -x i ∥ 2 2 , For evaluation, the robust empirical risk is computed over the test set. We assess the accuracy by computing the standard empirical risk ˆ R 0 ( θ ) . Computing the robust empirical risk is non-trivial since it requires finding adversarial perturbations for solving the inner maximization problem. This is explained next. We also study the computational cost of the different methods, which we measure in terms of GPU cost and time.

Finding adversarial perturbations. To evaluate the empirical risk and for robust training, we need to compute adversarial perturbations e = arg max ∥ e ∥ 2 ≤ ϵ ∥ f θ ( y + e ) -x ∥ 2 2 . We find the perturbations by running N a projected gradient ascent steps, starting with initial perturbation e 0 = 0 and iterate

Here, P B (0 ,ϵ ) is the projection into the ℓ 2 -ball B (0; ϵ ) of radius ϵ around the origin. The gradient is normalized to facilitate step size optimization with multiplier 2 . 5 such that the iteration can reach and move along the boundary, as suggested by Madry et al. (2018).

Training methods. Standard training minimizes the standard empirical risk ˆ R 0 . Adversarial training minimizes the empirical robust risk ˆ R ϵ . To minimize the empirical robust risk, we approximate the inner maximization, max ∥ e ∥ 2 ≤ ϵ ∥ f θ ( y i + e ) -x i ∥ 2 2 , with ∥ f θ (˜ y i ) -x i ∥ 2 2 , where ˜ y i is the adversarially perturbed measurement computed as described above. Training via jittering minimizes

where the jittering level σ w is chosen depending on the desired robustness level. We treat the jittering noise level as a hyperparameter optimized using the validation dataset (shown in the appendix). Jittering is practically implemented via performing the SGD update rule θ i +1 = θ i -η M ∑ M i =1 ∇ θ ∥ f θ ( y i + w i ) -x i ∥ 2 , with M samples per batch and learning rate η . For jittering, the network output is calculated on the noisy input y i + w i , instead of y i (standard training) or

0

.

2

0

.

4

0

.

6

0

0

.

2

0

.

4

Jittering level

σ

w

Robust risk

ˆ

R

ϵ

/n

Hyperparameter search

0

0

.

1

0

.

2

0

.

3

0

0

.

2

0

.

4

0

.

6

ϵ

2

/

E

[

∥

x

∥

2

2

]

Jittering level

σ

w

Jittering level

σ

w

(

ϵ

)

0

0

.

1

0

.

2

0

.

3

0.02

0.04

0.06

0.08

ϵ

2

/

E

[

∥

x

∥

2

2

]

Standard risk

ˆ

R

0

/n

Standard risk

Jittering

Prediction

Adversarial

Jittering (optimal)

Standard

y i + e i (adversarial training). To approximate the expectation in the jittering risk we draw independent jittering noise samples w i in each iteration of SGD.

-3

Throughout, we use PyTorch's Adam optimizer with learning rate 10 and batch size 50 for natural images, and 10 -2 and 1 for MRI data. As perturbation levels, we consider values within the practically interesting regime of ϵ 2 / E [ ∥ Ax ∥ 2 2 ] < 0 . 3 for natural images and 0 . 03 for MRI data. Note that for ϵ 2 > E [ ∥ x ∥ 2 2 ] , Theorem 3.1 predicts for denoising ( A = I ) that the optimal robust estimator is zero everywhere. Figure 7 in the appendix shows that for large perturbations ϵ the trained U-net denoiser also maps to zero.

4.2 Results

We now discuss the results of the denoising, deconvolution, and compressive sensing experiments.

Robust and standard performance. Figure 1, shows that the standard estimator is relatively robust for Gaussian denoising and increasingly sensitive for more ill-posed problems (deconvolution and compressive sensing). The experiments further show that jittering is effective for enhancing robustness, in particular relative to the sensitivity of the standard estimator. Nevertheless, as suggested by theory, we see a gap between the robust risk of adversarial training and jittering for image deconvolution and compressive sensing. For Gaussian denoising, however, Jittering is particularly effective and yields increasingly better performing networks in terms of standard risks for larger perturbations.

Choice of the jittering level. The results are based on choosing the jittering noise levels via hyperparameter search for each task. Figure 3 shows the results for Gaussian denoising: It can be seen that the choice of noise level is important for minimizing the robust risk. The estimated noise levels also aligns well the theoretical prediction. Details on this and the parameter choices for deconvolution and the compressive sensing experiments are in the appendix.

Computational complexity. We measured the GPU time until convergence and memory utilization of the methods and present the results in the Table 1 of the appendix. Performing adversarial training is by a factor of the projected gradient ascent steps more expensive than standard training. Moreover, training via jittering has similar computational cost as standard training, since it solely consists of drawing and adding Gaussian noise on the training data.

Target

clean

perturbed

Standard estimator

clean

perturbed

Adversarial estimator

clean

perturbed

Jittering estimator

Denoising

MRI-CS

Deconv.

Visual reconstructions. For the linear subspace setting adversarial training and jittering are equivalent. For Gaussian denoising with a neural network, however, they perform differently. For larger perturbations jittering tends to yield smoother images than networks trained adversarially, as can be seen in the example reconstructions shown in Figure 4. This effect is particularly noticeable for the Gaussian deconvolution task. We quantified the effect using the total variation norm and present the results in the appendix. Moreover, we discuss an approximation to jittering, Jacobian regularization, which similarly enhances robustness. It is computationally more expensive, but yields less smooth reconstructions.

5 Conclusion

In this paper, we characterized the optimal worst-case robust estimator for Gaussian subspace denoising and found that the optimal estimator can be provably learned with jittering. Our results for training neural networks for Gaussian denoising of images show that jittering enables the training of neural networks that are as robust as neural networks trained adversarially, but at a fraction of the computational cost, and without the hassle of having to find adversarial perturbations. While we demonstrated that jittering can yield suboptimal robust estimators in general, in practice, jittering is effective at improving the robustness for compressive sensing and image deconvolution. Moreover, our results imply that training on real data that contains slight measurement noise is robustness enhancing.

Reproducability The repository at https://github.com/MLI-lab/robust_ reconstructors_via_jittering contains the code to reproduce all results in the main body of this paper.

Acknowledgments A.K. and R.H. are supported by the Institute of Advanced Studies at the Technical University of Munich, the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) - 456465471, 464123524, the DAAD, and the German Federal Ministry of Education and Research, and the Bavarian State Ministry for Science and the Arts. M.S. is supported by the Packard Fellowship in Science and Engineering, a Sloan Research Fellowship in Mathematics, an NSF-CAREER under award #1846369, DARPA FastNICS programs, and NSF-CIF awards #1813877 and #2008443.

References

Antun, F. Renna, C. Poon, B. Adcock, and A. C. Hansen. On instabilities of deep learning in image reconstruction and the potential costs of AI. In Proceedings of the National Academy of Sciences , 2020.

M. Bishop. Training with Noise is Equivalent to Tikhonov Regularization. Neural Computation , 7 (1):108-116, 1995.

Boyd and L. Vandenberghe. Convex Optimization . Cambridge University Press, 2004.

Brooks, B. Mildenhall, T. Xue, J. Chen, D. Sharlet, and J. T. Barron. Unprocessing Images for Learned Raw Denoising. In IEEE Conference on Computer Vision and Pattern Recognition , 2019.

Carlini and D. Wagner. MagNet and "Efficient Defenses Against Adversarial Attacks" are Not Robust to Adversarial Examples. arXiv:1711.08478 , 2017.

Carlini, F. Tramer, K. D. Dvijotham, L. Rice, M. Sun, and J. Z. Kolter. (certified!!) adversarial robustness for free! In International Conference on Learning Representations , 2023.

J.-H. Choi, H. Zhang, J.-H. Kim, C.-J. Hsieh, and J.-S. Lee. Evaluating Robustness of Deep Image Super-Resolution Against Adversarial Attacks. In IEEE International Conference on Computer Vision , 2019.

J.-H. Choi, H. Zhang, J.-H. Kim, C.-J. Hsieh, and J.-S. Lee. Deep Image Destruction: Vulnerability of Deep Image-to-Image Models against Adversarial Attacks. In International Conference on Pattern Recognition , 2022.

Cohen, E. Rosenfeld, and Z. Kolter. Certified Adversarial Robustness via Randomized Smoothing. In International Conference on Machine Learning , 2019.

Z. Darestani, A. S. Chaudhari, and R. Heckel. Measuring robustness in deep learning based compressive sensing. In International Conference on Machine Learning , 2021.

Fawzi, S.-M. Moosavi-Dezfooli, and P. Frossard. Robustness of Classifiers: From Adversarial to Random Noise. In Advances in Neural Information Processing Systems , 2016.

Fawzi, O. Fawzi, and P. Frossard. Analysis of classifiers' robustness to adversarial perturbations. Mach Learn , 107(3):481-508, 2018.

V. Gandikota, P. Chandramouli, and M. Moeller. On adversarial robustness of deep image deblurring. In IEEE International Conference on Image Processing , 2022.

Genzel, J. Macdonald, and M. Marz. Solving Inverse Problems With Deep Neural Networks - Robustness Included. IEEE Transactions on Pattern Analysis and Machine Intelligence , 45: 1119-1134, 2022.

Gilmer, N. Ford, N. Carlini, and E. Cubuk. Adversarial Examples Are a Natural Consequence of Test Error in Noise. In International Conference on Machine Learning , 2019.

Holmstrom and P. Koistinen. Using additive noise in back-propagation training. IEEE Transactions on Neural Networks , 3(1):24-38, 1992.

Huang, T. Würfl, K. Breininger, L. Liu, G. Lauritsch, and A. Maier. Some Investigations on Robustness of Deep Learning in Limited Angle Tomography. In Medical Image Computing and Computer Assisted Intervention , 2018.

Javanmard, M. Soltanolkotabi, and H. Hassani. Precise tradeoffs in adversarial training for linear regression. In Conference on Learning Theory , 2020.

H. Jin, M. T. McCann, E. Froustey, and M. Unser. Deep Convolutional Neural Network for Inverse Problems in Imaging. IEEE Transactions on Image Processing , 26(9):4509-4522, 2017.

Kannan, A. Kurakin, and I. Goodfellow. Adversarial Logit Pairing. arXiv:1803.06373 [cs] , 2018.

D. Lee, T. T. C. K. Zhang, H. Hassani, and N. Matni. Adversarial Tradeoffs in Linear Inverse Problems and Robust State Estimation. arXiv:2111.08864 [cs, eess] , 2021.

Madry, A. Makelov, L. Schmidt, D. Tsipras, and A. Vladu. Towards Deep Learning Models Resistant to Adversarial Attacks. In International Conference on Learning Representations , 2018.

Raj, Y. Bresler, and B. Li. Improving Robustness of Deep-Learning-Based Image Reconstruction. In International Conference on Machine Learning , 2020.

Reed and R. J. Marks II. Neural Smithing: Supervised Learning in Feedforward Artificial Neural Networks . MIT Press, 1999.

Reed, R. Marks, and S. Oh. Similarities of error regularization, sigmoid gain scaling, target smoothing, and training with jitter. IEEE Transactions on Neural Networks , 6(3):529-538, 1995.

Ronneberger, P. Fischer, and T. Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation. In Medical Image Computing and Computer-Assisted Intervention , 2015.

Rusak, L. Schott, R. S. Zimmermann, J. Bitterwolf, O. Bringmann, M. Bethge, and W. Brendel. A Simple Way to Make Neural Networks Robust Against Diverse Image Corruptions. In European Conference on Computer Vision , 2020.

Salman, G. Yang, J. Li, P. Zhang, H. Zhang, I. Razenshteyn, and S. Bubeck. Provably Robust Deep Learning via Adversarially Trained Smoothed Classifiers. In Advances in Neural Information Processing Systems , 2019.

Salman, M. Sun, G. Yang, A. Kapoor, and J. Z. Kolter. Denoised Smoothing: A Provable Defense for Pretrained Classifiers. In Advances in Neural Information Processing Systems , 2020.

Shafahi, W. R. Huang, C. Studer, S. Feizi, and T. Goldstein. Are adversarial examples inevitable? In International Conference on Learning Representations , 2019.

Sriram, J. Zbontar, T. Murrell, A. Defazio, C. L. Zitnick, N. Yakubova, F. Knoll, and P. Johnson. End-to-End Variational Networks for Accelerated MRI Reconstruction. In Medical Image Computing and Computer Assisted Intervention , 2020.

Szegedy, W. Zaremba, I. Sutskever, J. Bruna, D. Erhan, I. J. Goodfellow, and R. Fergus. Intriguing properties of neural networks. In International Conference on Learning Representations , 2014.

Tsipras, S. Santurkar, L. Engstrom, A. Turner, and A. Madry. Robustness may be at odds with accuracy. In International Conference on Learning Representations , 2019.

Wong, L. Rice, and J. Z. Kolter. Fast is better than free: Revisiting adversarial training. In International Conference on Learning Representations , 2020.

Yan, J. Zhang, J. Feng, M. Sugiyama, and V. Y. F. Tan. Towards Adversarially Robust Deep Image Denoising. In International Joint Conference on Artificial Intelligence , 2022.

Zantedeschi, M.-I. Nicolae, and A. Rawat. Efficient Defenses Against Adversarial Attacks. In ACM Workshop on Artificial Intelligence and Security , 2017.

Zbontar, F. Knoll, A. Sriram, M. J. Muckley, M. Bruno, A. Defazio, M. Parente, K. J. Geras, J. Katsnelson, H. Chandarana, Z. Zhang, M. Drozdzal, A. Romero, M. Rabbat, P. Vincent, J. Pinkerton, D. Wang, N. Yakubova, E. Owens, C. L. Zitnick, M. P. Recht, D. K. Sodickson, and Y. W. Lui. fastMRI: An Open Dataset and Benchmarks for Accelerated MRI. arXiv:1811.08839 [cs.CV] , 2018.

A Proof of Theorem 3.1

In the main body we stated an analytical characterization of the optimal worst-case robust denoiser. We present the proof in the following and show that the risk

is minimized by H = σ UU T with

The scaling factor σ as a function of the noise level ϵ of the optimal worst-case estimator is shown in Figure 2.

We start by proving a lower bound of the risk, which relies on a characterization of the maximization and many unexpected applications of Jensens inequality.

We then compute the risk for H = σ UU T and show that it is equivalent to the lower bound on the risk.

A.1 Lower bounding the risk

Towards lower bounding the risk we define for notational convenience

With we get

Here, we defined e ∥ = U T e and z ∥ = U T z for notational convenience, and the last inequality follows by adding U T ⊥ e = 0 as a constraint to the maximization, which gives a lower bound.

Next, note that since z = [ z ∥ z ⊥ ] ∼ N (0 , σ 2 z /n I ) is Gaussian, the vector [ z ∥ -z ⊥ ] is equally Gaussian distributed N (0 , σ 2 z /n I ) . This implies that

Then

where the last step follows from Jensens inequality. Thus, we have shown that

For simplicity of exposition, we drop the ∥ -notation. Thus, with a slight abuse of notation, the vectors on the left and right hand side have different dimensions. On the left hand side, z , e ∈ R n , while on the right hand side z , e ∈ R d and z throughout has iid N (0 , σ 2 z /n ) entries.

Next, we'll apply the lemma below for characterizing the maximization inside of the expectation. The proof is in Section A.3. Similar computations as used to prove the lemma are on page 19-20 in the paper Lee et al. (2021) for deriving robustness-accuracy trade-off bounds.

Lemma A.1. For any z ,

Using Lemma A.1 the term within the expectation is

Now let's rewrite using the singular value decomposition H ∥ = VΣW T , with orthogonal matrices V and W . With this notation, the terms above can be rewritten as

where w i and v i are the i -th column vectors of the orthogonal matrices W and V , respectively. Using similar calculations for the other summands, the optimization problem on the right hand side of inequality (7) is

where minimization above is over the orthonormal matrices V , W ∈ R d × d and over the singular values σ i . Equality (i) follows by noting that ( σ i w T i -v i T ) c is Gaussian with variance ( σ 2 i +1 -2 σ i ( v i T w i )) σ 2 c d and ( σ i w T i ) z is Gaussian with variance σ 2 i σ 2 z n and therefore we can replace the expectation with respect to ( c , z ) with the expectation with respect to g ∼ N (0 , I ) and suitable scaling. Moreover, inequality (ii) follows from the fact that σ i ≥ 0 and g 2 i ≥ 0 .

Continuing, first note that the function x 2 y is convex in ( x, y ) when y > 0 . Also the extended value function is increasing in the first input and decreasing in the second input. Furthermore the mappings ( x, z ) ↦→ x -1 and ( x, z ) ↦→ 1 -x 2 z are convex and concave. Thus by the composition rule of convex functions we conclude that the functions

are jointly convex in ( x, z ) .

Jensen's inequality states that for a convex function ψ we have ∑ i z 2 i ψ ( x i ) ∑ i z 2 i ≥ ψ ( ∑ i z 2 i x i ∑ i z 2 i ) . Thus by Jensen's inequality the sum in the expectation in the right-hand-side of the equation above can be lower-bounded as

where ¯ σ ( g ) = ∑ d i =1 g 2 i ∥ g ∥ 2 2 σ i .

Now consider the event E = { g : ∥ g ∥ 2 2 ≥ (1 -δ ) d } which holds with probability at least 1 -e -δ 2 2 d . On this event, we have

Using the same argument as before the right hand side of the above inequality is jointly convex in ( λ, ¯ σ ( c )) . Since partial minimization of a jointly convex function preserves convexity we conclude that the function

is convex in ¯ σ . Thus by using convexity in terms of ¯ σ , applying Jensen's inequality (i.e., E [ ψ (¯ σ )] ≥ ψ ( E [¯ σ ]) ) we have

where we defined ¯ σ = ∑ d i =1 E g |E [ g 2 i ∥ g ∥ 2 2 ] σ i . Putting things together, we have shown that

where the last inequality follows from σ max ≥ ¯ σ . For d →∞ , we can choose δ arbitrarily small, which yields

Solving the optimization problem (9) : Consider the inner minimization problem in equation (9), i.e.,

where c ( σ ) = σ 2 c ( σ -1) 2 + σ 2 z d n σ 2 for notational convenience. Since f is differentiable on ( σ 2 , ∞ ) we can calculate its critical point λ ∗ by setting the derivative with respect to λ to zero, which yields

From this expression, we see that the constraint λ ∗ > σ 2 is satisfied and by convexity of f we know that λ ∗ is the unique minimizer. Hence, we have:

It follows that

Calculating the derivative of g ( σ ) and setting it to zero, we get:

The left factor is non-negative, and the right factor is zero if σ 2 c > ϵ 2 and if

If σ 2 c ≤ ϵ 2 , then the function g ( σ ) is monotonically increasing on [0 , ∞ ) and hence σ ∗ = 0 is the minimizer.

A.2 Upper bound for the risk of the estimator H = σ UU T

We upper bound the risk of the estimator H = σ UU T . For H = σ UU T the risk (6) becomes

Using the optimal λ ∗ , by equation 10, we get

where equation (i) follows by using Jensen's inequality once again (specifically, using ( E [ ∥ g ∥ 2 ]) 2 ≤ E [ ∥ g ∥ 2 2 ] = d ). Noting that (ii) is equal to the lower bound of the risk for any symmetric H in equation (11) shows that H = σ UU T with the optimal parameter σ ∗ derived above is optimal.

A.3 Proof of Lemma A.1

The optimization problem

The corresponding Lagrangian is, for λ ≥ 0

The Lagrange Dual Function is can be written as

Using that min e 2 c T e + e T Be = -c T B -1 c , we get, provided that λ ≥ λ max ( H T H ) ,

where the last equality follows from the Woodburry Identity. Thus, the dual problem is

Even though the primal problem (13) is non-convex, strong duality holds, since the primal program is a quadratic program that is strictly feasible, see Boyd and Vandenberghe (2004, Appendix B.1). The primal problem is strictly feasible if there exists a vector e such that ϵ 2 -e T e > 0 . This is trivially satisfied as long as ϵ > 0 .

B Additional proofs for denoising

We state two more proofs on optimal worst-case denoisers and jittering.

B.1 Proof of Corollary 3.2

In the main text it was stated that linear estimators f ( y ) = Hy that minimize the jittering risk J σ w with noise level chosen as a function of the desired noise level ϵ as

also minimizes the worst-case risk R ϵ .

The result follows from Theorem 3.1. For that let f r and f j be linear estimators minimizing the robust risk R ϵ and the jittering risk J σ w , respectively. By Theorem 3.1, the two estimators are scaled projections onto the subspace, i.e., f r ( y ) = α r UU T and f j ( y ) = α j UU T with

Setting α r = α j and solving for the standard deviation σ w yields the result.

B.2 Form of perturbations for estimators H = α UU T

We noted in the main body that for estimators H = α UU T worst-case perturbations can be computed in closed form for fixed y and x . We calculate:

Thus the perturbation points into the direction of the signal plus noise lying in the signal subspace.

C Theory for general linear inverse problems

In the following, we consider linear inverse problems y = Ax + z , with a linear forward operator A ∈ R m × n and noise z ∼ N (0 , σ 2 z /m I ) . For denoising ( A = I ) we stated an explicit analytical characterization of the optimal worst-case estimator and presented a proof in the appendix. The proof, however, does not generalize in a straightforward manner to more general inverse problems. We conjecture the worst-case optimal linear estimator for large dimensions d and present numerical simulation results. Moreover, we state the proof for the optimal jittering estimator, and demonstrate that it can yield suboptimal worst-case robust estimators in general.

C.1 Optimal worst-case robust estimator

As formalized by Lemma A.1, the robust-risk (1) of the estimator f can be written as an expectation involving a minimization problem over a single variable (instead of a maximization over an n -

dimensional variable, as in the original definition):

̸

Here, σ i are the singular values of the matrix H . In order to find the optimal robust estimator we wish to solve the optimization problem arg min H R ϵ ( H ) . The difficulty in solving this optimization problem is that we can't solve the minimization problem within the expectation (14) in closed form. In order to prove Theorem 3.1 for denoising (i.e., for A = I ) we derived an upper and a matching lower bound of the risks using several unusual applications of Jensen's inequality. The proof does not generalize in a straightforward manner to the more general case where A = I . However, for large d , the random variable v T ( I -1 λ HH T ) -1 v concentrates around it's expectation, and thus we conjecture that for large d , we can exchange expectation and minimization, which yields:

Based on equation (15) we derive a characterization of the optimal worst-case estimator. We proceed similar to the denoising case and start with rearranging the terms in expectation:

Now, let AU = WΛV T be the singular value decomposition of the matrix AU and H = UVΣW T the singular value decomposition of H . We then have:

For the individual parts in the summation it follows:

where we define c i = v i T c using the i -th column vector v i of the matrix V . Similarly, we set z i = w i T z , with w i the i -th column vector of W , and get for the other parts:

Hence, for the term in expectation in Eq. (15) we get:

where we note that c i ( σ i λ i -1)+ z i σ i are iid zero-mean gaussian with variance σ 2 c d ( λ i σ i -1) 2 + σ 2 z m σ 2 i and g i ∼ N (0 , 1) . From this it follows, assuming the robust risk conjecture (15), that the optimal robust estimator minimizes:

We first calculate the unconstrained minimizer of the function f and get:

̸

For λ i = 0 we obtain ∂f ∂σ i = 0 ⇒ σ i = 0 and for λ i = 0 :

The solutions σ ∗ i, ± obtained are for suitable a ( λ ) of the form:

On further examining the constraint σ 2 i ≤ λ in the problem (16) note that for real numbers a > b it holds √ a 2 -b 2 ≥ a -b . Hence, it follows:

Therefore, we can rule out σ i, + and set σ ∗ i = σ ∗ i, -. Finally, inserting σ ∗ i for 1 ≤ i ≤ d into the robust risk (17) yields:

The optimization problem involved is convex and box-constrained and can thus be solved numerically. Besides the argument above, we confirmed our conjecture with numerical simulations.

C.2 Optimal jittering estimator

We now derive the optimal jittering estimator, i.e., the estimator f ( y ) = Hy that minimizes the jittering risk

where as before the signal x is assumed to lie within a subspace, i.e., x = Uc . We first calculate the expectation by using that c ∼ N (0 , σ 2 c /d I ) , z ∼ N (0 , σ 2 z /m I ) and w ∼ N (0 , σ 2 w I ) are Gaussian distributed. The jittering risk is

Hence, the optimal jittering estimator minimizes

where X = σ 2 c d AU ( AU ) T + ( σ 2 w + σ 2 z /m ) I and Y = σ 2 c d AUU T . Using matrix calculus we calculate the optimal estimator as

which yields H ∗ = Y T X -1 . Now, let AU = WΛV T be the singular value decomposition of AU . Then:

We get the optimal jittering estimator as:

C.3 Numerical simulations supporting Conjecture 3.3

In the following we present numerical simulations for general linear inverse problems in the subspace model to support the Conjecture 3.3 further. The results show that the (empirical) optimal robust risk, obtained via adversarial training, is the same as the robust risk of the conjectured optimal estimator. Moreover, the results show that jittering can yield suboptimal robust estimators in some cases.

We consider linear inverse problems y = Ax + z with the signal x lying in a subspace x = Uc with c ∼ N (0 , σ 2 c /d I ) and Gaussian noise z ∼ N (0 , σ 2 z /n I ) . For the simulations, we choose σ c / √ d = 1 and σ z / √ n = 0 . 1 with dimensions d = 50 and n = 100 . Moreover, we consider two diagonal forward operators A = diag ( λ i ) : an operator with linear decaying singular values ( λ i = i n ) and one with geometrically decaying singular values λ i = 0 . 7 i .

Optimal worst-case estimator. We first estimate the optimal robust risks by performing adversarial training and compare it with the robust risk of the conjectured optimal estimator. Adversarial training is performed as described in Section 4, where we generate data { ( y i , x i ) } using the respective forward model in the subspace. The estimator f ( y ) = Hy can be viewed as a neural network with one layer and without bias. Figure 5 shows that the robust risk of the conjectured estimator is essentially the same as the (empirical) optimal robust risk.

)

0.0

0.05

0.1

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

ϵ

2

/σ

2

c

Robust risk

R

ϵ

/n

Linear decay (

λ

i

=

i/n

)

0.0

0.01

0.02

0

.

44

0

.

46

0

.

48

0

.

5

ϵ

2

/σ

2

c

Geometric decay (

λ

i

= 0

.

7

i

Conjectured

Adv. training

Suboptimality of jittering. We further investigate whether optimal robust estimators can be obtained via jittering. To that end, we calculate the minimal robust risk attainable via jittering

where f ( y ) = Hy is a linear reconstruction operator as before. We make use of the analytic expression of the jittering estimator in Eq. (21) and calculate the attainable robust risk by minimizing the robust risk of jittering with respect to the jittering noise level. The results are compared to the (conjectured) optimal robust risks via Eq. (18) and the robust risk of the standard estimator. Figure 6 shows the results of the calculations for the forward operator with linear and geometrically decaying singular values described above. It can be seen that the standard estimator is noticeable less robust compared to denoising setups. Moreover, a gap can be observed between the robust risk obtained via jittering and the optimal robust risks.

0.0

0.03

0.06

0.09

0.1

1.0

ϵ

2

/σ

2

c

Robust risk

R

ϵ

/n

Linear decay

0.0

0.06

0.12

0.18

ϵ

2

/σ

2

c

Geometric decay

Standard

Adversarial

Jittering

D Details on the experimental results and further experimental results

In this section, we present details on the experimental results in Section 4, and present further experimental results on U-nets trained with robustness-enhancing methods for image denoising, deconvolution, and compressive sensing. The setup and methods considered are as described in the main body.

D.1 Optimal robust denoiser for large perturbation levels

In the main body, we presented empirical results for perturbation levels in the range ϵ 2 / E [ ∥ x ∥ 2 2 ] ∈ [0 , 0 . 3] , since this is the practically most relevant regime. Figure 7 shows the risk of linear estimators and U-nets trained adversarially for perturbation levels in the range ϵ 2 / E [ ∥ x ∥ 2 2 ] ∈ [0 , 1 . 5] . From those plots, we see that the transition at ϵ 2 / E [ ∥ x ∥ 2 2 ] = 1 predicted by Theorem 3.1 for the estimator to map to zero occurs for the subspace model (as predicted by the theory) as well as for the U-net for image denoising.

D.2 Convolution kernel for deconvolution experiments

In addition to experiments on denoising image data, we consider a deconvolution setup y = k ⋆ x + z in this work. The kernel k is Gaussian, applied channel-wise and visualized in Figure 8.

D.3 Quantifying the smoothing effect of Jittering

In the main text we presented visual example reconstructions and observed that neural networks trained via Jittering yield smoother reconstructions compared to adversarial and standard training.

0

0

.

5

1

1

.

5

0

0

.

5

1

1

.

5

Perturbation level

ϵ

2

/

E

[

∥

x

∥

2

]

Robust risk

ˆ

R

ϵ

/n

Linear subspace denoising

0

0

.

5

1

1

.

5

Perturbation level

ϵ

2

/

E

[

∥

x

∥

2

]

U-Net image denoising

No noise

Medium noise

High noise

0.01

0.02

0.03

0.04

We further quantify this observation by calculating the total variation norm of reconstructed images for the deconvolution task. The results are shown Figure 9 and confirm the observation that networks trained via Jittering yield smoother images.

D.4 Hyperparameter selection

In the experiments we treat the Jittering noise level σ w as a hyperparameter, which we optimize over a validation dataset to obtain robust estimators at the desired perturbation levels. The hyperparameter search is performed by choosing a grid of jittering noise levels for each task. For each noise level neural networks (U-Nets) are trained via Jittering, and subsequently evaluated on the considered perturbation levels. Figures 10, 11 and 12 show the robust risks of jittering for the considered tasks, as well as the derived jittering choice rule. The smooth curves on the left panels represent the robust risk at a particular robustness level and are obtained by applying uniform filters on the evaluation results. It can be seen that for image denoising the empirical jittering choice is close to the prediction from theory.

0

0

.

1

0

.

2

0

.

3

0.4

0.6

0.8

1.0

Perturbation level

ϵ

2

/

E

[

∥

Ax

∥

2

]

Total variation TV

(

f

ϵ

)

/

10

5

Smoothness of reconstructions

Standard

Adversarial

Jittering

0

.

2

0

.

4

0

.

6

0

0

.

2

0

.

4

Jittering level

σ

w

Robust risk

ˆ

R

ϵ

/n

Robust risks of jittering (Denoising)

0

0

.

1

0

.

2

0

.

3

0

0

.

2

0

.

4

0

.

6

Perturbation level

ϵ

2

/

E

[

∥

x

∥

2

]

Jittering level

σ

w

Jittering level choice (Denoising)

D.5 Computational complexity

We measured the GPU time until convergence and memory utilization of the robustness-enhancing schemes on the task of Gaussian denoising of colorized images. Figure 13 shows the training error of adversarial training, training via jittering and standard training as a function of the number of epochs. It shows networks trained at two perturbation levels for adversarial training and jittering (parameter choice taken from Figure 10). We find that all methods require a similar number of epochs for convergence (roughly 600 epochs). Table 1 presents the measured GPU time until convergence and average memory consumption. It can be seen that adversarial training is by a factor of the projected gradient ascent steps ( 3 in this plot) more expensive than jittering. Moreover, training via jittering has similar computational cost as standard training in terms of GPU time. All three methods require a similar amount of GPU memory to train.

]

0

0

.

5

1

1

.

5

0

.

5

1

1

.

5

Jittering level

σ

w

/

10

-

5

Robust Risk

ˆ

R

ϵ

/n

Robust risks of jittering (MRI)

0.00

0.01

0.02

0.03

0

0

.

5

1

1

.

5

2

Perturbation level

ϵ

2

/

E

[

∥

Ax

∥

2

2

Jittering level

σ

w

/

10

-

5

Jittering level choice (MRI)

2

3

4

0

.

5

1

1

.

5

Jittering level

σ

w

Robust risk

ˆ

R

ϵ

/n

Robust risks of jittering (Deconvolution)

0

0

.

1

0

.

2

0

.

3

0

2

4

Perturbation level

ϵ

2

/

E

[

∥

Ax

∥

2

2

]

Jittering level

σ

w

Jittering level choice (Deconvolution)

D.6 Required compute

The experimental results presented in this paper were computed using on-premise infrastructure equipped with Nvidia RTX A6000 GPUs. In total, 18 GPU days were spent for the presented colorized image denoising, 20 GPU days for colorized image deconvolution and 15 GPU days for the compressive sensing experiments.

E Discussion of the related work on randomized smoothing

Randomized smoothing is a very successful technique for obtaining robust classifiers (Cohen et al., 2019; Carlini et al., 2023). Randomized smoothing constructs a smoothed classifier based on a base classifier by averaging the base classifier's outputs under Gaussian noise perturbation. The smoothed classifier allows for certified radii in which it is provably robust, without making any restrictions on the base classifier. However, Salman et al. (2020) demonstrated that it can give loose bounds, since the base classifier is not trained to be robust to Gaussian noise. For that reason Salman et al.

10

0

10

1

10

2

10

3

10

-

1

.

2

10

-

1

10

-

0

.

8

Epochs trained

Training error

ˆ

R

ϵ

/n

10

0

10

1

10

2

10

3

10

-

2

10

-

1

.

5

Epochs trained

Training error

ˆ

J

σ

w

/n

Jittering

ϵ

2

/

E

[

∥

x

∥

2

] = 0

.

06

ϵ

2

/

E

[

∥

x

∥

2

] = 0

.

03

Standard training

(2020) propose denoised smoothing, which considers a composition of the base classifier with a denoising method. At first sight, randomized smoothing might sound similar to the Jittering approach investigated here. However, as we argue below, randomized smoothing is conceptually very different from Jittering.

Given a classifier f : R d →{ 1 , . . . , K } randomized smoothing constructs a smoothed classifier g from the classifier f as:

̸

where the parameter σ 2 controls the robustness-accuracy tradeoff.

For an inverse problem, where we aim to reconstruct a signal x ∈ R n from a measurement y ∈ R m using a given reconstruction method f , replacing the 0 / 1 by the ℓ 2 loss yields:

For a linear estimator f ( y ) = Hy we see that g ( y ) = f ( y ) , so for the linear setting considered in the theory part of this paper randomized smoothing would not change the original estimator. If one considers f : R n → [0 , 1] n the smoothed estimator g ( y ) differs from f ( y ) and robustness gains can be expected, which follows from Salman et al. (2019), Lemma 1. In summary, randomized smoothing is very different to jittering in that it constructs a surrogate smoothed model g ( y ) based on a given fixed estimator f ( y ) , whereas jittering is a training technique.

F Regularizing beyond jittering for enhancing robustness

Training neural networks via jittering, with noise levels chosen for larger perturbations, yields smoother reconstructions compared to adversarial training (see results in Section 4). In this section,

we investigate two related regularization methods, ℓ 2 -regularization and Jacobian regularization, discuss the connection to regularization with jittering, and present experimental results for denoising grayscale images.

F.1 ℓ 2 - and Jacobian regularization in the subspace model

In the subspace model we established that the optimal jittering estimator is also worst-case optimal, when using a suitable choice of noise level σ w ( ϵ ) . It turns out, that jittering can further be approximated with an explicit regularizer, Jacobian regularization. For the linear setup considered here, this approximation becomes exact and therefore Jacobian regularization also enables training a worst-case robust estimator. Specifically, using the linear approximation of the function f around the point y , we get

Here, J y is the Jacobian of the function f at y . The approximation is good for small values of the noise variance σ 2 w , and is exact for the linear estimator f ( y ) = Hy we consider in this section. The approximate relation (22) motivates the Jacobian regularized risk, defined as

The connection between jittering and Jacobian regularization is well known in the literature and discussed by Reed et al. (1995). Recall that for the linear estimator considered in this section the approximation in equation (22) is exact, and therefore Jacobian regularization is equivalent to jittering. Thus, Jacobian regularization yields a provably robust estimator, if the regularization parameter is chosen as λ = σ 2 w ( ϵ ) according to corollary 3.2.

For the linear case, Jacobian regularization is even equivalent to ℓ 2 -regularization, since the Jacobian of the function f ( y ) = Hy is J y = H , and thus even ℓ 2 -regularization yields a robust estimator.

F.2 Experimental results on grayscale image denoising

In the following we present results on Gaussian denoising of grayscale images. While ℓ 2 regularization is equivalent to jittering in the subspace model, we find that the parameter choice λ = σ 2 w ( ϵ ) does not yield robust neural networks using ℓ 2 regularization. In contrast, Jacobian regularization turns out to be quite effective for learning neural network denoisers, but is computationally demanding compared to jittering.

F.2.1 Problem setup

We consider once again the dataset of natural image and convert the images to grayscale. We perform Gaussian denoising, i.e. the problem is to reconstruct the image x from a measurement y = x + z , with z ∼ N (0 , σ 2 z /n I ) and σ z / √ n = 0 . 2 .

The estimators are chosen as neural networks (U-nets) with the same architecture as for colorized images. We consider adversarial training, jittering and standard training as baseline and compare against ℓ 2 and Jacobian regularization:

ℓ 2 regularization. Implemented as weight-decay in PyTorch's SGD optimizer to minimize

Jacobian regularization. We train networks with the Jacobian regularized empirical risk

where J y i is the Jacobian of the network f θ with respect to it's input (not it's parameters) at y i . This regularization can be viewed as an approximation of the jittering risk, as described in

0

0

.

1

0

.

2

0

.

3

0

0

.

2

0

.

4

0

.

6

Perturbation level

ϵ

2

/

E

[

∥

x

∥

2

2

]

Robust risk

ˆ

R

ϵ

/n

0.0

0.1

0.2

0.3

0.01

0.03

0.05

Perturbation level

ϵ

2

/

E

[

∥

x

∥

2

2

]

Standard Risk

ˆ

R

0

/n

Jacobian

Jittering

ℓ

2

-reg.

subsection F.1.Calculating the full Jacobian J y i with PyTorch requires n -many calls of the backward function, which is very expensive, since n is large. To mitigate this cost, we approximate the norm of the Jacobian, ∥ J y i ∥ 2 F with ∥ ∥ ∥ J y i T w ∥ ∥ ∥ 2 2 , where w ∼ N (0 , I ) . This approximation of the norm concentrates around the actual squared norm of the Jacobian, and only costs one call of the PyTorch-backward function.

For the experiments we use stochastic gradient descent (SGD) with learning rate 10 -2 , momentum 0 . 9 and batch size 100 . We evaluate using the empirical pixel-wise robust risk ˆ R ϵ /n .

F.2.2 Results

The experimental results, plotted in Figure 14, show that the networks trained with jittering and Jacobian regularization have similar robust risks compared to the adversarial trained one. Weight-decay or ℓ 2 regularization yields worse performing estimators than jittering and Jacobian regularization. While for the linear subspace setting, adversarial training, Jacobian and ℓ 2 regularization are equivalent, for Gaussian denoising they perform differently. Figure 15 shows that Jacobian regularization, unlike jittering, does not yield smoothed images for larger perturbations. However, Jacobian regularization requires approximately 1 . 4 as much GPU memory and 4 times more time per epoch.

(a) Original

(b) Noisy Image

(c) Standard

(d) Robust Tr. (

medium

ϵ

)

(e) Jittering (

medium

ϵ

)

medium

ϵ

)

(g) Robust Tr. (

large

ϵ

)

(h) Jittering (

large

ϵ

)

large

ϵ

)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work considers introducing Jittering to the inverse problems for the benefits of robustness to the l2-worst-case. In this work, some analytical results are provided in proving the robust estimation with some assumptions, which I think are reasonable. Then, empirical experiments verify the effectiveness on the robustness improvement at a relatively lower computational cost compared to adversarial training.\n\nSTRENGTHS:\n1. The motivation of this work is quite nice, towards a provable robust estimator.\n2. The assumptions are reasonable.\n3. Existing results show somewhat promising performance.\n\nWEAKNESSES:\n1. Though some analytical results are presented, the actual realization and numerical implementation of the algorithm  are not articulated well I think, e.g., the assumed energy level of signals and noises are normally unknown, so that more details and algorithmic implementation discussions can be presented in a more systematic way, rather than just being mentioned in fractures in Section 4.\n2. Several experiments are conduced, but the presented results are limited in each experiment, e.g., cherry pick can be done if only a few images are presented.\n3. Some parts of the paper I read a few times, but I still found a bit difficult back and forth to relate all analytical results to the experiments in practical applications.\n\nQUESTIONS:\n1. Some statements lack clarity, e.g., ‘we study jittering, a simple regularization technique that adds noise during training as a robustness-enhancing technique for inverse problems’……as a regularization technique, I would expect a concrete implementation or algorithmic description to make the paper more readable to general readers. (Similar as the bullet 3 in *Weakness); \n\n2. Line 90-94, page 3, [1][2] are mentioned. I quickly checked [1], and it seems like it already implemented the Jittering to the task of inverse problem. It would be nice to provide more details on the difference and relation to [1]. Though [1] has been cited, no further explanations or comparisons are provided later on in this work. As far as I have understood, this work is posed as a step forward upon [1] with some analytical results? This I believe should be clarified. \n\n3. It mentioned the random smoothing method in Section 2 and also the appendix. I understand that RS is significantly different this work. In RS, lots of samplings are required in evaluations. I’m wondering as an expectation over ‘w’ as in (2), there should also be some procedures requiring samplings either for training or evaluations or for sufficing assumptions. Could you please elaborate this part? (Similar as the bullet 3 in *Weakness) \n\n4. In the analytical results of this work, $\\sigma_c$, $\\sigma_z$ are the essentials. Could you summarize or provide  a general procedure of the relevant implementation procedures or simulation verifications on these variables when they are unknown in the practical settings.  (Similar as the bullet 3 in *Weakness), and the same applies to Section 3.3. \n\n5.  In section 4, besides the bullet 2 in *Weakness, the authors can provide some synthetic experiments to show and explain more detailed results or analysis, as these mentioned variables in the assumptions can be known explicitly and also at least can help the readability. \n\n6. It seems that the hyper-parameter tuning can be time consuming, if too many searches are needed to bring good results.\n\n7. One major claimed benefit is the efficiency, however, only Table 1 is provided roughly, where the complexity analysis is lacking. It should also be nice to mention the computational part involved in finding the hyper parameters. I roughly checked the codes of defining “param” argument and it seems that it just pre-gave a  energy value and heuristically proceeds the setup of other relevant params to run the training. (Similar as the bullet 3 in *Weakness) . For instance, though the compared the adversarial training is more expensive, it doesn't need too many pre-requisite or tuning and can be flexibly applied in various practical tasks.\n\nIf this paper can be put in a clearer way connecting the theoretical results and the implementation, and also the empirical experiments can be more practically feasible and convincing in the results, this work would deserve an improved score and I would like to increase my evaluation.\n\n[1] M. Genzel, J. Macdonald, and M. Marz. Solving Inverse Problems With Deep Neural Networks- Robustness Included. IEEE TPAMI 2022\n[2] K. V. Gandikota, P. Chandramouli, and M. Moeller. On adversarial robustness of deep image deblurring. In IEEE International Conference on Image Processing, 2022.\n\nLIMITATIONS:\nThe authors should verify how convenient and efficient the proposed method can be applied in practical settings and explain some pipelines, in line with the claimed provability of the robust estimator. It seems unclear to me that how promising it could be in generally practical applications, especially considering the claimed analytical aspects.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work studies the design of estimators for the solutions of inverse problems that are adversarially robust, in the sense that they minimize the maximum MSE after being contaminated with an additive and L2 bounded perturbation. The authors show that, for linear estimators and for signals lying on linear subspaces, the optimal solution (i.e. optimally robust) is attained by training the estimator with “Jiterring”, i.e. by minimizing a loss over the randomly perturbed inputs. For more general inverse problems (when the forward operator is not an identity), the authors show that the resulting estimator from Jittering is not optimally robust in general, but their difference can be small. Moreover, they show numerically that this gap is small in practice, thus leading to similar performance (as that obtained by adversarial training) alas with significant increase in computational efficiency.\n\nSTRENGTHS:\n- Neat and elegant idea.\n- Very clear presentation - this was a pleasure to read.\n- The results are novel and interesting.\n\nWEAKNESSES:\n- The results are somewhat limited, and hold for linear estimators with potentially overly simple signal models.\n\nQUESTIONS:\n1. As the authors mention, estimating the robust risk is non-trivial, since this involves the optimization of a non-convex/non-concave problem for non-convex functions $f(\\cdot)$ (as in the case of U-net). In light of this, the plots in Fig. 1 are not “exactly” the robust risk, but rather numerical approximations to it (except for the linear case). The authors might wish to clarify this in the caption and/or the description of Fig. 1 in the text.\n2. The role of the symmetry of H is not completely clear to me. When describing their results colloquially, the authors mention that, broadly speaking, training linear estimators with Jittering provides an estimator with minimal robust risk if the signal lies on a subspace. However, as written, the formal statement of Theorem 3.1 requires the estimator not only to be linear but also symmetric. So is, is symmetry required for Thm 3.1 to hold, or does symmetry arise as a property of the optimal estimator? I have a related question for Conjecture 3.3 - is symmetry here required, and in particular, $H = UV\\Sigma W^T$ need not be symmetric, as written.\n3. Their results are stated for $d\\to \\infty$. It is unclear why this is needed, or how the other dimensions and parameters scale with d. \n4. In Eq (3), do the authors mean to require $\\lambda \\geq max_i \\sigma_i^2$?\n5. Can the authors comment some more on the choice of norms? More precisely, the authors have focused on the MSE as loss (L2) as well as L2 bounded perturbations. I wonder to what extent their conclusions might extend to other norms. Moreover, it is possible that for different choices, some of the limitations in the analysis (e.g. in proving their conjecture) might be resolved. On the same vein, and because of their signal model, I wonder if there are any connections to the work of [Awasthi, Pranjal, et al. \"Adversarially robust low dimensional representations.\" COLT 2021], that the authors could leverage.\n\nLIMITATIONS:\nSome limitations are commented on throughout the text. Might be nice to stress or clarify them further.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study the effectiveness of jittering in the setting of inverse problems, specifically considering denoising, compressive sensing, and deconvolution problems. Jittering is a well-known regularization technique for classification problems. The authors prove in the linear setting that the robust risk estimator can be probed to be the optimal estimator when learned with jittering for Gaussian subspace denoising. In addition, the authors demonstrate experimentally jittering is effective at improving the robustness of compressive sensing and image deconvolution, even though it yields suboptimal estimates.\n\nSTRENGTHS:\n1). The theoretical contribution is straightforward but has a strong contribution in the denoising setting. Even though the authors studied a linear model they were able to provide insights in the non-linear model (U-net) that was studied empirically. Refer to Figure 3 (Middle Plot), Cor.3.2 can accurately predict the optimal jittering level.\n\n2). The authors extend theoretical contribution for general inverse problems, section 3.3. \n\n3). The authors do a thorough investigation on the effectiveness of jittering empirically and theoretically for denoising inverse problems.\n\nWEAKNESSES:\n1). The motivation of the paper seems a bit unclear, it appears to be motivated by results from classification problems rather than image reconstruction results in inverse problems. This approach seems slightly ad-hoc because the authors choose a popular regularization technique for classification problems and tested if it was effective for inverse problems without stating their intuition for it being effective or being of interest to the community. \n\n2). Section 4.2 lacks some metrics, specifically, Section 4.2 Figure 4 states \"Jittering yields robust estimators, but at the same time yields smoother reconstructions.\". This claim seems pretty strong considering it is not reported over a large test set with appropriate metrics for image reconstruction. \n\n3).  The section for 3.3 appears disconnected from section 4, it's unclear whether the linear model studied in this section (3.3) is representative of the behavior of the non-linear model utilized in section 4. \n\n\n\n\n\nMinor Weakness-\nTypo Figure 3- \"The jittering estimators estimators are similarly robust as adversarial\ntraining (Figure 1), but attain lower standard risks (right panel)\" Repeated word estimators estimators\n\nQUESTIONS:\n1). Could you provide metrics for a test set confirming jittering yields smoother reconstructions? \n\n2). Could the authors justify why they studied a linear model instead a non-linear model, a model more similar to a U-net? \n\n3). Could the authors explain the connection of section 3.3 to experiments in Section 4? Does the linear model exhibit enough similar behavior as the non-linear model to justify it as an appropriate model for general linear inverse problems? If so, could you please provide some numerical evidence?\n\nLIMITATIONS:\nI do not see any negative societal impacts of this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}