Robust Learning with Jacobian Regularization

1 INTRODUCTION . Stability analysis lies at the heart of many scientific and engineering disciplines . In an unstable system , infinitesimal perturbations amplify and have substantial impacts on the performance of the system . It is especially critical to perform a thorough stability analysis on complex engineered systems deployed in practice , or else what may seem like innocuous perturbations can lead to catastrophic consequences such as the Tacoma Narrows Bridge collapse ( Amman et al. , 1941 ) and the Space Shuttle Challenger disaster ( Feynman and Leighton , 2001 ) . As a rule of thumb , well-engineered systems should be robust against any input shifts – expected or unexpected . Most models in machine learning are complex nonlinear systems and thus no exception to this rule . For instance , a reliable model must withstand shifts from training data to unseen test data , bridging the so-called generalization gap . This problem is severe especially when training data are strongly biased with respect to test data , as in domain-adaptation tasks , or when only sparse sampling of a true underlying distribution is available , as in few-shot learning . Any instability in the system can further be exploited by adversaries to render trained models utterly useless ( Szegedy et al. , 2013 ; Goodfellow et al. , 2014 ; Moosavi-Dezfooli et al. , 2016 ; Papernot et al. , 2016a ; Kurakin et al. , 2016 ; Madry et al. , 2017 ; Carlini and Wagner , 2017 ; Gilmer et al. , 2018 ) . It is thus of utmost importance to ensure that models be stable against perturbations in the input space . Various regularization schemes have been proposed to improve the stability of models . For linear classifiers and support vector machines ( Cortes and Vapnik , 1995 ) , this goal is attained via an L2 regularization which maximizes classification margins and reduces overfitting to the training data . This regularization technique has been widely used for neural networks as well and shown to promote generalization ( Hinton , 1987 ; Krogh and Hertz , 1992 ; Zhang et al. , 2018 ) . However , it remains unclear whether or not L2 regularization increases classification margins and stability of a network , especially for deep architectures with intertwining nonlinearity . In this paper , we suggest ensuring robustness of nonlinear models via a Jacobian regularization scheme . We illustrate the intuition behind our regularization approach by visualizing the classification margins of a simple MNIST digit classifier in Figure 1 ( see Appendix A for more ) . Decision cells of a neural network , trained without regularization , are very rugged and can be unpredictably unstable ( Figure 1a ) . On average , L2 regularization smooths out these rugged boundaries but does not necessarily increase the size of decision cells , i.e. , does not increase classification margins ( Figure 1b ) . In contrast , Jacobian regularization pushes decision boundaries farther away from each training data point , enlarging decision cells and reducing instability ( Figure 1c ) . The goal of the paper is to promote Jacobian regularization as a generic scheme for increasing robustness while also being agnostic to the architecture , domain , or task to which it is applied . In support of this , after presenting the Jacobian regularizer , we evaluate its effect both in isolation as well as in combination with multiple existing approaches that are intended to promote robustness and generalization . Our intention is to showcase the ease of use and complimentary nature of our proposed regularization . Domain experts in each field should be able to quickly incorporate our regularizer into their learning pipeline as a simple way of improving the performance of their state-of-the-art system . The rest of the paper is structured as follows . In Section 2 we motivate the usage of Jacobian regularization and develop a computationally efficient algorithm for its implementation . Next , the effectiveness of this regularizer is empirically studied in Section 3 . As regularlizers constrain the learning problem , we first verify that the introduction of our regularizer does not adversely affect learning in the case when input data remain unperturbed . Robustness against both random and adversarial perturbations is then evaluated and shown to receive significant improvements from the Jacobian regularizer . We contrast our work with the literature in Section 4 and conclude in Section 5 . 2 METHOD . Here we introduce a scheme for minimizing the norm of an input-output Jacobian matrix as a technique for regularizing learning with stochastic gradient descent ( SGD ) . We begin by formally defining the input-output Jacobian and then explain an efficient algorithm for computing the Jacobian regularizer using standard machine learning frameworks . 2.1 STABILITY ANALYSIS AND INPUT-OUTPUT JACOBIAN . Let us consider the set of classification functions , f , which take a vectorized sensory signal , x ∈ RI , as input and outputs a score vector , z = f ( x ) ∈ RC , where each element , zc , is associated with likelihood that the input is from category , c.1 In this work , we focus on learning this classification function as a neural network with model parameters θ , though our findings should generalize to any parameterized function . Our goal is to learn the model parameters that minimize the classification objective on the available training data while also being stable against perturbations in the input space so as to increase classification margins . 1Throughout the paper , the vector z denotes the logit before applying a softmax layer . The probabilistic output of the softmax pc relates to zc via pc ≡ e zc/T∑ c′ e z c′/T with temperature T , typically set to unity . The input-output Jacobian matrix naturally emerges in the stability analysis of the model predictions against input perturbations . Let us consider a small perturbation vector , ∈ RI , of the same dimension as the input . For a perturbed input x̃ = x+ , the corresponding output values shift to z̃c = fc ( x+ ) = fc ( x ) + I∑ i=1 i · ∂fc ∂xi ( x ) +O ( 2 ) = zc + I∑ i=1 Jc ; i ( x ) · i +O ( 2 ) ( 1 ) where in the second equality the function was Taylor-expanded with respect to the input perturbation and in the third equality the input-output Jacobian matrix , Jc ; i ( x ) ≡ ∂fc ∂xi ( x ) , ( 2 ) was introduced . As the function f is typically almost everywhere analytic , for sufficiently small perturbations the higher-order terms can be neglected and the stability of the prediction is governed by the input-output Jacobian . 2.2 ROBUSTNESS THROUGH INPUT-OUTPUT JACOBIAN MINIMIZATION . From Equation ( 1 ) , it is straightforward to see that the larger the components of the Jacobian are , the more unstable the model prediction is with respect to input perturbations . A natural way to reduce this instability then is to decrease the magnitude for each component of the Jacobian matrix , which can be realized by minimizing the square of the Frobenius norm of the input-output Jacobian,2 ||J ( x ) ||2F ≡ ∑ i , c [ Jc ; i ( x ) ] 2 . ( 3 ) For linear models , this reduces exactly to L2 regularization that increases classification margins of these models . For nonlinear models , however , Jacobian regularization does not equate to L2 regularization , and we expect these schemes to affect models differently . In particular , predictions made by models trained with the Jacobian regularization do not vary much as inputs get perturbed and hence decision cells enlarge on average . This increase in stability granted by the Jacobian regularization is visualized in Figure 1 , which depicts a cross section of the decision cells for the MNIST digit classification problem using a nonlinear neural network ( LeCun et al. , 1998 ) . The Jacobian regularizer in Equation ( 3 ) can be combined with any loss objective used for training parameterized models . Concretely , consider a supervised learning problem modeled by a neural network and optimized with SGD . At each iteration , a mini-batch B consists of a set of labeled examples , { xα , yα } α∈B , and a supervised loss function , Lsuper , is optimized possibly together with some other regularizerR ( θ ) – such as L2 regularizer λWD2 θ 2 – over the function parameter space , by minimizing the following bare loss function Lbare ( { xα , yα } α∈B ; θ ) = 1 |B| ∑ α∈B Lsuper [ f ( xα ) ; yα ] +R ( θ ) . ( 4 ) To integrate our Jacobian regularizer into training , one instead optimizes the following joint loss LBjoint ( θ ) = Lbare ( { xα , yα } α∈B ; θ ) + λJR 2 [ 1 |B| ∑ α∈B ||J ( xα ) ||2F ] , ( 5 ) where λJR is a hyperparameter that determines the relative importance of the Jacobian regularizer . By minimizing this joint loss with sufficient training data and a properly chosen λJR , we expect models to learn both correctly and robustly . 2Minimizing the Frobenius norm will also reduce the L1-norm , since these norms satisfy the inequalities ||J ( x ) ||F ≤ ∑ i , c ∣∣Jc ; i ( x ) ∣∣ ≤ √IC||J ( x ) ||F . We prefer to minimize the Frobenius norm over the L1-norm because the ability to express the former as a trace leads to an efficient algorithm [ see Equations ( 6 ) through ( 8 ) ] . 2.3 EFFICIENT APPROXIMATE ALGORITHM . In the previous section we have argued for minimizing the Frobenius norm of the input-output Jacobian to improve robustness during learning . The main question that follows is how to efficiently compute and implement this regularizer in such a way that its optimization can seamlessly be incorporated into any existing learning paradigm . Recently , Sokolić et al . ( 2017 ) also explored the idea of regularizing the Jacobian matrix during learning , but only provided an inefficient algorithm requiring an increase in computational cost that scales linearly with the number of output classes , C , compared to the bare optimization problem ( see explanation below ) . In practice , such an overhead will be prohibitively expensive for many large-scale learning problems , e.g . ImageNet classification has C = 1000 target classes ( Deng et al. , 2009 ) . ( Our scheme , in contrast , can be used for ImageNet : see Appendix H. ) Here , we offer a different solution that makes use of random projections to efficiently approximate the Frobenius norm of the Jacobian.3 This only introduces a constant time overhead and can be made very small in practice . When considering such an approximate algorithm , one naively must trade off efficiency against accuracy for computing the Jacobian , which ultimately trades computation time for robustness . Prior work by Varga et al . ( 2017 ) briefly considers an approach based on random projection , but without providing any analysis on the quality of the Jacobian approximation . Here , we describe our algorithm , analyze theoretical convergence guarantees , and verify empirically that there is only a negligible difference in model solution quality between training with the exact computation of the Jacobian as compared to training with the approximate algorithm , even when using a single random projection ( see Figure 2 ) . Given that optimization is commonly gradient based , it is essential to efficiently compute gradients of the joint loss in Equation ( 5 ) and in particular of the squared Frobenius norm of the Jacobian . First , we note that automatic differentiation systems implement a function that computes the derivative of a vector such as z with respect to any variables on which it depends , if the vector is first contracted with another fixed vector . To take advantage of this functionality , we rewrite the squared Frobienus norm as ||J ( x ) ||2F = Tr ( JJT ) = ∑ { e } eJJTeT = ∑ { e } [ ∂ ( e · z ) ∂x ] 2 , ( 6 ) where a constant orthonormal basis , { e } , of the C-dimensional output space was inserted in the second equality and the last equality follows from definition ( 2 ) and moving the constant vector inside the derivative . For each basis vector e , the quantity in the last parenthesis can then be efficiently computed by differentiating the product , e · z , with respect to input parameters , x . Recycling that computational graph , the derivative of the squared Frobenius norm with respect to the model parameters , θ , can be computed through backpropagation with any use of automatic differentiation . Sokolić et al . ( 2017 ) essentially considers this exact computation , which requires backpropagating gradients through the model C times to iterate over the C orthonormal basis vectors { e } . Ultimately , this incurs computational overhead that scales linearly with the output dimension C. Instead , we further rewrite Equation ( 6 ) in terms of the expectation of an unbiased estimator ||J ( x ) ||2F = C Ev̂∼SC−1 [ ||v̂ · J ||2 ] , ( 7 ) where the random vector v̂ is drawn from the ( C − 1 ) -dimensional unit sphere SC−1 . Using this relationship , we can use samples of nproj random vectors v̂µ to estimate the square of the norm as ||J ( x ) ||2F ≈ 1 nproj nproj∑ µ=1 [ ∂ ( v̂µ · z ) ∂x ] 2 , ( 8 ) which converges to the true value as O ( n−1/2proj ) . The derivation of Equation ( 7 ) and the calculation of its convergence make use of random-matrix techniques and are provided in Appendix B . Finally , we expect that the fluctuations of our estimator can be suppressed by cancellations within a mini-batch . With nearly independent and identically distributed samples in a mini-batch of size 3In Appendix C , we give an alternative method for computing gradients of the Jacobian regularizer by using an analytically derived formula . Algorithm 1 Efficient computation of the approximate gradient of the Jacobian regularizer . Inputs : mini-batch of |B| examples xα , model outputs zα , and number of projections nproj . Outputs : Square of the Frobenius norm of the Jacobian JF and its gradient∇θJF . JF = 0 for i = 1 to nproj do { vαc } ∼ N ( 0 , I ) . ( |B| , C ) -dim tensor with each element sampled from a standard normal . v̂α = vα/||vα|| . Uniform sampling from the unit sphere for each α. zflat = Flatten ( { zα } ) ; vflat = Flatten ( { v̂α } ) . Flatten for parallelism . Jv = ∂ ( zflat · vflat ) /∂xα JF += C||Jv||2/ ( nproj|B| ) end for ∇θJF = ∂JF /∂θ return JF , ∇θJF |B| 1 , we expect the error in our estimate to be of order ( nproj|B| ) −1/2 . In fact , as shown in Figure 2 , with a mini-batch size of |B| = 100 , single projection yields model performance that is nearly identical to the exact method , with computational cost being reduced by orders of magnitude . The complete algorithm is presented in Algorithm 1 . With a straightforward implementation in PyTorch ( Paszke et al. , 2017 ) and nproj = 1 , we observed the computational cost of the training with the Jacobian regularization to be only ≈ 1.3 times that of the standard SGD computation cost , while retaining all the practical benefits of the expensive exact method.4

Summary: Stability is one of the important aspects of machine learning.  This paper views Jacobian regularization as a scheme to improve the stability, and studies the behavior of Jacobian regularization under random input perturbations, adversarial input perturbations, train/test distribution shift, and simply as a regularization tool for the classical setting without any distribution shifts nor perturbations.  There are already several related works that propose to use Jacobian regularization, but previous works didn’t have an efficient algorithm and also did not have theoretical convergence guarantee.  This paper offers a solution that efficiently approximate the Frobenius norm of the Jacobian and also show the optimal convergence rate for the proposed method.  Various experiments show that the behavior of Jacobian regularization and show that it is robust.