Efficient Training of Low-Curvature Neural Networks
Abstract
Standard deep neural networks often have excess non-linearity, making them susceptible to issues such as low adversarial robustness and gradient instability.
Common methods to address these downstream issues, such as adversarial training, are expensive and often sacrifice predictive accuracy.
In this work, we address the core issue of excess non-linearity via curvature, and demonstrate low-curvature neural networks (LCNNs) that obtain drastically lower curvature than standard models while exhibiting similar predictive performance.
This leads to improved robustness and stable gradients, at a fraction of the cost of standard adversarial training.
To achieve this, we decompose overall model curvature in terms of curvatures and slopes of its constituent layers.
To enable efficient curvature minimization of constituent layers, we introduce two novel architectural components: first, a non-linearity called centered-softplus that is a stable variant of the softplus non-linearity, and second, a Lipschitz-constrained batch normalization layer.
Our experiments show that LCNNs have lower curvature, more stable gradients and increased off-the-shelf adversarial robustness when compared to standard neural networks, all without affecting predictive performance.
Our approach is easy to use and can be readily incorporated into existing neural network architectures.
Code to implement our method and replicate our experiments is available at https://github.com/kylematoba/lcnn .
Introduction
The high degree of flexibility present in deep neural networks is critical to achieving good performance in complex tasks such as image classification, language modelling and generative modelling of images [1–3].
However, excessive flexibility is undesirable as this can lead to model underspecification [4] which results in unpredictable behaviour on out-of-domain inputs, such as vulnerability to adversarial examples.
Such under-specification can be avoided in principle via Occam’s razor, which requires training models that are as simple as possible for the task at hand, and not any more.
Motivated by this principle, in this work we focus on training neural network models without excess non-linearity in their input-output map (e.g.: see Fig 1b), such that predictive performance remains unaffected.
[Figure]
softplus( =1.0) softplus( =0.2)
c-softplus( =1.0) c-softplus( =0.2)
[Figure]
LCNN recovers highly regular decision boundaries in contrast to the standard NN
Comparison of softplus and centered-softplus non-linearities (defined in §3.3).
These behave similarly for large β values, and converge to linear maps for low β values.
However, softplus diverges while centeredsoftplus stays close to the origin.
Central to this work is a precise notion of curvature, which is a mathematical quantity that encodes the flexibility or the degree of non-linearity of a function at a point.
In deep learning, the curvature of a function at a point is often quantified as the norm of the Hessian at that point [5–7].
Hessian norms are zero everywhere if and only if the function is linear, making them suitable to measure the degree of non-linearity.
However, they suffer from a dependence on the scaling of model gradients, which makes them unsuitable to study its interplay with model robustness.
In particular, robust models have small gradient norms [8], which naturally imply smaller Hessian norms.
But are they truly more linear as a result?
To be able to study robustness independent of non-linearity, we propose normalized curvature , which normalizes the Hessian norm by its corresponding gradient norm, thus disentangling the two measures.
Surprisingly, we find that normalized curvature is a stable measure across train and test samples (see Table 3), whereas the usual curvature is not.
A conceptually straightforward approach to train models with low-curvature input-output maps [6, 9] involves directly penalizing curvature locally at training samples.
However, these methods involve expensive Hessian computations, and only minimize local point-wise curvature and not curvature everywhere.
A complementary approach is that of Dombrowski et al.
[7], who propose architectures that have small global curvature, but do not explicitly penalize this curvature during training.
In contrast, we propose efficient mechanisms to explicitly penalize the normalized curvature globally.
In addition, while previous methods [6, 7] penalize the Frobenius norm of the Hessian, we penalize its spectral norm, thus providing tighter and more interpretable robustness bounds.
Our overall contributions in this paper are:
In §3.1 we propose to measure curvature of deep models via normalized curvature , which is invariant to scaling of the model gradients, in particular, the magnitude of linear layers.
In §3.2 we show that normalized curvature of neural networks can be upper bounded via a decomposition in terms of normalized curvatures and slopes of individual layers.
We introduce develop an architecture for training LCNNs combining a novel activation function called the centered-softplus (§3.3), and recent innovations to constrain the Lipschitz constant of convolutional (§3.4.1) and batch normalization (§3.4.2) layers.
In §4 we prove bounds on the relative gradient robustness and adversarial robustness of models in terms of normalized curvature, showing that controlling normalized curvature directly controls these properties.
In §5, we show experiments demonstrating that our proposed innovations are successful in training low-curvature models without sacrificing training accuracy, and that such models have robust gradients, and are more robust to adversarial examples out-of-the-box.
Related Work
Adversarial Robustness of Neural Networks:
The well-known phenemenon of adversarial vulnerabilities of neural networks [10, 11] shows that adding small amounts of imperceptible noise can cause deep neural networks to misclassify points with high confidence.
The canonical method to defend against this vulnerability is adversarial training [12] which trains models to accurately classify adversarial examples generated via an attack such as projected gradient descent (PGD).
However, this approach is computationally expensive and provides no formal guarantees on robustness.
Cohen et al.
[13] proposed randomized smoothing, which provides a formal guarantee on robustness by generating a smooth classifer from any black-box classifier.
Hein and Andriushchenko [8] identified the local Lipschitz constant as critical quantity to prove formal robustness guarantees.
Moosavi Dezfooli et al.
[6] penalize the Frobenius norm of the Hessian, and show that they performs similarly to models trained via adversarial training.
Qin et al.
[9] introduce a local linearity regularizer, which also implicitly penalizes the Hessian.
Similar to these works, we enforce low curvature to induce robustness, but we focus on out of the box robustness of LCNNs.
Unreliable Gradient Interpretations in Neural Networks: Gradient explanations in neural networks can be unreliable.
Ghorbani et al.
[14], Zhang et al.
[15] showed that for any input, it is possible to find adversarial inputs such that the gradient explanations for these points that are highly dissimilar to each other.
Srinivas and Fleuret [16] showed that pre-softmax logit gradients however are independent of model behaviour, and as a result we focus on post-softmax loss gradients in this work.
Ros and Doshi-Velez [17] showed empirically that robustness can be improved by gradient regularization, however Dombrowski et al.
[18] showed that gradient instability is primarily due to large Hessian norms.
This suggests that the gradient penalization in Ros and Doshi-Velez [17] performed unintentional Hessian regularization, which is consistent with our experimental results.
To alleviate this, Dombrowski et al.
[7] proposed to train low curvature models via softplus activations and weight decay, broadly similar to our strategy.
However, while Dombrowski et al.
[7] focused on the Frobenius norm of the Hessian, we penalize the normalized curvature – a scaled version of the
Hessian spectral norm – which is strictly smaller than the Frobenius norm, and this results in a more sophisticated penalization strategy.
Lipschitz Layers in Neural Networks There has been extensive work on methods to bound the Lipschitz constant of neural networks.
Cisse et al.
[19] introduced Parseval networks, which penalizes the deviation of linear layers from orthonormality – since an orthonormal linear operator evidently has a Lipschitz constant of one, this shrinks the Lipschitz constant of a layer towards one.
Trockman and Kolter [20] use a reparameterization of the weight matrix, called the Cayley Transform, that is orthogonal by construction.
Miyato et al.
[21], Ryu et al.
[22] proposed spectral normalization, where linear layers are re-parameterized by dividing by their spectral norm, ensuring that the overall spectral norm of the parameterized layer is one.
Training Low-Curvature Neural Networks
In this section, we introduce our approach for training low-curvature neural nets (LCNNs).
Unless otherwise specified, we shall consider a neural network classifier f
that maps inputs x ∈ R d to logits which characterize the prediction, and can be further combined with the true label distribution and a loss function to give a scalar loss value f ( x ) ∈ R + .
3.1 Measuring Relative Model Non-Linearity via Normalized Curvature
We begin our analysis by discussing a desirable definition of curvature C f ( x ) ∈ R + .
While curvature is well-studied topic in differential geometry [23] where Hessian normalization is a common theme, our discussion will be largely independent of this literature, in a manner more suited to requirements in deep learning.
Regardless of the definition, a typical property of a curvature measure is that C f ( x ) =
∀ x ∈ R d ⇐⇒ f is linear, and the higher the curvature, the further from linear the function is. Hence max x ∈ R d C f ( x ) can be thought of as a measure of a model’s non-linearity.
A common way to define curvature in machine learning [5–7] has been via Hessian norms.
However, these measures are sensitive to gradient scaling, which is undesirable.
After all, the degree of model non-linearity must intuitively be independent of how large its gradients are.
For example, if two functions f, g are scaled ( f = k × g ), or rotated versions of each other ( ∇ f = k ×∇ g ), then we would like them to have similar curvatures in the interest of disentangling curvature (i.e., degree of non-linearity) from scaling.
It is easy to see that Hessian norms ∥∇ 2 x f ( x ) ∥ 2 do not have this property, as scaling the function also scales the corresponding Hessian.
We would like to be able to recover low-curvature models with steep gradients (that are non-robust), precisely to be able to disentangle their properties from low-curvature models with small gradients (that are robust), which Hessian norms do not allow us to do.
To avoid this problem, we propose a definition of curvature that is approximately normalized :
C f ( x ) = ∥∇ 2 f ( x ) ∥ 2 / ( ∥∇ f ( x ) ∥ 2 + ε ) .
Here ∥∇ f ( x ) ∥ 2 and ∥∇ 2 f ( x ) ∥ 2 are the ℓ 2 norm of the gradient and the spectral norm of the Hessian, respectively, where ∇ f ( x ) ∈ R d , ∇ 2 f ( x ) ∈ R d × d , and ε > 0 is a small constant to ensure well-behavedness of the measure.
This definition measures
Hessian norm relative to the gradient norm, and captures a notion of relative local linearity, which can be seen via an application of Taylor’s theorem:
[Equation]
Here the numerator on the left captures the local linearity error, scaled by the gradient norm in the denominator, which can be shown using Taylor’s theorem to be upper bounded by the normalized curvature.
We shall consider penalizing this notion of normalized curvature, and we shall simply refer to this
3.2 A Data-Free Upper Bound on Curvature
Directly penalizing the curvature is computationally expensive as it requires backpropagating an estimate of the Hessian norm, which itself requires backpropagating gradient-vector products.
This requires chaining the entire computational graph of the model at least three times.
Moosavi Dezfooli et al.
[6] reduce the complexity of this operation by computing a finite-difference approximation to the Hessian from gradients, but even this double-backpropagation is expensive.
We require an efficient penalization procedure takes a single backpropagation step.
To this end, we propose to minimize a data-free upper bound on the curvature.
To illustrate the idea, we first show this upper bound for the simplified case of the composition of one-dimensional functions ( f : R → R ).
Lemma 1.
Given a 1-dimensional compositional function f = f
L ◦ f L − 1 ◦ ... ◦
f 1 with f i : R → R for i = 1 , 2 , . . .
, L , the normalized curvature C f := |
f ′′ /f ′ | , is bounded by
[Equation]
Proof.
We first have f ′ =
(cid:81)
L i =
1 f ′ i .
Differentiating this expression we have f ′′ = (cid:80) Li
= 1 f ′′ i (cid:81) i j =1 f ′ j (cid:81)
Lk =1 ,k ̸
= i f ′ k .
Dividing by f ′ , taking absolute value of both sides, and using the triangle inequality, we have the intended result.
Extending the above expression to functions R n → R is not straightforward, as the intermediate layer Hessians are order three tensors.
We derive this using a result from Wang et al.
[24] that connects the spectral norms of order- n tensors to the spectral norm of their matrix “unfoldings”.
The full derivation is presented in the appendix, and the (simplified) result is stated below.
Theorem 1.
Given a function f = f
L ◦ f L − 1 ◦ . . .
◦ f 1 with f i : R n i − 1 → R n i , the curvature C f can be bounded by the sum of curvatures of individual layers
C f i ( x ) , i.e.,
C f (
[Equation]
The rightmost term is independent of x and thus holds uniformly across all data points.
This bound shows that controlling the curvature and Lipschitz constant of each layer of a neural network enables us to control the overall curvature of the model.
Practical neural networks typically consist of 1 linear maps such as convolutions, fully connected layers, and batch normalization layers, and non-linear activation functions.
Linear maps have zero curvature by definition, and non-linear layers often have bounded gradients ( = 1 ), which simplifies computations.
In the sections that follow, we shall see how to penalize the remaining terms, i.e., the curvature of the non-linear activations, and the Lipschitz constant of the linear layers.
3.3 Centered-Softplus: Activation Function with Trainable Curvature
Theorem 2 shows that the curvature of a neural network depends on the curvature of its constituent activation functions.
We thus propose to use activation functions with minimal curvature.
The kink at the origin of the ReLU function implies an undefined second derivative.
On the other hand, a smooth activation such as the softplus function, s ( x ; β )
= log(1 + exp( βx ))
/β is better suited to analyzing questions of curvature.
Despite not being a common baseline choice, softplus does see regular use, especially where its smoothness facilitates analysis [25].
The curvature of the softplus function is
[Equation]
Thus using softplus with small β values ensures low curvature.
However, we observe two critical drawbacks of softplus preventing its usage with small β : (1) divergence for small β , where s ( x ; β → 0) = ∞ which ensures that well-behaved low curvature maps cannot be recovered and, (2) instability upon composition i.e., s n ( x = 0; β )
=
s
[Equation]
s ◦ ... ◦
s (cid:124)
[Equation]
Thus we have that s n ( x = 0; β ) → ∞ as n → ∞ , which shows that composing softplus functions exacerbates instability around the origin.
This is critical for deep models with large number of layers n , which is precisely the scenario of interest to us.
To remedy this problem, we propose centeredsoftplus s 0 ( x ; β ) , a simple modification to softplus by introducing a normalizing term as follows.
[Equation]
This ensures that s ( x = 0; β ) =
s n ( x = 0; β ) =
for any positive integer n , and hence also ensures stability upon composition.
More importantly, we have s 0 ( x ; β → 0) =
x/ 2 which is a scaled linear map, while still retaining s 0 ( x ; β → ∞ )
=
ReLU ( x ) .
This ensures that we are able to learn both well-behaved linear maps, as well as highly non-linear ReLU-like maps if required.
We further propose to cast β is a learnable parameter and penalize its value, hence directly penalizing the curvature of that layer.
Having accounted for the curvature of the non-linearities, the next section discusses controlling the gradients of the linear layers.
3.4 Lipschitz Linear Layers
Theorem 2 shows that penalizing the Lipschitz constants of the constituent linear layers in a model is necessary to penalize overall model curvature.
There are broadly three classes of linear layers we consider: convolutions, fully connected layers, and batch normalization.
3.4.1 Spectrally Normalized Convolutions and Fully Connected Layers
We propose to penalize the Lipschitz constant of convolutions and fully connected layer via existing spectral normalization-like techniques.
For fully connected layers, we use vanilla spectral normalization (Miyato et al.
[21]) which controls the spectral norm of a fully connected layer by reparameterization – replacing a weight matrix W , by W/
|| W || 2 which has a unit spectral norm.
Ryu et al.
[22] generalize this to convolutions by presenting a power iteration method that works directly on the linear mapping implicit in a convolution: maintaining 3D left and right singular “vectors” of the 4D tensor of convolutional filters and developing the corresponding update rules.
They call this “real” spectral normalization to distinguish it from the approximation that [21] proposed for convolutions.
Using spectral normalization on fully connected layers and “real” spectral normalization on convolutional layers ensures that the spectral norm of these layers is exactly equal to one, further simplifying the bound in Theorem 2.
3.4.2 γ -Lipschitz Batch Normalization
In principle, at inference time batch normalization (BN) is multiplication by the diagonal matrix of inverse running standard deviation estimates.
Thus we can spectrally normalize the BN layer by computing the reciprocal of the smallest running standard deviation value across dimensions i , i.e., || BN || 2 = max x || BN ( x ) || 2 = 1 / min i running-std ( i ) .
In practice, we found that models with such spectrally normalized BN layers tended to either diverge or fail to train in the first place, indicating that the scale introduced by BN is necessary for training.
To remedy this, we introduce the γ - Lipschitz batch normalization, defined by
[Equation]
By clipping the scaling above at γ (equivalently, the running standard deviation below, at 1 /γ ), we can ensure that the Lipschitz constant of a batch normalization layer is at most equal to γ ∈ R + .
We provide a PyTorch-style code snippet in the appendix.
As with β , described in §3.3, we cast γ as a learnable parameter in order to penalize it during training.
Gouk et al.
[26] proposed a similar (simpler) solution, whereas they they fit a common γ to all BN layers, we let γ vary freely by layer.
3.5 Penalizing Curvature
We have discussed three architectural innovations – centered-softplus activations with a trainable β , spectrally normalized linear and convolution layers and a γ -Lipschitz batch normalization layer.
We now discuss methods to penalize the overall curvature of a model built with these layers.
Since convolutional and fully connected layers have spectral norms equal to one by construction, they contribute nothing to curvature.
Thus, we restrict attention to batch normalization and activation layers.
The set of which will be subsequently referred to respectively β SP and γ BN.
For models with batch normalization, naively using the upper bound in Theorem 2 is problematic due to the exponential growth in the product of Lipschitz constants of batch normalization layers.
To alleviate this, we propose to use a penalization R f where γ i terms are aggregated additively across batch normalization layers, and independent of the β i terms in the following manner:
[Equation]
An additive aggregation ensures that the penalization is well-behaved during training and does not grow exponentially.
Note that the underlying model is necessarily linear if the penalization term is zero, thus making it an appropriate measure of model non-linearity.
Also note that β i ≥ 0 , γ i ≥by construction.
We shall henceforth use the term “LCNN” to refer to a model trained with the proposed architectural components (centered-softplus, spectral normalization, and γ -Lipschitz batch normalization) and associated regularization terms on β, γ .
We next discuss the robustness and interpretability benefits that we obtain with LCNNs.
Why Train Low-Curvature Models?
In this section we discuss the advantages that low-curvature models offer, particularly as it pertains to robustness and gradient stability.
These statements apply not just to LCNNs, but low-curvature models in general obtained via any other mechanism.
4.1 Low Curvature Models have Stable Gradients
Recent work [14, 15] has shown that gradient explanations are manipulable, and that we can easily find inputs whose explanations differ maximally from those at the original inputs, making them unreliable in practice for identifying important features.
As we shall show, this amounts to models having a large curvature.
In particular, we show that the relative gradient difference across points ∥
ϵ ∥ 2 away in the ℓ 2 sense is upper bounded by the normalized curvature C f , as given below.
Proposition 1. Consider a model f with max x C f ( x ) ≤ δ C , and two inputs x and x + ϵ ( ∈ R d ).
The relative distance between gradients at these points is bounded by
[Equation]
The proof expands f ( x ) in a Taylor expansion around x + ϵ and bounds the magnitude of the second and higher order terms over the neighborhood of x , and the full argument is given in the appendix.
The upper bound is considerably simpler when we assume that the function locally quadratic, which corresponds to the rightmost term
∥ ϵ ∥ 2 C f ( x ) .
Thus the smaller the model curvature, the more locally stable are the gradients.
4.2 Low Curvature is Necessary for ℓ 2 Robustness
Having a small gradient norm is known to be an important aspect of adversarial robustness [8].
However, small gradients alone are not sufficient, and low curvature is necessary, to achieve robustness.
This is easy to see intuitively - a model may have low gradients at a point leading to robustness for small noise values, but if the curvature is large, then gradient norms at neighboring points can quickly increase, leading to misclassification for even slightly larger noise levels.
This effect is an instance of gradient-masking [27] , which provides an illusion of robustness by making models only locally robust.
In the result below, we formalize this intuition and establish an upper bound on the distance between two nearby points, which we show depends on both the gradient norm (as was known previously) and well as the max curvature of the underlying model.
Proposition 2. Consider a model f with max x C f ( x ) ≤ δ C , then for two inputs x and x + ϵ ( ∈ R d ), we have the following expression for robustness
[Equation]
The proof uses similar techniques to those of proposition 1, and is also given in the appendix.
This result shows that given two models with equally small gradients at data points, the greater robustness will be achieved by the model with the smaller curvature.
Experiments
In this section we perform experiments to (1) evaluate the effectiveness of our proposed method in training models with low curvature as originally intended, (2) evaluate whether low curvature models have robust gradients in practice, and (3) evaluate the effectiveness of low-curvature models for adversarial robustness.
Our experiments are primarily conducted on a base ResNet-18 architecture ([28]) using the CIFAR10 and CIFAR100 datasets ([29]), and using the Pytorch [30] framework.
Our methods entailed fairly modest computation – our most involved computations can be completed in under three GPU days, and all experimental results could be computed in less than 60 GPU-days.
We used a mixture of GPUs – primarily NVIDIA GeForce GTX 1080 Tis – on an internal compute cluster.
Baselines and Our Methods Our baseline model for comparison involves using a ResNet-18 model with softplus activation with a high β = 10 3 to mimic ReLU, and yet have well-defined curvatures.
Another baseline is gradient norm regularization (which we henceforth call ‘GradReg’)
[31], where the same baseline model is trained with an additional penalty on the gradient norm.
We train two variants of our approach - a base LCNN, which involves penalizing the curvature, and another variant combining LCNN penalty and gradient norm regularization (LCNN + GradReg), which controls both the curvature and gradient norm.
Our theory indicates that the LCNN + GradReg variant is likely to produce more robust models, which we verify experimentally.
We also compare with
CURE [6], softplus with weight decay [7] and adversarial training with ℓ 2 PGD [12] with noise magnitude of 0 . 1 and 3 iterations of PGD.
We provide experimental details in the appendix.
Parameter Settings All our models are trained for 200 epochs with an SGD + momentum optimizer, with a momentum of 0 .
9 and an initial learning rate of 0 . 1 which decays by a factor of 10 at 150 and 175 epochs, and a weight decay of 5 ×
10 − 4 .
5.1 Evaluating the Efficacy of Curvature Penalization
In this section, we evaluate whether LCNNs indeed reduce model curvature in practice.
Tablecontains our results, from which we make the following observations: (1) most of our baselines except CURE and adversarial training do not meaningfully lose predictive performance (2) GradReg and adversarially trained models are best at reducing gradient norm while LCNN-based models are best at penalizing curvature.
Overall, these experimental results show that LCNN-based models indeed minimize curvature as intended.
We also observe in Table 1 that GradReg [31] has an unexpected regularizing effect on the Hessian and curvature.
We conjecture that this is partly due to the following decomposition of the loss
Hessian, which can be written as ∇
[Equation]
, where LSE ( x ) is the LogSumExp function, and f l ( x ) is the pre-softmax logit output.
We observe that the first term strongly depends on the gradients, which may explain the Hessian regularization effects of GradReg, while the second term depends on the Hessian of the bulk of the neural network, which is penalized by the LCNN penalties.
This also explains why combining both penalizations (LCNN + GradReg) further reduces curvature.
We also measure the average per-epoch training times on a GTX 1080Ti, which are: standard models / softplus + weight decay ( ∼ 100 sec), LCNN ( ∼ 160 sec), GradReg ( ∼ 270 sec), LCNN+GradReg ( ∼ 350 sec), CURE / Adversarial Training ( ∼ 500 sec).
Note that the increase in computation for LCNN is primarily due to the use of spectral normalization layers.
The results show that LCNNs are indeed able to penalize curvature by only marginally ( 1 . 6 × ) increasing training time, and using LCNN+GradReg only increases time 1 .
3 × over GradReg while providing curvature benefits.
[Table]
5.2 Impact of Curvature on Gradient Robustness
In §4.
[Table]
we showed that low-curvature models tend to have more robust gradients.
Here we evaluate whether this prediction empirically by measuring the relative gradient robustness for the models with various ranges of curvature values and noise levels.
In particular
[Table]
we measure robustness to random noise at fixed magnitudes ranging logarithmically from 1 × 10 − 3 to 1 × 10 − 1 . We plot our results
[Figure]
[Figure]
as a function of added noise ϵ ∥ 2 on (a) CIFAR10 and (b) CIFAR100 with a ResNet-18 model.
We observe that low-curvature models lead to an order of magnitude improvement in gradient robustness, and this improvement closely follows the trend predicted by the theoretical upper bound in §4. in Figure 2 where we find that our results match theory (i.e, the simplified quadratic approximation in §4.1) quite closely in terms of the overall trends, and that low curvature models have an order of magnitude improvement in robustness over standard models.
5.3 Impact of Curvature on Adversarial Robustness
Our theory in §4.2 shows that having low curvature is necessary for robustness, along with having small gradient norms.
In this section we evaluate this claim empirically, by evaluating adversarial examples via ℓ 2 PGD [12] adversaries with various noise magnitudes.
We use the Cleverhans library [33] to implement PGD.
Our results are present in Table 2 where we find that LCNN+GradReg models perform on par with adversarial training, without a resulting accuracy loss.
[Table]
[Table]
5.4 Train-Test Discrepancy in Model Geometry
During our experiments, we observed a consistent phenomenon where the gradient norms and Hessian norms for test data points were much larger than those for the train data points, which hints at a form of overfitting with regards to these quantities.
We term this phenomenon as the train-test discrepancy in model geometry.
Interestingly, we did not observe any such discrepancy using our proposed curvature measure, indicating that our proposed measure may be a more reliable measure of model geometry.
We report our results in Table 3, where we measure the relative discrepancy, finding that the discrepancy for our proposed measure of curvature is multiple orders of magnitude smaller than the corresponding quantity for gradient and Hessian norms.
We leave further investi- gation of this phenomenon – regarding why curvature is stable across train and test – as a topic for future work.
[Table]
Summary of Experimental Results Overall, our experiments show that:
(1) LCNNs have lower curvature than standard models as advertised, and combining them with gradient norm regularization further decreases curvature (see Table 1).
The latter phenomenon is unexpected, as our curvature measure ignores gradient scaling.
(2) LCNNs combined with gradient norm regularization achieve an order of magnitude improved gradient robustness over standard models (see Figure 2).
(3) LCNNs combined with gradient norm regularization outperform adversarial training in terms of achieving a better predictive accuracy at a lower curvature (see Table 1), and are competitive in terms of adversarial robustness (see Table 2), while being ∼ 1 . 4 × faster.
(4) We observe that there exists a train-test discrepancy for standard geometric quantities like the gradient and Hessian norm, and this discrepancy disappears for our proposed curvature measure (see Table 3).
We also present ablation experiments, additional adversarial attacks, and evaluations on more datasets and architectures in the Appendix.
Discussion
In this paper, we presented a modular approach to remove excess curvature in neural network models.
Importantly, we found that combining vanilla LCNNs with gradient norm regularization resulted in models with the smallest curvature, the most stable gradients as well as those that are the most adversarially robust.
The current limitations of our approach are that we only consider convolutional and fully connected layers, and not self-attention or recurrent layers.
We also do not investigate the learning-theoretic benefits (or harms) of low-curvature models, or characterize how they may improve generalization for small number of training samples, or provide robustness to label noise (which we already observe in Fig. 1b).
Investigating these are important topics for future work.
Acknowledgments and Disclosure of Funding
The authors would like to thank the anonymous reviewers for their helpful feedback and all the funding agencies listed below for supporting this work.
SS and HL are supported in part by the NSF awards
# IIS-2008461 and # IIS-2040989, and research awards from Google, JP Morgan, Amazon,
Harvard Data Science Initiative, and D 3 Institute at Harvard.
KM and SS (partly) are supported by the Swiss National Science Foundation under grant number FNS-188758 “CORTI”.
HL would like to thank Sujatha and Mohan Lakkaraju for their continued support and encouragement.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A] We present foundational work on the geometry of neural network models and as such we do not foresee any negative societal impacts beyond those that already exist with the usage of neural network models.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] proofs are provided in the appendix 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] some training details are mentioned in the appendix (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects...
We did not use crowdsourcing or conducted research with human subjects
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Proof of Theorem
In this section, we provide the proof of the main theorem in the paper, which decomposes overall curvature into curvatures and slopes of constituent layers.
We state Theorem 1 below for reference.
Theorem 2.
Given a function f = f L ◦ f L − 1 ◦ . . . ◦
f 1 with f i : R n i − 1 → R n i , the curvature C f can be bounded as follows
[Equation]
This statement is slightly different than the one given in the paper, differing by a term of the width of each nonlinear layer.
Since we do not use or care about the units of curvature, only its minimization, and have elected to equally-weight each term of the sum, this is an inconsequential discrepancy.
A similar bound is constructed recursively in [34].
[7] gives a similar formula, albeit for the Frobenius norm.
The Frobenius norm is both simpler, because the sum of squared entries is independent of the layout of the data, and also weaker, since it cannot deliver a bound which holds uniformly in the data.
To our knowledge Equation 6 is the first explicit, easily-interpreted formula of its type.
We start with some preliminaries, with the actual proof being in subsection A.4.
A.1 Derivatives of compositional functions
For a function f : R d → R r , let ∇ f : R d → R d × r denote its gradient, and ∇ 2 f : R d → R d × r × d denote its Hessian.
We drop the argument to functions when possible, and all norms will be spectral norms.
Given L
[Equation]
: R n
[Equation]
R n k +
[Equation]
for 1 ≤ k ≤ k + j ≤ L .
This function composition will be our model of a deep neural network where f ℓ represents the action of the ℓ th layer.
If each f i is continuously differentiable, we have this formula for the gradient of f k,k + j
[Equation]
where we adopt the convention that f j,j ( x )
=
x in order to make the boundary conditions correct.
The product begins at the end, with ∇ f k + j and progresses forward through the indices – a straightforward consequence of the chain rule of differentiation.
Supposing moreover that each f i is twice-differentiable, we have this formula for the second derivative of f 1 ,k :
[Equation]
where ( ∇ 2 f 1 ,i − ∇ 2 f 1 ,i − 1 )
= ( ∇ 2 f i ) (cid:0) ∇ f 1 ,i − 1 , ∇ f ⊤ i +1 ,k , ∇ f 1 ,i − 1 (cid:1) ∈ R n 0 × n k × n(8) where we have used the covariant multilinear matrix multiplication notation: ∇ 2 f i is an order-three tensor ∈ R n i × n i × n i , with the first and third modes multiplied by ∇ f 1 ,i − 1 ∈ R n i × n 0 and the second mode multiplied by ∇ f ⊤ i +1 ,k ∈ R n i × n k .
A.2 Tensor calculus
In this section, we present a simplifed version of the notation from [24].
A k -linear map is a function of k (possibly multidimensional) variables such that if any k − 1 variables are held constant, the map is linear in the remaining variable.
A k -linear function is can be represented by an order- k tensor
A given elementwise by A = (cid:74)
a j 1 ...j
[Table]
(cid:75)
[Table]
The covariant multilinear matrix multiplication of a tensor with matrices (order 2 tensors)
M 1 = ( m (1) i 1 j 1 ) ∈ R
d 1 × s 1 , . . .
, M k
= ( m ( k ) i
k j k ) ∈
R
d k × s k is
A ( M 1 , . . . , M k ) =
(cid:116)
d 1 (cid:88)
i 1 = . . .
d k (cid:88)
i k a i 1 ...
i k m (1) i 1 j 1 . . .
m ( k ) i
k j k (cid:124) ∈
R s 1 × ...
× s k .
This operation can be implemented via iterated einsums as: def covariant_multilinear_matmul(a: torch.
Tensor, mlist: List[torch.
Tensor]) -> torch.
Tensor: order = a.ndim base_indices =
string.ascii_letters indices = base_indices[:order] next_index = base_indices[order] val = a for idx in range(order): resp_str = indices[:idx] + next_index + indices[idx+1:]
einsum_str = indices + f",{indices[idx]}{next_index}->{resp_str}" val =
torch.einsum(einsum_str, val, mlist[idx]) return val
For example, covariant multilinear matrix multiplication of an order two tensor is pre- and postmultiplication by its arguments: M ⊤ 1 A M 2 =
A ( M 1 , M 2 ) .
The generalization of the matrix spectral norm is ||A|| 2 = sup {A ( x 1 , . . .
, x k ) :
|| x i ||
= 1 , x i ∈ R
d i , i = 1 , 2 , . . . , k } .
The computation of order- k operator norms is hard in theory, and also in practice (cf.
In order to address this difficulty, we introduce an instance of the unfold operator.
For A ∈ R d 1 × d 2 × d 3 , unfold {{ 1 , 2 } , { 3 }} ( A )
∈
R
d 1 d 2 × d 3 is the matrix with the j th column being the flattened j th (in the final index)
d 2 × d 3 matrix.
1 Unfolding is useful because it it allows us to bound an order-3 operator norm in terms of order-2 operator norms – Wang et al.
[24] shows that ||A|| ≤ || unfold {{ 1 , 2 } , { 3 }} ( A ) || .
The upper bound – the operator norm of a matrix – can computed with standard largest singular-value routines.
A similar unfolding-based bound was used in the deep-learning context by [36] to give improved estimates on the spectral norm of convolution operators.
To facilitate the analysis of unfolded tensors, we coin operations that put to the diagonal of tensors:
• pdiag 2 : R d (cid:55)→ R d × d defined by pdiag 2( x ) ij = (cid:26) x j
if i =
j
otherwise.
• pdiag 3 : R d (cid:55)→ R d × d × d defined by pdiag 3( x ) ijk = (cid:26) x j
if
i =
j = k
otherwise.
Further, let 1 n ∈ R n be a vector of ones, I n = pdiag 2(1 n ) ∈ R n × n be the n -dimensional identity matrix, and I n = pdiag 3(1 n ) ∈ R n × n × n .
For two vectors a ∈ R n , b ∈ R n , let ab denote the elementwise product.
⊗ denotes the well-understood Kronecker product, so that, for example, 1 ⊤ n ⊗
I m is an m × nm matrix consisting of n copies of the m × m identity matrix stacked side by side.
Where it is redundant, we drop the subscripts indicating dimension.
We use the following facts about tensors, their unfoldings, and their operator norms, in what follows.
[Equation]
M 2 , pdiag 2( b ) M 3 ) 2.
unfold {{ 1 , 2 } , { 3 }} ( I ( M 1 , M 2 , M 3 )) =
( M 1 ⊗ I s 2 ) ⊤ unfold {{ 1 , 2 } , { 3 }} ( I ( I d 1 , M 2 , M 3 )) 3.
|| (1 ⊤ d 1 ⊗ I s 2 ) unfold {{ 1 , 2 } , { 3 }} ( I ( I d 1 , M 2 , M 3 ))
|| ≤ || M ⊤ 2 M 3 || s
Taken together Fact #2 and #3 imply that
[Equation]
which is our essential bound for the norm of an order-3 tensor in terms of order-2 tensors
[Equation]
A.3 Hessian increment bound
Let σ ( x ) = exp( x ) / (1 + exp( x )))
denote the (elementwise) logistic function R d (cid:55)→ R d .
The derivatives of s ( x ; β ) can be written as
[Equation]
Let the i th softplus layer have coefficient β
i , then the increment from Equation 8, can be bounded as follows:
[Equation]
Equation 13 follows by Fact #1 above, along with Equation 11.
Equation 14 is the standard unfolding bound by [24].
Equation 15 is our main bound on order-3 tensors in terms of order-2 matrices, Equation 9.
Equation 16 follows from the Cauchy-Schwartz inequality.
The replacement in the last term of the product is Equation 10.
Equation 17 rewrites Equation 16 using Equation 7: f ⊤ i +1 ,k ∇ f i ∇ f 1 ,i − 1 = ∇ f 1 ,k and the definition of the curvature of f i .
A.4 Putting it together
Ignoring the ϵ term,
[Equation]
≤
[Equation]
i = n i ×
[Equation]
Equation 19 substitutes Equation 8.
Equation 20 is the triangle inequality.
Equation 21 is Equation 17, along with cancelling the term of ||∇ f
|| top and bottom.
Equation 22 are Equation 23 are obvious and a standard simplification in the literature on controlling the Lipschitz constant of neural networks.
Because exactly computing the smallest Lipschitz constant of a general neural network is
NP-complete, a widely-used baseline measure of Lipschitz-smoothness is rather the product of the
Lipschitz constants of smaller components of the network, such as single layers ([37]).
B Loss curvature vs logit curvature
We have thusfar discussed how to assure a bound on curvature independent of the input x .
However, a truly data-independent bound must also not depend on the class label y .
In this short section we discuss the main consideration that accompanies this: the differences between loss and (pre-softmax) logit curvature.
Let the function including the loss for a class c be denoted by f c loss , and the logit corresponding to a class c ∈ [1 , C ] be f c logit , and let f 1 ,C logit denote the vector valued function corresponding to all the logits.
Then we have f c loss =
lsm c ◦ f 1 ,C logit , where lsm c ( x
[Equation]
= − log
[Equation]
the negative log softmax function.
The derivatives of this function are:
[Equation]
whose norm is upper bounded by 1 and 0 .
respectively.
This implies that bounding the logit curvature also ensures that the loss curvature is bounded, as the gradients and Hessians of the negative log softmax layer do not explode.
In other words, C f loss < C f logit .
However penalizing an upper bound may not be the most efficient way to penalize the loss
Hessian.
Consider the following decomposition of the loss Hessian, which can be written as ∇ 2 f ( x ) ∼ ∇ f logit ( x ) ∇ 2 f l LSE ( x ) ∇ f logit ( x ) ⊤ + ∇ f l LSE ( x ) ∇ 2 f logit ( x ) , where LSE ( x ) is the LogSumExp function.
Thus a more efficient way to penalize the loss Hessians is to both use our LCNN penalties, as well as penalize the gradient norm, which we also find to be true in our experiments.
C γ -Lipschitz Batchnorm
We present here a Pytorch-style pseudo-code for the γ -Lipschitz Batchnorm for clarity.
bn = torch.nn.
BatchNorm2d(in_channels, affine=False) log_lipschitz = torch.nn. Parameter(torch.tensor(init_lipschitz).log()) ...
#perform spectral normalization for BN in closed form bn_spectral_norm = torch.max(1 / (bn.running_var + 1e-5).sqrt()) one_lipschitz_bn = bn(x) / bn_spectral_norm # multiply normalized BN with a learnable scale scale = torch.min((bn.running_var + 1e-5) ** .5) one_lipschitz_part = bn(x)
* scale x = one_lipschitz_part * torch.minimum(1 / scale, log_lipschitz.exp())
Note that we parameterize γ in terms of its log value, to ensure it remains positive during training.
We employ the same method for fitting the centered softplus parameter, β .
D Proofs of LCNN Properties
Here we present proofs for the properties linking LCNNs to gradient smoothness and adversarial robustness.
To this end, we first prove a lemma that is used in both these results, which states that the gradient norm can rise exponentially in the worst case.
In the text below, B δ ( x ) denotes the δ -ball around x .
Lemma 2.
The ratio of gradient norms at nearby points rises exponentially with distance between those points.
For a function f satisfying max x ′ ∈
B δ ( x ) C f ( x ′ ) =
δ C , points x and x + ϵ that are ∥ ϵ ∥
= r away, we have
[Equation]
Proof.
Let g ( x ) =
log ∥∇ f ( x ) ∥ 22 .
Applying a first order Taylor series with Lagrange remainder / mean-value theorem, we have
[Equation]
for some ξ ∈ B
[Equation]
Note that to apply the Cauchy-Schwartz inequality, we first upper bound the term on the right by its 2-norm, which for a scalar is simply its absolute value.
Taking exponent on both sides of the final expression and taking square root, we have the intended result.
D.1 LCNNs have Robust Gradients
Proposition 3.
Let max x ′ ∈
B ϵ ( x ) C f
( x ′ ) ≤ δ C , then the relative distance between gradients at x and x + ϵ is
[Equation]
Approximation )
Proof.
We begin the proof by invoking the Taylor series approximation of ∇ f ( x + ϵ ) at x , and using the explicit Lagrange form of the Taylor error.
This is equivalent to using a form of the multivariate mean-value theorem.
Let ξ ∈ B ϵ ( x ) , then there exists some ξ such that the following holds
[Equation]
Taylor Series
[Equation]
- Schwartz
[Equation]
by gradnorm [Equation]
Plugging in the value of ∥∇ f
( x + ξ 0 ) ∥ 2 / ∥∇ f ( x ) ∥ 2 from Lemma 2, and further upper bounding
C f ( x + ξ ) ≤ δ C
we have the intended result.
To derive the simplified quadratic approximation, replace the first step in the Taylor series with ∇ f ( x + ϵ ) − ∇ f ( x ) = ∇ 2 f ( x ) ⊤ ϵ , i.e., use ξ = 0 . D
.2 Curvature
is Necessary for Robustness
Proposition 4.
Let max x ′ ∈
B ϵ ( x ) C f
( x ′ ) ≤ δ C , then for two points x and x + ϵ ,
[Equation]
Proof.
Let ξ ∈ B ϵ
( x ) , then there exists some ξ such that the following holds
[Equation]
Taylor Series
[Equation]
( Defn of C f ( x + ξ )) Factoring out ∥∇ f ( x ) ∥ 2 ∥ ϵ ∥ 2 in the RHS, and using Lemma 2, and further upper bounding C f
( x + ξ ) ≤ δ C
we have the intended result.
To derive the simplified quadratic approximation, replace the first step in the Taylor series with f ( x + ϵ ) − f ( x ) = ∇ f ( x ) ⊤ ϵ +
12 ϵ ⊤ ∇ 2 f ( x ) ϵ , i.e., use ξ = 0 .
E Experimental Settings
In this section we elaborate on the hyper-parameter settings used for our tuning our models.
For the standard ResNet-18, we use standard hyper-parameter settings as indicated in the main paper, and we do not modify this for the other variants.
For LCNNs, we chose regularizing constants as λ β = 10 − 4 and λ γ = 10 − 5 .
For GradReg, we use λ grad = 10 − 3 , and for LCNNs + GradReg, we chose λ β = 10 − 4 , λ γ = 10 − 5 , λ grad = 10 − 3 .
We performed a coarse grid search and chose the largest regularizing constants that did not affect predictive performance.
F Additional Experiments F.1 Ablation Experiments
In this section
[Table] we present ablation studies where we train models with each of the proposed modifications separately, i [Table] e., we train a model with only spectral norm for the convolution layers, γLipschitz Batchnorm or centered softplus.
Our results in Table 4 show that for the resnet18 architecture considered,
[Table]
(1) performing spectral normalization had no effect of curvature as presumably the batchnorm layers are able to compensate for the lost flexibility, and (2) either penalizing the batchnorm alone or
the softplus alone
performs almost as well as LCNN, or sometimes even better in terms of curvature reduction.
Note that the most expensive computational step is the spectral norm for the convolutional layers, indicating that avoiding this in practice may yield speedups in practice.
While in practice for Resnet-18 only γ -Lipschitz or centered softplus is sufficient for curvature reduction, in theory we must regularize all components to avoid scaling restrictions in one layer being compensated by other layers, as dictated by the upper bound.
In particular, this means that the strategy of penalizing only a subset of layers may not generalize to other architectures.
F.2 Robustness Evaluation on RobustBench / Autoattack
The attack presented in section 5 was relatively “weak” – a network could be truly susceptible to adversarial attack, but by only testing against a weak adversary, we could fail to notice.
Short of a comprehensive verification (e.g. [38]), which is known to be computationally intractable at scale, there is no fully satisfactory way to guarantee robustness.
However, one common method to develop confidence in a model is to demonstrate robustness in the face of a standard set of nontrivially capable attackers.
This is what we do in Table 5, where we use the Robustbench software library [39] to evaluate both a white-box (having access to the internal details of the model), and a black-box (using only function evaluations) attacks.
[Table]
LCNN+GradReg 52.
CURE 50. Softplus+wt decay 23. Adv Training 55. CURE + GradReg 60.
LCNN + GradReg + Adv Training 59.
These experimental results show overall that:
(1) LCNN + GradReg is still on par with adversarial training even against stronger attacks such as APGD-t and Square, as they are with PGD.
(2) Combining LCNN + GradReg with adversarial training (in the last row) further improves robustness at the cost of predictive accuracy.
(3) Combining CURE with GradReg (in the penultimate row) improves robustness at the cost of further deteriorating predictive accuracy.
F.3 Additional Evaluations on more Architectures / Dataset Combinations
We present results on the following dataset - architecture pairs:
(1) In Table 6, we present results on SVHN dataset and VGG-11 model (2) In Table 7, we present results on SVHN dataset and ResNet-18 model (3) In Table 8, we present results on CIFAR-100 dataset and VGG-11 model
In all cases, we find that our results are on par with our experiment done on the CIFAR- and
ResNet-18 setup, which confirms that generality of our approach.
[Table]
[Table] [Table]