SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Anselm Krainovic
Technical University of Munich anselm.krainovic@tum.de
University of Southern California  
soltanol@usc.edu
Reinhard Heckel
reinhard.heckel@tum.de
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.
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.
Figure 1: Jittering yields worst-case robust reconstruction methods. The plots show the pixelwise empirical robust risks  $\hat{R}_{\epsilon} / n$  of models trained to minimize the robust risk  $R_{\epsilon}$  and jittering risk  $J_{\sigma_w}$ , respectively, with suitable choices of jittering levels  $\sigma_w(\epsilon)$ . The optimal robust risk is approximated via adversarial training. The shaded areas are  $66\%$  confidence intervals. Left panel is for subspace denoising where jittering and robust training are provably equivalent, the panels from second left to right are for image reconstruction problems with the U-net, where jittering is particularly effective for denoising.
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\colon \mathbb{R}^m\to \mathbb{R}^n$  be an estimator (a neural network in practice) for a signal (often an image)  $\mathbf{x}\in \mathbb{R}^n$  based on the measurement  $\mathbf{y} = \mathbf{A}\mathbf{x} + \mathbf{z}\in \mathbb{R}^{m}$ , where  $\mathbf{A}$  is a measurement matrix and  $\mathbf{z}$  is random noise. We want to learn an estimator that has small robust risk defined as
$$
R _ {\epsilon} (f) = \mathbb {E} _ {(\mathbf {x}, \mathbf {y})} \left[ \max  _ {\| \mathbf {e} \| _ {2} \leq \epsilon} \| f (\mathbf {y} + \mathbf {e}) - \mathbf {x} \| _ {2} ^ {2} \right]. \tag {1}
$$
The robust risk is the expected worst-case error with respect to a  $\ell_2$ -perturbation of norm at most  $\epsilon$  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_{\epsilon}(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
$$
J _ {\sigma_ {w}} (f) = \mathbb {E} _ {(\mathbf {x}, \mathbf {y}), \mathbf {w}} \left[ \| f (\mathbf {y} + \mathbf {w}) - \mathbf {x} \| _ {2} ^ {2} \right], \tag {2}
$$
where  $\mathbf{w} \sim \mathcal{N}(0, \sigma_w^2 \mathbf{I})$  is Gaussian jittering noise with noise variance  $\sigma_w^2$  that depends on the desired robustness level  $\epsilon$ . 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  $\mathbf{y} = \mathbf{A}\mathbf{x} + \mathbf{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.
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  $\ell_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  $\ell_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  $\ell_{\infty}$ -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.
We consider a signal reconstruction problem, where the goal is to estimate a signal  $\mathbf{x}$  based on a noisy measurement  $\mathbf{y} = \mathbf{A}\mathbf{x} + \mathbf{z}$ , where  $\mathbf{z} \sim \mathcal{N}(0, \sigma_z^2 1/m\mathbf{I})$  is Gaussian noise and  $\mathbf{A} \in \mathbb{R}^{m \times n}$  a measurement or forward operator. The random noise is scaled so that the expected noise energy is  $\mathbb{E}\left[\| \mathbf{z}\|_2^2\right] = \sigma_z^2$ . The random noise is denoted by  $\mathbf{z}$ , to distinguish it from the adversarial noise or worst-case error, denoted by  $\mathbf{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  $\mathbf{x} = \mathbf{U}\mathbf{c}$ , where  $\mathbf{c} \sim \mathcal{N}(0, \sigma_c^2 1/d\mathbf{I})$  is Gaussian and  $\mathbf{U} \in \mathbb{R}^{n \times d}$  is an orthonormal basis for a  $d$ -dimensional subspace of  $\mathbb{R}^n$ . The expected signal energy is  $\mathbb{E}\left[\| \mathbf{x}\|_2^2\right] = \sigma_c^2$ .
We consider a linear estimator of the form  $f(\mathbf{y}) = \mathbf{H}\mathbf{y}$  for estimating the signal from the measurement. For the standard reconstruction problem of estimating the signal  $\mathbf{x}$  from the measurement  $\mathbf{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  $\ell_2$ -perturbation, defined in equation (1), and given by
$$
R _ {\epsilon} (f) = \mathbb {E} _ {(\mathbf {x}, \mathbf {y})} \left[ \max  _ {\| \mathbf {e} \| _ {2} \leq \epsilon} \| \mathbf {H} (\mathbf {y} + \mathbf {e}) - \mathbf {x} \| _ {2} ^ {2} \right].
$$
For  $\epsilon = 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.,  $\mathbf{A} = \mathbf{I}$ . The following theorem characterizes the optimal worst-case robust denoiser.
Figure 2: Robust reconstruction of signals lying in a subspace. The left panel shows the scaling factor  $\alpha$  of the optimal robust denoiser according to Theorem 3.1 at noise levels  $\sigma_z\sqrt{d / n}\in$ $\{0,0.4,1.2\}$  and signal energy  $\sigma_c^2 = 1$ . The right panel depicts the robust risks of standard training and jittering as well as the optimal robust risk for an inverse problem beyond denoising, with the details stated in subsection 3.3.
Theorem 3.1. For  $d\to \infty$  , the optimal worst-case estimator, i.e., the estimator minimizing the worst-case risk  $R_{\epsilon}(f)$  amongst all estimators of the form  $f(\mathbf{y}) = \mathbf{H}\mathbf{y}$  , is  $\mathbf{H} = \alpha \mathbf{U}\mathbf{U}^T$  , where
$$
\alpha = \left\{ \begin{array}{l l} \sigma_ {c} ^ {2} - \frac {\epsilon \sigma_ {c} \sigma_ {z} \sqrt {\frac {d}{n}}}{\sqrt {\sigma_ {c} ^ {2} + \sigma_ {z} ^ {2} \frac {d}{n} - \epsilon^ {2}}} & i f \sigma_ {c} ^ {2} > \epsilon^ {2} \\ \hline \sigma_ {c} ^ {2} + \sigma_ {z} ^ {2} \frac {d}{n} & \\ 0 & e l s e. \end{array} \right.
$$
The worst-case optimal estimator projects onto the signal subspace, and then shrinks towards zero, by a factor determined by the noise variance  $\sigma_z^2$  and the worst-case noise energy  $\epsilon^2$ . We consider the asymptotic setup where  $d \to \infty$  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  $\alpha = \frac{\sigma_c^2}{\sigma_c^2 + \sigma_z^2\frac{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 zero—just like the optimal estimator for the noise-free case—but this time by a factor  $\alpha$ , that decreases in the adversarial noise energy  $\epsilon^2$ .
The proof of Theorem 3.1 is in the appendix. Note that for estimators  $\mathbf{H} = \alpha \mathbf{U}\mathbf{U}^T$ , a worst-case perturbation can be computed in closed form for a fixed  $\mathbf{y}$  and  $\mathbf{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.,  $\mathbf{e} = \mathbf{U}\epsilon \frac{(1 - \alpha)\mathbf{c} + \alpha\mathbf{U}^T\mathbf{z}}{\|\langle(1 - \alpha)\mathbf{c} + \alpha\mathbf{U}^T\mathbf{z}\|_2}$ . However, for a general estimator  $\mathbf{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.
Recall from the introduction that regularization via jittering simply adds Gaussian noise  $\mathbf{w} \sim \mathcal{N}(0, \sigma_w^2 \mathbf{I})$  to the measurement during training. The jittering risk (2) of the estimator  $f(\mathbf{y}) = \mathbf{H}\mathbf{y}$  for denoising is  $J_{\sigma_w}(f) = \mathbb{E}_{(\mathbf{x}, \mathbf{y}), \mathbf{w}}\left[\| f(\mathbf{y} + \mathbf{w}) - \mathbf{x}\|_2^2\right]$ . Choosing the variance of the jittering noise level accordingly as a function of the desired robustness level  $\epsilon$  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  $\epsilon^2 < \sigma_c^2$ , the linear estimator  $f(\mathbf{y}) = \mathbf{H}\mathbf{y}$  that minimizes the jittering risk  $J_{\sigma_w}$  with noise level chosen as a function of the desired noise level  $\epsilon$  as  $\sigma_w(\epsilon) = \sqrt{\frac{\epsilon^2\sigma_z^2\frac{d}{n} + \sigma_z\sqrt{\frac{d}{n}}\sigma_c\epsilon\sqrt{\sigma_c^2 - \epsilon^2 + \sigma_z^2\frac{d}{n}}}{d(\sigma_c^2 - \epsilon^2)}}$  also minimizes the worst-case risk  $R_{\epsilon}$ .
Hence, if we aim for a robustness level  $\epsilon < \sigma_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  $\sigma_w(\epsilon)$  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 worst-case robustness comes at a loss of accuracy. In the practically relevant regime of  $0 \leq \epsilon^2 < \sigma_c^2$  the standard risk of the optimal worst-case estimator  $f_{\alpha(\epsilon)}$  is  $R_0(f_{\alpha(\epsilon)}) = \sigma_c^2 \cdot \frac{\sigma_z^2 \frac{d}{n}}{\sigma_c^2 + \sigma_z^2 \frac{d}{n} - \epsilon^2}$ . This expression yields the optimal standard error for  $\epsilon = 0$  and is strictly monotonically increasing in  $\epsilon$ , 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  $\mathbf{x}$  from a measurement  $\mathbf{y} = \mathbf{A}\mathbf{x} + \mathbf{z}$ , with a forward operator  $\mathbf{A}$  (with  $\mathbf{A} \neq \mathbf{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  $\mathbf{A}\mathbf{U} = \mathbf{W}\boldsymbol{\Lambda}\mathbf{V}^T$  be the singular value decomposition of the matrix  $\mathbf{A}\mathbf{U}$  with singular values  $\lambda_{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):
$$
R _ {\epsilon} (\mathbf {H}) = \mathbb {E} _ {\mathbf {v}} \left[ \min  _ {\lambda \geq \sigma_ {\max } ^ {2}} \lambda \epsilon^ {2} + \mathbf {v} ^ {T} \left(\mathbf {I} - \frac {1}{\lambda} \mathbf {H} \mathbf {H} ^ {T}\right) ^ {- 1} \mathbf {v} \right], \quad \mathbf {v} = (\mathbf {H} \mathbf {A} - \mathbf {I}) \mathbf {x} + \mathbf {H} \mathbf {z}. \tag {3}
$$