SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 10: Jittering hyperparameter results for image denoising. Neural networks (U-nets) are trained via jittering on a grid of jittering noise levels  $\sigma_w$  (with fixed noise level  $\sigma_z / \sqrt{n} = 0.25$ ). The models are subsequently evaluated on the robust risk  $R_{\epsilon}$ , which is calculated for each of the perturbation levels  $\epsilon$ . Each line in the left panel corresponds to the robust risk at one perturbation level. The derived choice rule is displayed in the right panel. It can be seen that for denoising the empirical jittering choice matches the prediction (dashed line) from theory well (Corollary 3.2, with  $n = 128 \cdot 128 \cdot 3 = 49152$  and subspace dimension  $d = 32000$ ).
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.
Figure 11: Jittering hyperparameter search for compressive sensing. Neural networks (U-nets) are trained via jittering on a grid of noise levels  $\sigma_w$ . The models are subsequently evaluated on the robust risk  $R_{\epsilon}$ , which is calculated for each of the perturbation levels  $\epsilon$ . Each line in the left panel corresponds to the robust risk at one perturbation level. The derived choice rule is displayed in the right panel.
Figure 12: Jittering hyperparameter results for image deconvolution. Neural networks (U-nets) are trained via jittering on a grid of jittering noise levels  $\sigma_w$  (with fixed noise level  $\sigma_z / \sqrt{n} = 0.25$ ). The models are subsequently evaluated on the robust risk  $R_{\epsilon}$ , which is calculated for each of the perturbation levels  $\epsilon$ . Each line in the left panel corresponds to the robust risk at one perturbation level. The derived choice rule is displayed in the right panel.
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.
Figure 13: Training metrics of U-nets for adversarial training, jittering and standard training. The plot shows that the convergence rates are very similar for the considered setup. The curves show the training of networks for perturbation levels  $\epsilon^2 / \mathbb{E}[\| \mathbf{x} \|^2] = 0.03$  (red) and 0.06 (green). Training the methods takes roughly 600 epochs for convergence.
Method	Total GPU hours	Memory
Adversarial training	12.2 h	9971 MiB
Training via jittering	3.0 h	9523 MiB
Standard training	3.0 h	9523 MiB
Table 1: GPU time until convergence and memory utilization of adversarial training, jittering and standard training for the task of denoising colorized images. The required number of epochs is estimated from the data visualized in Figure 13. GPU time and memory consumption are measured on a Nvidia RTX A6000 GPU.
(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\colon \mathbb{R}^d\to \{1,\ldots ,K\}$  randomized smoothing constructs a smoothed classifier  $g$  from the classifier  $f$  as:
$$
g (\mathbf {x}) = \arg \min _ {k \in \{1, \dots , K \}} \mathbb {E} _ {\mathbf {e} \sim \mathcal {N} (0, \sigma^ {2} \mathbf {I})} \left[ \mathbb {1} \{f (\mathbf {x} + \mathbf {e}) \neq k \} \right],
$$
where the parameter  $\sigma^2$  controls the robustness-accuracy tradeoff.
For an inverse problem, where we aim to reconstruct a signal  $\mathbf{x} \in \mathbb{R}^n$  from a measurement  $\mathbf{y} \in \mathbb{R}^m$  using a given reconstruction method  $f$ , replacing the  $0/1$  by the  $\ell_2$  loss yields:
$$
\begin{array}{l} g (\mathbf {y}) = \arg \min  _ {\mathbf {x}} \mathbb {E} _ {\mathbf {e} \sim \mathcal {N} (0, \sigma^ {2} \mathbf {I})} \left[ \| f (\mathbf {y} + \mathbf {e}) - \mathbf {x} \| _ {2} ^ {2} \right] \\ = \mathbb {E} _ {\mathbf {e} \sim \mathcal {N} (0, \sigma^ {2} \mathbf {I})} [ f (\mathbf {y} + \mathbf {e}) ]. \\ \end{array}
$$
For a linear estimator  $f(\mathbf{y}) = \mathbf{H}\mathbf{y}$  we see that  $g(\mathbf{y}) = f(\mathbf{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\colon \mathbb{R}^n\to [0,1]^n$  the smoothed estimator  $g(\mathbf{y})$  differs from  $f(\mathbf{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(\mathbf{y})$  based on a given fixed estimator  $f(\mathbf{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,  $\ell_2$ -regularization and Jacobian regularization, discuss the connection to regularization with jittering, and present experimental results for denoising grayscale images.
F.1  $\ell_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  $\sigma_w(\epsilon)$ . 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  $\mathbf{y}$ , we get
$$
\mathbb {E} _ {\mathbf {w}} \left[ \| f (\mathbf {y} + \mathbf {w}) - \mathbf {x} \| _ {2} ^ {2} \right] \approx \mathbb {E} _ {\mathbf {w}} \left[ \| f (\mathbf {y}) + \mathbf {J} _ {\mathbf {y}} \mathbf {w} - \mathbf {x} \| _ {2} ^ {2} \right] = \| f (\mathbf {y}) - \mathbf {x} \| _ {2} ^ {2} + \sigma_ {w} ^ {2} \| \mathbf {J} _ {\mathbf {y}} \| _ {F} ^ {2}. \tag {22}
$$
Here,  $\mathbf{J}_{\mathbf{y}}$  is the Jacobian of the function  $f$  at  $\mathbf{y}$ . The approximation is good for small values of the noise variance  $\sigma_w^2$ , and is exact for the linear estimator  $f(\mathbf{y}) = \mathbf{H}\mathbf{y}$  we consider in this section. The approximate relation (22) motivates the Jacobian regularized risk, defined as
$$
\operatorname {J a c} _ {\lambda} (f) = \mathbb {E} _ {(\mathbf {x}, \mathbf {y})} \left[ \| f (\mathbf {y}) - \mathbf {x} \| _ {2} ^ {2} + \lambda \| \mathbf {J} _ {\mathbf {y}} \| _ {F} ^ {2} \right] \tag {23}
$$
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  $\lambda = \sigma_w^2 (\epsilon)$  according to corollary 3.2.
For the linear case, Jacobian regularization is even equivalent to  $\ell_2$ -regularization, since the Jacobian of the function  $f(\mathbf{y}) = \mathbf{H}\mathbf{y}$  is  $\mathbf{J}_{\mathbf{y}} = \mathbf{H}$ , and thus even  $\ell_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  $\ell_2$  regularization is equivalent to jittering in the subspace model, we find that the parameter choice  $\lambda = \sigma_w^2 (\epsilon)$  does not yield robust neural networks using  $\ell_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  $\mathbf{x}$  from a measurement  $\mathbf{y} = \mathbf{x} + \mathbf{z}$ , with  $\mathbf{z} \sim \mathcal{N}(0, \sigma_z^2 / n\mathbf{I})$  and  $\sigma_z / \sqrt{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  $\ell_2$  and Jacobian regularization:
$\ell_2$  regularization. Implemented as weight-decay in PyTorch's SGD optimizer to minimize
$$
\hat {W} _ {\lambda} (\boldsymbol {\theta}) = \sum_ {i = 1} ^ {N} \| f _ {\boldsymbol {\theta}} (\mathbf {y} _ {i}) - \mathbf {x} _ {i} \| _ {2} ^ {2} + \lambda \| \boldsymbol {\theta} \| _ {2} ^ {2}.
$$
Jacobian regularization. We train networks with the Jacobian regularized empirical risk
$$
\widehat {\operatorname {J a c}} _ {\lambda} (\boldsymbol {\theta}) = \sum_ {i = 1} ^ {N} \| f _ {\boldsymbol {\theta}} (\mathbf {y} _ {i}) - \mathbf {x} _ {i} \| _ {2} ^ {2} + \lambda \| \mathbf {J} _ {\mathbf {y} _ {i}} \| _ {F} ^ {2},
$$
where  $\mathbf{J}_{\mathbf{y}_i}$  is the Jacobian of the network  $f_{\theta}$  with respect to it's input (not it's parameters) at  $\mathbf{y}_i$ . This regularization can be viewed as an approximation of the jittering risk, as described in
Figure 14: Pixel-wise robust (left) and standard risk (right) of U-nets trained via jittering,  $\ell_2$  and Jacobian regularization on the task of denoising grayscale images (at noise level  $\sigma_z / \sqrt{n} = 0.2$ ). The plot shows that  $\ell_2$  regularization yields less robust and accurate estimators compared to jittering. In contrast, Jacobian regularization obtains similarly robust estimators, but with slightly weaker performance in standard risk.
subsection F.1.Calculating the full Jacobian  $\mathbf{J}_{\mathbf{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,  $\| \mathbf{J}_{\mathbf{y}_i}\| _F^2$  with  $\left\| \mathbf{J}_{\mathbf{y}_i}^T\mathbf{w}\right\| _2^2$  , where  $\mathbf{w}\sim \mathcal{N}(0,\mathbf{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  $\hat{R}_{\epsilon} / n$ .
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  $\ell_2$  regularization yields worse performing estimators than jittering and Jacobian regularization. While for the linear subspace setting, adversarial training, Jacobian and  $\ell_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.
(b) Noisy Image
(c) Standard
(d) Robust Tr. (medium  $\epsilon$ )
(e) Jittering (medium  $\epsilon$ )
(f) Jacobian (medium  $\epsilon$ )
(g) Robust Tr. (large  $\epsilon$ )
(h) Jittering (large  $\epsilon$ )
(i) Jacobian (large  $\epsilon$ )
Figure 15: Example reconstructions using U-Nets, trained with different robustness-enhancing schemes at noise level  $\sigma_z / \sqrt{n} = 0.2$ . The second row depicts the results of using neural networks trained with methods tuned on a medium perturbation level of  $\epsilon^2 / \sigma_c^2 = 0.03$ , whereas the third row shows results for a large perturbation level  $\epsilon^2 / \sigma_c^2 = 0.3$ . The plot shows that jittering yields smooth reconstructions, whereas adversarial training and Jacobian regularization yield less smooth reconstructions.