SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Suraj Srinivas*  
Harvard University  
Cambridge, MA  
ssrinivas@seas.harvard.edu
Sebastian Bordt*  
University of Tübingen, Tübingen AI Center  
Tübingen, Germany  
sebastian.bordt@uni-tuebingen.de
Himabindu Lakkaraju  
Harvard University  
Cambridge, MA  
hlakkaraju@hbs.edu
One of the remarkable properties of robust computer vision models is that their input-gradients are often aligned with human perception, referred to in the literature as perceptually-aligned gradients (PAGs). Despite only being trained for classification, PAGs cause robust models to have rudimentary generative capabilities, including image generation, denoising, and in-painting. However, the underlying mechanisms behind these phenomena remain unknown. In this work, we provide a first explanation of PAGs via off-manifold robustness, which states that models must be more robust off- the data manifold than they are on-manifold. We first demonstrate theoretically that off-manifold robustness leads input gradients to lie approximately on the data manifold, explaining their perceptual alignment. We then show that Bayes optimal models satisfy off-manifold robustness, and confirm the same empirically for robust models trained via gradient norm regularization, randomized smoothing, and adversarial training with projected gradient descent. Quantifying the perceptual alignment of model gradients via their similarity with the gradients of generative models, we show that off-manifold robustness correlates well with perceptual alignment. Finally, based on the levels of on- and off-manifold robustness, we identify three different regimes of robustness that affect both perceptual alignment and model accuracy: weak robustness, bayes-aligned robustness, and excessive robustness. Code is available at https://github.com/tml-tuebingen/pags.
An important desideratum for machine learning models is robustness, which requires that models be insensitive to small amounts of noise added to the input. In particular, adversarial robustness requires models to be insensitive to adversarially chosen perturbations of the input. Tsipras et al. [1] first observed an unexpected benefit of such models, namely that their input-gradients were "significantly more human-aligned" (see Figure 1 for examples of perceptually-aligned gradients). Santurkar et al. [2] built on this observation to show that robust models could be used to perform rudimentary image synthesis - an unexpected capability of models trained in a purely discriminative manner. Subsequent works have made use of the perceptual alignment of robust model gradients to improve zero-shot
Figure 1: A demonstration of the perceptual alignment phenomenon. The input-gradients of robust classifiers ("robust gradient") are perceptually similar to the score of diffusion models [10], while being qualitatively distinct from input-gradients of standard models ("standard gradient"). Best viewed in digital format.
object localization [3], perform conditional image synthesis [4], and improve classifier robustness [5, 6]. Kaur et al. [7] coined the term perceptually-aligned gradients (PAGs), and showed that it occurs not just with adversarial training [8], but also with randomized smoothed models [9]. Recently, Ganz et al. [6] showed that the relationship between robustness and perceptual alignment can also work in the opposite direction: approximately enforcing perceptual alignment of input gradients can increase model robustness.
Despite these advances, the underlying mechanisms behind the phenomenon of perceptually aligned gradients in robust models are still unclear. Adding to the confusion, prior works have used the same term, PAGs, to refer to slightly different phenomena. We ground our discussion by first identifying variations of the same underlying phenomenon.
Phenomenon 1 (Perceptual Alignment). The gradients of robust models highlight perceptually relevant features [1], and highlight discriminative input regions while ignoring distractors [11].
Phenomenon 2 (Generative Capabilities). Robust models have rudimentary image synthesis capabilities, where samples are generated by maximizing a specific output class by iteratively following the direction dictated by the input-gradients [2].
Phenomenon 3 (Smoothgrad Interpretability). Smoothgrad visualizations (which are gradients of randomized smooth models) tend to be visually sharper than the gradients of standard models [12].
While Kaur et al. [7] refer to Phenomenon 2 as PAGs, Ganz et al. [6] use this term to refer to Phenomenon 1. Since they are all related to the gradients of robust models, we choose to collectively refer to all of these phenomena as PAGs.
While the various phenomena arising from PAGs are now well-documented, there is little to no work that attempts to explain the underlying mechanism. Progress on this problem has been hard to achieve because PAGs have been described via purely qualitative criteria, and it has been unclear how to make these statements quantitative. In this work, we address these gaps and make one of the first attempts at explaining the mechanisms behind PAGs. Crucially, we ground the discussion by attributing PAGs to gradients lying on a manifold, based on an analysis of Bayes optimal classifiers. We make the following contributions:
1. We establish the first-known theoretical connections between Bayes optimal predictors and the perceptual alignment of classifiers via off-manifold robustness. We also identify the manifold w.r.t. which this holds, calling it the signal manifold.
2. We experimentally verify that models trained with gradient-norm regularization, noise augmentation and randomized smoothing all exhibit off-manifold robustness.
3. We quantify the perceptual alignment of gradients by their perceptual similarity with the score of the data distribution, and find that this measure correlates well with off-manifold robustness.
4. We identify three regimes of robustness: weak robustness, bayes-aligned robustness and excessive robustness, that differently affect both perceptual alignment and model accuracy.
Robust training of neural networks Prior works have considered two broad classes of model robustness: adversarial robustness and robustness to normal noise. Adversarial robustness is achieved by training with adversarial perturbations generated using project gradient descent [8], or by randomized smoothing [9], which achieves certified robustness to adversarial attacks by locally averaging with normal noise. Robustness to normal noise is achieved by explicitly training with noise, a technique that is equivalent to Tikhonov regularization for linear models [13]. For non-linear models, gradient-norm regularization [14] is equivalent to training with normal noise under the limit of training with infinitely many noise samples [15]. In this paper, we make use of all of these robust training approaches and investigate their relationship to PAGs.
Gradient-based model explanations Several popular post hoc explanation methods[16, 17, 12] estimate feature importances by computing gradients of the output with respect to input features and aggregating them over local neighborhoods [18]. However, the visual quality criterion used to evaluate these explanations has given rise to methods that produce visually striking attribution maps, while being independent of model behavior [19]. While [20] attribute visual quality to implicit score-based generative modeling of the data distribution, [21] propose that it depends on explanations lying on the data manifold. While prior works have attributed visual quality to generative modeling of data distribution or explanations lying on data manifold, our work demonstrates for the first time that (1) off-manifold robustness is a crucial factor, and (2) it is not the data manifold / distribution, rather the signal manifold / distribution (defined in Section 3.2) that is the critical factor in explaining the phenomenon of PAGs.
3 Explaining Perceptually-Aligned Gradients
Our goal in this section is to understand the mechanisms behind PAGs. We first consider PAGs as lying on a low-dimensional manifold, and show theoretically that such on-manifold alignment of gradients is equivalent to off-manifold robustness of the model. We then argue that Bayes optimal models achieve both off-manifold robustness and on-manifold gradient alignment. In doing so, we introduce the distinction between the data and the signal manifold, which is key to understanding the manifold structure of PAGs. Finally, we present arguments for why empirical robust models also exhibit off-manifold robustness.
Notation Throughout this paper, we consider the task of image classification with inputs  $\mathbf{x} \in \mathbb{R}^d$  where  $\mathbf{x} \sim \mathcal{X}$  and  $y \in [1,2,\dots C]$  with  $C$ -classes. We consider deep neural networks  $f: \mathbb{R}^d \to \triangle^{C-1}$  which map inputs  $\mathbf{x}$  onto a  $C$ -class probability simplex. We assume that the input data  $\mathbf{x}$  lies on a  $k$ -dimensional manifold in the  $d$ -dimensional ambient space. Formally, a  $k$ -dimensional differential manifold  $\mathcal{M} \subset \mathbb{R}^d$  is locally Euclidean in  $\mathbb{R}^k$ . At every point  $\mathbf{x} \in \mathcal{M}$ , we can define a projection matrix  $\mathbb{P}_x \in \mathbb{R}^{d \times d}$  that projects points onto the  $k$ -dimensional tangent space at  $\mathbf{x}$ . We denote  $\mathbb{P}_x^\perp = \mathbb{I} - \mathbb{P}_x$  as the projection matrix to the subspace orthogonal to the tangent space.
3.1 Off-Manifold Robustness  $\Leftrightarrow$  On-Manifold Gradient Alignment
We now show via geometric arguments that off-manifold robustness of models and on-manifold alignment of gradients are identical. We begin by discussing definitions of on- and off-manifold noise. Consider a point  $\mathbf{x}$  on manifold  $\mathcal{M}$ , and a noise vector  $\mathbf{u}$ . Then,  $\mathbf{u}$  is off-manifold noise, if  $\mathbb{P}_x(\mathbf{x} + \mathbf{u}) = \mathbf{x}$ , which we denote  $\mathbf{u} \coloneqq \mathbf{u}_{\mathrm{off}}$ , and  $\mathbf{u}$  is on-manifold noise, if  $\mathbb{P}_x(\mathbf{x} + \mathbf{u}) = \mathbf{x} + \mathbf{u}$ , which we denote  $\mathbf{u} \coloneqq \mathbf{u}_{\mathrm{on}}$ . In other words, if the noise vector lies on the tangent space then it is on-manifold, otherwise it is off-manifold. Given this definition, we can define relative off-manifold robustness. For simplicity, we consider a scalar valued function, which can correspond to one of the  $C$  output classes of the model.
Definition 1. (Relative off-manifold robustness) A model  $f: \mathbb{R}^d \to \mathbb{R}$  is  $\rho_1$ -off-manifold robust, wrt some noise distribution  $\mathbf{u} \sim U$  if
$$
\frac {\mathbb {E} _ {\mathbf {u} _ {o f f}} \left(f (\mathbf {x} + \mathbf {u} _ {o f f}) - f (\mathbf {x})\right) ^ {2}}{\mathbb {E} _ {\mathbf {u}} (f (\mathbf {x} + \mathbf {u}) - f (\mathbf {x})) ^ {2}} \leq \rho_ {1} \quad \text {w h e r e} \quad \mathbf {u} _ {o f f} = \mathbb {P} _ {x} ^ {\perp} (\mathbf {u})
$$
While this definition states that the model is more robust off-manifold than it is overall, it can equivalently be interpreted as being more robust off-manifold than on-manifold (with a factor  $\frac{\rho_1}{1 - \rho_1}$ ). Let us now define on-manifold gradient alignment.
Definition 2. (On-manifold gradient alignment) A model  $f: \mathbb{R}^d \to \mathbb{R}$  has  $\rho_2$ -on-manifold aligned gradients if
$$
\frac {\| \nabla_ {\mathbf {x}} f (\mathbf {x}) - \nabla_ {\mathbf {x}} ^ {o n} f (\mathbf {x}) \| ^ {2}}{\| \nabla_ {\mathbf {x}} f (\mathbf {x}) \| ^ {2}} \leq \rho_ {2} \quad w h e r e \quad \nabla_ {\mathbf {x}} ^ {o n} f (\mathbf {x}) = \mathbb {P} _ {x} \nabla_ {\mathbf {x}} f (\mathbf {x})
$$
This definition captures the idea that the difference between the gradients and its on-manifold projection is small. If  $\rho_{2} = 0$ , then the gradients exactly lie on the tangent space. We are now ready to state the following theorem.
Theorem 1 (Equivalence between off-manifold robustness and on-manifold alignment). A function  $f: \mathbb{R}^d \to \mathbb{R}$  exhibits on-manifold gradient alignment if and only if it is off-manifold robust wrt normal noise  $\mathbf{u} \sim \mathcal{N}(0, \sigma^2)$  for  $\sigma \rightarrow 0$  (with  $\rho_1 = \rho_2$ ).
Proof Idea. Using a Taylor series expansion, we can re-write the off-manifold robustness objective in terms of gradients. We further use linear algebraic identities to simplify that yield the on-manifold gradient alignment objective. The full proof is given in the supplementary material.  $\square$
When the noise is small, on-manifold gradient alignment and off-manifold robustness are identical. To extend this to larger noise levels, we would need to make assumptions regarding the curvature of the data manifold. Broadly speaking, the less curved the underlying manifold is (the closer it is to being linear), the larger the off-manifold noise we can add, without it intersecting at another point on the manifold. In practice, we expect image manifolds to be quite smooth, indicating relative off-manifold robustness to larger noise levels [22]. In this next part, we will specify the properties of the specific manifold w.r.t. which these notions hold.
3.2 Connecting Bayes Optimal Classifiers and Off-Manifold Robustness
In this subsection, we aim to understand the manifold structure characterizing the gradients of Bayes optimal models. We proceed by recalling the concept of Bayes optimality.
Bayes optimal classifiers. If we perfectly knew the data generating distributions for all the classes, i.e.,  $p(\mathbf{x} \mid y = i)$  for all  $C$  classes, we could write the Bayes optimal classifier as  $p(y = i \mid \mathbf{x}) = p(\mathbf{x} \mid y = i) / \sum_{j=1}^{C} p(\mathbf{x} \mid y = j)$ . This is an oracle classifier that represents the "best possible" model one can create from data with perfect knowledge of the data generating process. Given our assumption that the data lies on a low-dimensional data manifold, the Bayes optimal classifier is uniquely defined on the data manifold. However, outside this manifold, its behaviour is undefined as all of the class-conditional probabilities  $p(\mathbf{x} \mid y)$  are zero or undefined themselves. In order to link Bayes-optimality and robustness, which is inherently about the behavior of the classifier outside the data manifold, we introduce a ground-truth perturbed data distribution that is also defined outside the data manifold. While we might consider many perturbed generative models [23]. Here, the stochasticity ensures that the data-generating probabilities are defined everywhere, not just on the data manifold. This approach allows us to estimate the data manifold, assuming that the autoencoder has a bottleneck layer with  $k$  features for a  $d$ -dimensional ambient space [22].
Figure 2: An illustration of a signal-distractor decomposition for a bird classification task. The signal represents the discriminative parts of the input, while the distractor represents the non-discriminative parts.
Distinguishing Between the Data and the Signal Manifold. Given a classification problem, one can often decompose the inputs into a signal component and a distractor component. For intuition, consider the binary task of classifying cats and dogs. Given an oracle data generating distribution of cats and dogs in diverse backgrounds, the label must be statistically independent of the background and depend purely on the object (cat or dog) in the image. In other words, there exist no spurious correlations between the background and the output label. In such a case, we can call the object the signal and the background the distractor. Formally, for every input  $\mathbf{x}$  there exists a binary mask  $\mathbf{m}(\mathbf{x}) \in \{0,1\}^d$  such that the signal is given by  $s(\mathbf{x}) = \mathbf{x} \odot \mathbf{m}(\mathbf{x})$  and the distractor is given by  $d(\mathbf{x}) = \mathbf{x} \odot (1 - \mathbf{m}(\mathbf{x}))$ . The signal and distractor components are orthogonal to one another ( $s(\mathbf{x})^\top d(\mathbf{x}) = 0$ ), and we can decompose any input in this manner ( $\mathbf{x} = s(\mathbf{x}) + d(\mathbf{x})$ ). We provide an illustrative figure in Figure 2. Using this, we can now define the signal-distractor distributions.
Definition 3. Given a data distribution  $p(\mathbf{x} \mid y)$  for  $y \in [1, C]$ , we have masking functions  $\mathbf{m}(\mathbf{x})$  such that the resulting distribution  $p(\mathbf{x} \odot (1 - \mathbf{m}) \mid y) = p(\mathbf{x} \odot (1 - \mathbf{m}))$  is statistically independent of  $y$ . The sparsest such masking function (such that  $\mathbf{m}^* = \arg \min_{\mathbf{m}} \mathbb{E}_{\mathbf{x} \sim p(\mathbf{x} \mid y)} \| \mathbf{m}(\mathbf{x}) \|_0$ ), yields a corresponding distribution  $p(\mathbf{x} \odot (1 - \mathbf{m}^*(\mathbf{x})))$ , which is the distractor distribution, and its counterpart  $p(\mathbf{x} \odot \mathbf{m}^*(\mathbf{x}) \mid y)$  the signal distribution.
While it is possible to define signals and distractors more generally as any linear projections of the data, in this work, we restrict ourselves to masking to relate to feature attributions [11]. While the subject of finding such optimal masks is the topic of feature attribution [24], in this discussion, we shall assume that the optimal masks  $\mathbf{m}^*$  and the corresponding signal and distractor distributions are known. Similar to decomposing any point on the manifold, we can also decompose any vector on the tangent space on the data manifold into signal and distractor components using the optimal mask  $\mathbf{m}^*$ . In other words, we can write
$$
\nabla_ {\mathbf {x}} p (\mathbf {x} \mid y) = \underbrace {\nabla_ {\mathbf {x}} p (\mathbf {x} \mid y) \odot \mathbf {m} ^ {*} (\mathbf {x})} _ {\text {s i g n a l} s (\mathbf {x}) \text {(h a s i n f o r m a t i o n a b o u t} y)} + \underbrace {\nabla_ {\mathbf {x}} p (\mathbf {x} \mid y) \odot (1 - \mathbf {m} ^ {*} (\mathbf {x}))} _ {\text {d i s t r a c t o r} d (\mathbf {x}) \text {(i n d e p e n d e n t o f y)}}
$$
Finally, we note that this signal-distractor decomposition does not meaningfully exist for all classification problems in that it can be trivial with the entire data distribution being equal to the signal, with zero distractors. Two examples of such cases are: (1) ordinary MNIST classification has no distractors due to the simplicity of the task, as the entire digit is predictive of the true label and the background is zero. (2) multi-class classification with a large set of classes with diverse images also have no distractors due to the complexity of the task, as the background can already be correlated with class information. For example, if the dataset mixes natural images with deep space images, there is no single distractor distribution one can find via masking that is independent of the class label. This disadvantage arises partly because we aim to identify important pixels, whereas a more general definition based on identifying subspaces can avoid these issues. We leave the exploration of this more general class of methods to future work.
Given this definition of the signal and distractor distributions, we are ready to make the following theorem: the input-gradients of a Bayes optimal classifier lie on the signal manifold, as opposed to the general data manifold.
Proof Idea. We first show that the gradients of the Bayes optimal predictor lie in the linear span of the gradients of the gradients of the data distribution, thus lying on the tangent space of the data distribution. Upon making a signal-distractor decomposition of the gradients of the Bayes optimal classifier, we find that the distractor term is always zero, indicating that the gradients of the Bayes optimal classifier always lie on the signal manifold. The proof is given in the supplementary material.
The full proof is given in the Supplementary material. This theorem can be intuitively thought of as follows: the optimal classifier only needs to look at discriminative regions of the input in order to classify optimally. In other words, changing the input values at discriminative signal regions is likely to have a larger effect on model output than changing the inputs slightly at unimportant distractor regions, indicating that the gradients of the Bayes optimal classifier highlight the signal. Thus, the Bayes optimal classifier does not need to model the distractor; only the signal is sufficient. This fact inspires us to make the following hypothesis to help us ground the discussion on the perceptual alignment of gradients:
Hypothesis 1. The input gradients of classifiers are perceptually aligned if and only if they lie on the signal manifold, thus highlighting only discriminative parts of the input.
This indicates that Bayes optimal model's gradients are perceptually aligned. Recall that the signal component only provides information about discriminative components of the input, and thus, we expect this to connect to perceptual alignment from Phenomenon 1. Next, we ask whether the gradients of practical models, particularly robust models, are also off-manifold robust.
3.3 Connecting Robust Models and Off-Manifold Robustness
Previously, we saw that Bayes optimal classifiers have the property of relative off-manifold robustness. Here, we argue that off-manifold robustness also holds for robust models in practice. This is a nontrivial and a perhaps surprising claim: common robustness objectives such as adversarial training, gradient-norm regularization, etc are isotropic, meaning that they do not distinguish between on- and off-manifold directions.
Hypothesis 2. Robust models are off-manifold robust w.r.t. the signal manifold.
We provide two lines of argument in support for this hypothesis. Ultimately, however, our evidence is empirical and presented in the next section.
**Argument 1:** Combined robust objectives are non-isotropic. While robustness penalties itself are isotropic and do not prefer robustness in any direction, they are combined with the cross-entropy loss on data samples, which lie on the data manifold. Let us consider the example of gradient norm regularization. The objective is given by:
$$
\mathbb {E} _ {\mathbf {x}} \left(\ell (f (\mathbf {x}), y (\mathbf {x})) + \lambda \| \nabla_ {\mathbf {x}} f (\mathbf {x}) \| ^ {2}\right) = \mathbb {E} _ {\mathbf {x}} \left(\underbrace {\ell (f (\mathbf {x}) , y (\mathbf {x})) + \lambda \| \nabla_ {\mathbf {x}} ^ {\mathrm {o n}} f (\mathbf {x}) \| ^ {2}} _ {\text {o n - m a n i f o l d o b j e c t i v e}} + \lambda \underbrace {\| \nabla_ {\mathbf {x}} ^ {\mathrm {o f f}} f (\mathbf {x}) \| ^ {2}} _ {\text {o f f - m a n i f o l d o b j e c t i v e}}\right)
$$
Here we have used,  $\| \nabla_{\mathbf{x}}f(\mathbf{x})\| ^2 = \| \nabla_{\mathbf{x}}^{\mathrm{on}}f(\mathbf{x}) + \nabla_{\mathbf{x}}^{\mathrm{off}}f(\mathbf{x})\| ^2 = \| \nabla_{\mathbf{x}}^{\mathrm{on}}f(\mathbf{x})\| ^2 +\| \nabla_{\mathbf{x}}^{\mathrm{off}}f(\mathbf{x})\| ^2$  due to  $\nabla_{\mathbf{x}}^{\mathrm{on}}f(\mathbf{x})^{\top}\nabla_{\mathbf{x}}^{\mathrm{off}}f(\mathbf{x}) = 0$ , which is possible because the on-manifold and off-manifold parts of the gradient are orthogonal to each other. Assuming that we are able to decompose models into on-manifold and off-manifold parts, these two objectives apply to these decompositions independently. This argument states that in this robust training objective, there exists a trade-off between the cross-entropy loss and the on-manifold robustness term, whereas there is no such trade-off for the off-manifold term, indicating that it is much easier to minimize off-manifold robustness than on-manifold. This argument also makes the prediction that increased on-manifold robustness must be accompanied by higher train loss and decreased out-of-sample performance, and we will test this in the experiments section.
Theorem 2. The input-gradients of Bayes optimal classifiers lie on the signal manifold.