SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yuhao Mao, Mark Niklas Müller, Marc Fischer, Martin Vechev
{yuhao.mao, mark.mueller, marc.fischer, martin.vechev}@inf.ethz.ch
Training certifiably robust neural networks remains a notoriously hard problem. While adversarial training optimizes under-approximations of the worst-case loss, which leads to insufficient regularization for certification, sound certified training methods, optimize loose over-approximations, leading to over-regularization and poor (standard) accuracy. In this work, we propose TAPS, an (unsound) certified training method that combines IBP and PGD training to optimize more precise, although not necessarily sound, worst-case loss approximations, reducing over-regularization and increasing certified and standard accuracies. Empirically, TAPS achieves a new state-of-the-art in many settings, e.g., reaching a certified accuracy of  $22\%$  on TINYIMAGENet for  $\ell_{\infty}$ -perturbations with radius  $\epsilon = 1/255$ . We make our implementation and networks public at github.com/eth-sri/taps.
Adversarial robustness, i.e., a neural network's resilience to small input perturbations (Biggio et al., 2013; Szegedy et al., 2014), has established itself as an important research area.
Neural Network Certification can rigorously prove such robustness: While complete verification methods (Tjeng et al., 2019; Bunel et al., 2020; Zhang et al., 2022a; Ferrari et al., 2022) can decide every robustness property given enough (exponential) time, incomplete methods (Wong and Kolter, 2018; Singh et al., 2019a; Zhang et al., 2018) trade precision for scalability.
Adversarial training methods, such as PGD (Madry et al., 2018), aim to improve robustness by training with samples that are perturbed to approximately maximize the training loss. This can be seen as optimizing an under-approximation of the worst-case loss. While it empirically improves robustness significantly, it generally does not induce sufficient regularization for certification and has been shown to fail in the face of more powerful attacks (Tramér et al., 2020).
Certified Training methods, in contrast, optimize approximations of the worst-case loss, thus increasing certified accuracies at the cost of over-regularization that leads to reduced standard accuracies. In this work, we distinguish two certified training paradigms. Sound methods (Mirman et al., 2018; Gowal et al., 2018; Shi et al., 2021) compute sound over-approximations of the worst-case loss via bound propagation. The resulting approximation errors induce a strong (over-)regularization that makes certification easy but causes severely reduced standard accuracies. Interestingly, reducing these approximation errors by using more precise bound propagation methods, empirically, results in strictly worse performance, as they induce harder optimization problems (Jovanović et al., 2022). This gave rise to unsound methods (Balunovic and Vechev, 2020; Palma et al., 2022; Müller et al., 2022a), which aim to compute precise but not necessarily sound approximations of the worst-case loss, reducing (over)-regularization and resulting in networks that achieve higher standard and certified accuracies, but can be harder to certify. Recent advances in certification techniques, however, have made their certification practically feasible (Ferrari et al., 2022; Zhang et al., 2022a).
We illustrate this in Figure 1, where we compare certified training methods with regard to their worst-case loss approximation errors and the resulting trade-off between certified and standard accuracy. On the left, we show histograms of the worst-case loss approximation error over test set samples (see Section 4.2 for more details). Positive values (right of the y-axis) correspond to over- and negative values (left of the y-axis) to under-approximations. As expected, we observe that the sound IBP (Gowal et al., 2018) always yields over-approximations (positive values) while the unsound SABR (Müller et al., 2022a) yields a more precise (6-fold reduction in mean error) but unsound approximation of the worst-case loss. Comparing the resulting accuracies (right), we observe that this more precise approximation of the actual optimization objective, i.e., the true worst-case loss, by SABR  $(\diamond)$  yields both higher certified and standard accuracies than the over-approximation by IBP  $(\bullet)$ . Intuitively, reducing the over-regularization induced by a systematic underestimation of the network's robustness allows it to allocate more capacity to making accurate predictions.
The core challenge of effective certified training is, thus, to compute precise (small mean error and low variance) worst-case loss approximations that induce a well-behaved optimization problem.
This Work proposes Training via Adversarial Propagation through Subnetworks (TAPS), a novel (unsound) certified training method tackling this challenge, thereby increasing both certified and standard accuracies. Compared to SABR (the current state-of-the-art), TAPS enjoys a further 5-fold mean approximation error reduction and significantly reduced variance (Figure 1 left), leading to improved certified and natural accuracies (right). The key technical insight behind TAPS is to combine IBP and PGD training via a gradient connector, a novel mechanism that allows training the whole network jointly such that the over-approximation of IBP
Figure 1: Histograms of the worst-case loss approximation errors over the test set (left) for different training methods show that TAPS (our work) achieves the most precise approximations and highest certified accuracy (right). Results shown here are for a small CNN3.
and under-approximations of PGD cancel out. We demonstrate in an extensive empirical study that TAPS yields exceptionally tight worst-case loss approximations which allow it to improve on state-of-the-art results for MNIST, CIFAR-10, and TINYIMAGENET.
2 Background on Adversarial and Certified Training
Here, we provide the necessary background on adversarial and certified training. We consider a classifier  $F\colon \mathcal{X}\mapsto \mathcal{Y}$  parameterized by weights  $\pmb{\theta}$  and predicting a class  $y_{\mathrm{pred}}\coloneqq F(\pmb {x})\coloneqq \arg \max_{y\in \mathcal{Y}}f_y(x)$  for every input  $\pmb {x}\in \mathcal{X}\subseteq \mathbb{R}^d$  with label  $y\in \mathcal{V}\coloneqq \{1,\dots ,K\}$  where  $\pmb {f}\colon \mathcal{X}\mapsto \mathbb{R}^{|\mathcal{Y}|}$  is a neural network, assigning a numerical logit  $o_i\coloneqq f_i(\pmb {x})$  to each class  $i$ .
Adversarial Robustness We call a classifier adversarially robust on an  $\ell_p$ -norm ball  $\mathcal{B}_p(\boldsymbol{x},\epsilon)$  if it classifies all elements within the ball to the correct class, i.e.,  $F(\boldsymbol{x}^{\prime}) = y$  for all perturbed inputs  $\boldsymbol{x}^{\prime}\in \mathcal{B}_{p}(\boldsymbol{x},\epsilon)$ . In this work, we focus on  $\ell_{\infty}$ -robustness with  $\mathcal{B}_{\infty}(\boldsymbol{x},\epsilon)\coloneqq \{\boldsymbol{x}^{\prime}\mid \| \boldsymbol{x}^{\prime} - \boldsymbol {x}\|_{\infty}\leq \epsilon \}$  and thus drop the subscript  $\infty$ .
Neural Network Certification is used to formally prove robustness properties of a neural network, i.e., that all inputs in the region  $\mathcal{B}(\pmb{x},\epsilon)$  yield the correct classification. We call samples  $\pmb{x}$  where this is successful, certifiably robust and denote the portion of such samples as certified accuracy, forming a lower bound to the true robustness of the analyzed network.
Interval bound propagation (IBP) (Mirman et al., 2018; Gowal et al., 2018) is a particularly simple yet effective certification method. Conceptually, it computes an over-approximation of a network's reachable set by propagating the input region  $\mathcal{B}(\boldsymbol{x},\epsilon)$  through the network, before checking whether all reachable outputs yield the correct classification. This is done by, first, over-approximating the input region  $\mathcal{B}(\boldsymbol{x},\epsilon)$  as a Box  $[\underline{\boldsymbol{x}}^0,\overline{\boldsymbol{x}}^0]$  (each dimension is described as an interval), centered at  $c^0 = x$  and with radius  $\delta^0 = \epsilon$ , such that we have the  $i^{\mathrm{th}}$  dimension of the input  $x_i^0\in [c_i^0 -\delta_i^0,c_i^0 +\delta_i^0 ]$ . We then
propagate it through the network layer-by-layer (for more details, see (Mirman et al., 2018; Gowal et al., 2018)), until we obtain upper and lower bounds  $[\underline{o}^{\Delta},\overline{o}^{\Delta}]$  on the logit differences  $o^{\Delta}\coloneqq o - o_{y}\mathbf{1}$ . If we can now show dimensionwise that  $\overline{o}^{\Delta} < 0$  (except for  $\overline{o}_y^\Delta = 0$ ), this proves robustness. Note that this is equivalent to showing that the maximum margin loss  $\mathcal{L}_{\mathrm{MA}}(\pmb{x}',y)\coloneqq \max_{i\neq y}\overline{o}_i^\Delta$  is less than 0 for all perturbed inputs  $\pmb{x}^{\prime}\in \mathcal{B}(\pmb {x},\epsilon)$ .
Training for Robustness aims to find a model parametrization  $\theta$  that minimizes the expected worst-case loss for some loss-function  $\mathcal{L}$ :
$$
\boldsymbol {\theta} = \underset {\boldsymbol {\theta}} {\arg \min } \mathbb {E} _ {\boldsymbol {x}, y} \left[ \max  _ {\boldsymbol {x} ^ {\prime} \in \mathcal {B} (\boldsymbol {x}, \epsilon)} \mathcal {L} \left(\boldsymbol {x} ^ {\prime}, y\right) \right]. \tag {1}
$$
As the inner maximization objective in Equation (1) can generally not be solved exactly, it is often under- or over-approximated, giving rise to adversarial and certified training, respectively.
Adversarial Training optimizes a lower bound on the inner maximization problem in Equation (1) by training the network with concrete samples  $\pmb{x}^{\prime} \in \mathcal{B}(\pmb{x},\epsilon)$  that (approximately) maximize the loss function. A well-established method for this is Projected Gradient Descent (PGD) training (Madry et al., 2018) which uses the Cross-Entropy loss  $\mathcal{L}_{\mathrm{CE}}(\pmb{x},y) \coloneqq \ln \left(1 + \sum_{i \neq y} \exp(f_i(\pmb{x}) - f_y(\pmb{x}))\right)$ . Starting from a random initialization point  $\hat{\pmb{x}}_0 \in \mathcal{B}(\pmb{x},\epsilon)$ , it performs  $N$  update steps
$$
\hat {\boldsymbol {x}} _ {n + 1} = \Pi_ {\mathcal {B} (\boldsymbol {x}, \epsilon)} \hat {\boldsymbol {x}} _ {n} + \eta \operatorname {s i g n} \left(\nabla_ {\hat {\boldsymbol {x}} _ {n}} \mathcal {L} \left(\hat {\boldsymbol {x}} _ {n}, y\right)\right)
$$
with step size  $\eta$  and projection operator  $\Pi$ . Networks trained this way typically exhibit good empirical robustness but remain hard to formally certify and vulnerable to stronger or different attacks (Tramèr et al., 2020; Croce and Hein, 2020).
Certified Training, in contrast, is used to train certifiably robust networks. In this work, we distinguish two classes of such methods: while sound methods optimize a sound upper bound of the inner maximization objective in Equation (1), unsound methods sacrifice soundness to use an (in expectation) more precise approximation. Methods in both paradigms are often based on evaluating the cross-entropy loss  $\mathcal{L}_{\mathrm{CE}}$  with upper bounds on the logit differences  $\overline{o}^{\Delta}$ .
IBP (a sound method) uses sound Box bounds on the logit differences, yielding
$$
\mathcal {L} _ {\mathrm {I B P}} (x, y, \epsilon) := \ln \left(1 + \sum_ {i \neq y} \exp \left(\bar {o} _ {i} ^ {\Delta}\right)\right). \tag {2}
$$
SABR (an unsound method) (Müller et al., 2022a), in contrast, first searches for an adversarial example  $\pmb{x}' \in \mathcal{B}(\pmb{x}', \epsilon - \tau)$  and then computes Box-bounds only for a small region  $\mathcal{B}(\pmb{x}', \tau) \subset \mathcal{B}(\pmb{x}, \epsilon)$  (with  $\tau < \epsilon$ ) around this adversarial example  $\pmb{x}'$  instead of the original input  $\pmb{x}$
$$
\mathcal {L} _ {\mathrm {S A B R}} := \max  _ {\boldsymbol {x} ^ {\prime} \in \mathcal {B} \left(\boldsymbol {x} ^ {\prime}, \epsilon - \tau\right)} \mathcal {L} _ {\mathrm {I B P}} \left(\boldsymbol {x} ^ {\prime}, y, \tau\right). \tag {3}
$$
This generally yields a more precise (although not sound) worst-case loss approximation, thereby reducing over-regularization and improving both standard and certified accuracy.
3 Precise Worst-Case Loss Approximation
In this section, we first introduce TAPS, a novel certified training method combining IBP and PGD training to obtain more precise worst-case loss estimates, before showing that this approach is orthogonal and complementary to current state-of-the-art methods.
3.1 TAPS - Combining IBP and PGD
The key insight behind TAPS is that adversarial training with PGD and certified training with IBP complement each other perfectly: (i) both yield well-behaved optimization problems, as witnessed by their empirical success, and (ii) we can combine them such that the over-approximation errors incurred during IBP are compensated by the under-approximations of PGD. TAPS harnesses this as follows: For every sample, we first propagate the input region part-way through the network using
Figure 2: Overview of TAPS training. First, forward propagation  $(\rightarrow)$  of a region  $B(\pmb{x},\epsilon)$  (□, left) around an input  $\pmb{x}$  (×) through the feature extractor  $\pmb{f}_{E}$  yields the exact reachable set (□, middle) and its IBP approximation  $[\underline{z},\overline{z}]$  (□, middle) in the embedding space. Further IBP propagation through the classifier  $\pmb{f}_{C}$  would yield an imprecise box approximation (□, right) of the reachable set (□, right). Instead, TAPS conducts an adversarial attack  $(\rightarrow)$  in the embedding space IBP approximation (□) yielding an under-approximation ( $\mathbf{\Pi}^{\mathrm{red}}$ ) of its reachable set (□, right). We illustrate the points realizing the worst-case loss in every output region with  $\bullet$  and enable back-propagation ( $\cdot \rightarrow$ ) through the adversarial attack by introducing the gradient connector (discussed in Section 3.2).
IBP and then conduct PGD training within the thus obtained Box approximation. The key technical challenge with this approach lies in connecting the gradients of the two propagation methods and thus enabling joint training of the corresponding network portions.
We now explain TAPS in more detail along the illustration in Figure 2. We first partition a neural network  $\pmb{f}$  with weights  $\pmb{\theta}$  into a feature extractor  $\pmb{f}_{E}$  and a classifier  $\pmb{f}_{C}$  with parameters  $\pmb{\theta}_{E}$  and  $\pmb{\theta}_{C}$ , respectively, such that we have  $\pmb{f}_{\theta} = \pmb{f}_{C} \circ \pmb{f}_{E}$  and  $\pmb{\theta} = \pmb{\theta}_{E} \cup \pmb{\theta}_{C}$ . We refer to the output space of the feature extractor as the embedding space. Given an input sample  $\pmb{x}$  (illustrated as  $\times$  in Figure 2) and a corresponding input region  $\mathcal{B}(\pmb{x}, \epsilon)$  (in the input panel), training proceeds as follows: During the forward pass (black dashed arrows  $\rightarrow$ ), we first use IBP to compute a BOX over-approximation  $[\underline{z}, \overline{z}]$  (dashed box) of the feature extractor's exact reachable set (blue region), shown in the middle panel of Figure 2. Then, we conduct separate adversarial attacks  $(\rightarrow)$  within this region in the embedding space (□) to bound all output dimensions of the classifier. This yields latent adversarial examples  $\hat{z} \in [\underline{z}, \overline{z}]$ , defining the TAPS bounds  $\overline{o}_{\mathrm{TAPS}}^{\Delta}$  (dotted lines " in the output space) on the network's output. This way, the under-approximation of the classifier via PGD, partially compensates the over-approximation of the feature extractor via IBP. Full IBP propagation, in contrast, continues to exponentially accumulate approximation errors (Müller et al., 2022a; Shi et al., 2021), yielding the much larger dashed box [□]. We now compute the TAPS loss  $\mathcal{L}_{\mathrm{TAPS}}$  analogously to  $\mathcal{L}_{\mathrm{IBP}}$  (Equation (2)) by plugging the TAPS bound estimate  $\overline{o}_{\mathrm{TAPS}}^{\Delta}$  into the Cross-Entropy loss. Comparing the resulting losses (illustrated as  $\bullet$  and growing towards the top right), we see that while the TAPS bounds are not necessarily sound, they yield a much better approximation of the true worst-case loss.
During the backward pass (orange dotted arrows  $\rightarrow$  in Figure 2), we compute the gradients w.r.t. the classifier's parameters  $\theta_C$  and the latent adversarial examples  $\hat{z}$  (classifier input) as usual. However, to compute the gradients w.r.t. the feature extractor's parameters  $\theta_F$ , we have to compute (pseudo) gradients of the latent adversarial examples  $\hat{z}$  w.r.t. the box bounds  $\underline{z}$  and  $\overline{z}$ . As these gradients are not well defined, we introduce the gradient connector, discussed next, as an interface between the feature extractor and classifier, imposing such pseudo gradients. This allows us to train  $f_E$  and  $f_C$  jointly, leading to a feature extractor that minimizes approximation errors and a classifier that is resilient to the spurious points included in the remaining approximation errors.
3.2 Gradient Connector
The key function of the gradient connector is to enable gradient computation through the adversarial example search in the embedding space. Using the chain rule, this only requires us to define the (pseudo) gradients  $\frac{\partial\hat{z}}{\partial\underline{z}}$  and  $\frac{\partial\hat{z}}{\partial\overline{z}}$  of the latent adversarial examples  $\hat{z}$  w.r.t. the box bounds  $\underline{z}$  and  $\overline{z}$  on the feature extractor's outputs. Below, we will focus on the  $i^{\mathrm{th}}$  dimension of the lower box bound  $\underline{z}_i$  and note that all other dimensions and the upper bounds follow analogously.
As the latent adversarial examples can be seen as multivariate functions in the box bounds, we obtain the general form  $\frac{d\mathcal{L}}{d\bar{z}_i} = \sum_j\frac{d\mathcal{L}}{d\hat{z}_j}\frac{\partial\hat{z}_j}{\partial\underline{z}_i}$ , depending on all dimensions of the latent adversarial example.
We now consider a single PGD step and observe that bounds in the  $i^{\mathrm{th}}$  dimension have no impact on the  $j^{\mathrm{th}}$  coordinate of the resulting adversarial example as they impact neither the gradient sign nor the projection in this dimension, as Box bounds are axis parallel. We thus assume independence of the  $j^{\mathrm{th}}$  dimension of the latent adversarial example  $\hat{z}_j$  from the bounds in the  $i^{\mathrm{th}}$  dimension  $\underline{z}_i$  and  $\overline{z}_i$  (for  $i\neq j$ ), which holds rigorously (up to initialization) for a single step attack and constitutes a mild assumption for multi-step attacks. Therefore, we have  $\frac{\partial\hat{z}_j}{\partial\underline{z}_i} = 0$  for  $i\neq j$  and obtain  $\frac{d\mathcal{L}}{d\underline{z}_i} = \frac{d\mathcal{L}}{d\hat{z}_i}\frac{\partial\hat{z}_i}{\partial\underline{z}_i}$ , leaving only  $\frac{\partial\hat{z}_i}{\partial\underline{z}_i}$  for us to define.
The most natural gradient connector is the binary connector, i.e., set  $\frac{\partial\hat{z}_i}{\partial\underline{z}_i} = 1$  when  $\hat{z}_i = \underline{z}_i$  and 0 otherwise, as it is a valid sub-gradient for the projection operation in PGD. However, the latent adversarial input often does not lie on a corner (extremal vertex) of the Box approximation, leading to sparse gradients and thus a less well-behaved optimization problem. More importantly, the binary connector is very sensitive to the distance between (local) loss extrema and the box boundary and thus inherently ill-suited to gradient-based optimization. For example, a local extremum at  $\hat{z}_i$  would induce  $\frac{\partial\hat{z}_i}{\partial\underline{z}_i} = 1$  in the box  $[\hat{z}_i,0]$ , but  $\frac{\partial\hat{z}_i}{\partial\underline{z}_i} = 0$  for  $[\hat{z}_i - \epsilon ,0]$ , even for arbitrarily small  $\epsilon$ .
To alleviate both of these problems, we consider a linear connector, i.e., set  $\frac{\partial\hat{z}_i}{\partial\underline{z}_i} = \frac{\overline{z}_i - \hat{z}_i}{\overline{z}_i - \underline{z}_i}$ . However, even when our latent adversarial example is very close to one bound, the linear connector would induce nonzero gradients w.r.t. to the opposite bound. To remedy this undesirable behavior, we propose the rectified linear connector, setting  $\frac{\partial\hat{z}_i}{\partial\underline{z}_i} = \max (0,1 - \frac{\hat{z}_i - z_i}{c(\overline{z}_i - z_i)})$  where  $c\in [0,1]$  is a constant (visualized in Figure 3 for  $c = 0.3$ ). Observe that it recovers the binary connector
Figure 3: Gradient connector visualization.
for  $c = 0$  and the linear connector for  $c = 1$ . To prevent gradient sparsity ( $c \leq 0.5$ ) while avoiding the above-mentioned counterintuitive gradient connections ( $c \geq 0.5$ ), we set  $c = 0.5$  unless indicated otherwise. When the upper and lower bounds are identical in the  $i^{\text{th}}$  dimension, PGD turns into an identity function. Therefore, we set both gradients to  $\frac{\partial \hat{z}_i}{\partial \bar{z}_i} = \frac{\partial \hat{z}_i}{\partial \bar{z}_i} = 0.5$  turning the gradient connector into an identity function for the backward pass.
3.3 TAPS Loss & Multi-estimator PGD
The standard PGD attack, used in adversarial training, henceforth called single-estimator PGD, is based on maximizing the Cross-Entropy loss  $\mathcal{L}_{\mathrm{CE}}$  of a single input. In the context of TAPS, this results in the overall loss
$$
\mathcal {L} _ {\mathrm {T A P S}} ^ {\text {s i n g l e}} (\boldsymbol {x}, y, \epsilon) = \max  _ {\hat {\boldsymbol {z}} \in [ \underline {{\boldsymbol {z}}}, \overline {{\boldsymbol {z}}} ]} \ln \left(1 + \sum_ {i \neq y} \exp (f _ {C} (\hat {\boldsymbol {z}}) _ {i} - f _ {C} (\hat {\boldsymbol {z}}) _ {y})\right),
$$
where the embedding space bounding box  $[\underline{z},\overline{z} ]$  is obtained via IBP. However, this loss is not necessarily well aligned with adversarial robustness. Consider the example illustrated in Figure 4, where only points in the lower-left quadrant are classified correctly (i.e.,  $o_i^\Delta \coloneqq o_i - o_y < 0$ ). We compute the latent adversarial example  $\hat{z}$  by conducting a standard adversarial attack on the Cross-Entropy loss over the reachable set (optimally for illustration purposes) and observe that the
corresponding output  $f(\hat{z})(\times)$  is classified correctly. However, if we instead use the logit differences  $o_1^\Delta$  and  $o_2^\Delta$  as attack objectives, we obtain two misclassified points  $(\times)$ . Combining their dimension-wise worst-case bounds  $(\dots)$ , we obtain the point  $\bullet$ , which realizes the maximum loss over an optimal box approximation of the reachable set. As the correct classification of this point (when computed exactly) directly corresponds to true robustness, we propose the multi-estimator PGD variant of  $\mathcal{L}_{\mathrm{TAPS}}$ , which estimates the upper bounds on the logit differences  $o_i^\Delta$  using separate samples and then computes the loss function using the per-dimension worst-cases as:
$$
\mathcal {L} _ {\mathrm {T A P S}} (\boldsymbol {x}, y, \epsilon) = \ln \Big (1 + \sum_ {i \neq y} \exp \Big (\max  _ {\hat {\boldsymbol {z}} \in [ \underline {{\boldsymbol {z}}}, \overline {{\boldsymbol {z}}} ]} f _ {C} (\hat {\boldsymbol {z}}) _ {i} - f _ {C} (\hat {\boldsymbol {z}}) _ {y} \Big) \Big).
$$
Figure 4: Illustration of the bounds on  $o_i^{\Delta} \coloneqq o_i - o_t$  obtained via single estimator  $(-)$  and multi-estimator  $(\uparrow)$  PGD and the points maximizing the corresponding losses:  $\times$  for  $\mathcal{L}_{\mathrm{TAPS}}^{\mathrm{single}}$  and  $\bullet$  for  $\mathcal{L}_{\mathrm{TAPS}}$ .
3.4 Training Objective & Regularization
While complete certification methods can decide any robustness property, this requires exponential time. Therefore, networks should not only be robust but also certifiable. Thus, we propose to combine the IBP loss for easy-to-learn and certify samples with the TAPS loss for harder samples as follows: