Fast Jacobian-Vector Product for Deep Networks

Deep (Neural) Networks (DNs) are used throughout the spectrum of machine learning tasks ranging from image classification (Lawrence et al., 1997; Lin et al., 2013; Zagoruyko & Komodakis, 2016) to financial fraud prediction (Rushin et al., 2017; Roy et al., 2018; Zheng et al., 2018). The ability of DNs to reach super-human performances added great momentum into the deployment of those techniques into real life applications e.g. self-driving cars (Rao & Frtunikj, 2018), fire prediction (Kim & Lee, 2019) or seismic activity prediction (Seydoux et al., 2020). The main ingredients behind this success roughly fall into three camps: (i) specialized optimization algorithms such as Adam (Kingma & Ba, 2014) coupled with carefully designed learning rate schedules (Smith, 2017) and/or mini-batch size (Devarakonda et al., 2017; Wang et al., 2017), (ii) novel regularization techniques which can be explicit as with orthogonal/sparse weights penalties (Cohen & Welling, 2016; Jia et al., 2019; Wang et al., 2020) or implicit as with dropout (Molchanov et al., 2017) or data-augmentation (Taylor & Nitschke, 2017; Perez & Wang, 2017; Hernández-García & König, 2018; Shorten & Khoshgoftaar, 2019); and (iii) adaptivity of the DN architecture to the data and task at hand thanks to expert knowledge (Graves et al., 2013; Seydoux et al., 2020) or automated Neural Architecture Search (Elsken et al., 2019).

Despite DNs’ astonishing performance, their practical deployment in the real-world remains hazardous, requiring tremendous tweaking and the implementation of various safe guard measures (Marcus, 2018) to avoid failure cases that can lead to injuries or death, but also to ensure fairness and unbiased behaviors (Gianfrancesco et al., 2018). One path for overcoming those challenges in a more principled way is to develop novel, specialized methods for each part of the deep learning pipeline (optimization/regularization/architecture).

Recent techniques aiming at tackling those shortcomings have been proposed. Most of those techniques rely on Jacobian-vector products (JVPs) (Hirsch & Smale, 1974), which are defined as the projection of a given vector onto the Jacobian matrix of an operator, i.e. the DN in our case. The JVP captures crucial information on the local geometry of the DN input-output mapping which is one of the main reason behind its popularity. For example, Dehaene et al. (2020) leverage the JVP to better control the denoising/data generation process of autoencoders via explicit manifold projections (generated samples are repeatedly projected onto the autoencoder Jacobian matrix); Balestriero & Baraniuk (2020) assess the sensitivity of DNs to adversarial examples (noise realizations are projected onto the DN Jacobian matrix); Cosentino et al. (2020) guarantees the generalization performance of DNs based on a novel JVP based measure (neighboring data samples are projected onto the DN Jacobian matrix and compared); Márquez-Neila et al. (2017) trains DNs with explicit input-output constraints (the constraints are cast into JVPs and evaluated as part of the DN training procedure). Beyond those recent developments, JVPs have a long history of being employed to probe models and get precise understandings of the models’ properties e.g. input sensitivity, intrinsic dimension, conditioning (Gudmundsson et al., 1995; Bentbib & Kanber, 2015; Balestriero et al., 2020). To enable the deployment of these techniques in practical scenarios, it is crucial to evaluate JVPs as fast as possible and without relying on automatic differentiation, a feature primarily implemented in development libraries but often missing in software deployed in production.

In this paper, we propose a novel method to evaluate JVPs that is inspired by recent theoretical studies of DNs employing Continuous Piecewise Affine (CPA) nonlinearities and by Forward-mode Differentiation (FD). In short, the CPA property allows to specialize FD to be evaluate only by feeding two inputs to a DN while holding the DN internal state fixed, and combining the two produced outputs such that the final result corresponds to the desired JVP. Our solution offers the following benefits:

• •

  Speed: $2\times$ faster on average over the fastest alternative solution with empirical validation on $13$ different DN architectures (including Recurrent Neural Networks, ResNets, DenseNets, VGGs, UNets) and with speed-ups consistent across hardware.

• •

  Simple implementation: Implementing our solution for any DN requires to change only a few lines of code in total, regardless of the architecture being considered. This allows our method to be employed inside the architecture search loop if needed.

• •

  Does not employ automatic differentiation: Our solution can be implemented in any software/library including the ones that are not equipped with automatic differentiation, allowing deployment of the method on production specific hardware/software.

While our method is most simple to implement for CPA DNs (a constraint embracing most of the current state-of-the-art architectures), we also demonstrate how to employ it for smooth DN’s layers such as batch-normalization, which is CPA only during testing.

This paper is organized as follows: Sec. 2 reviews current JVPs solutions and thoroughly presents notations from CPA DNs, Sec. 3 presents our proposed solution where detailed implementation is provided in Sec. 3.2 and careful empirical validations are performed in Sec. 3.3; Sec. 4 presents further results on how JVPs are crucial to evaluate many insightful statistics of CPA DNs such as the largest eigenvalues of the DN Jacobian matrix.

Jacobian-Vector Products (JVPs). The JVP operation, also called Rop (short for right-operation), is abbreviated as ${\rm Rop}(f,\bm{x},v)$ and computes the directional derivative, with direction $\bm{u}\in\mathbb{R}^{D}$, of the multi-dimensional operator $f:\mathbb{R}^{D}\mapsto\mathbb{R}^{K}$ with respect to the input $\bm{x}\in\mathbb{R}^{D}$ as

Throughout our study, we will assume that $f$ is differentiable at $\bm{x}$, ${\bm{J}}_{f}(\bm{x})$ denotes the Jacobian of $f$ evaluated at $\bm{x}$. The ${\rm Rop}$ operator defines a linear operator from $\mathbb{R}^{D}$ to $\mathbb{R}^{K}$ with respect to the $\bm{u}$ argument, additional details can be found in Appendix A; throughout this paper we will use the terms JVP and Rop interchangeably. The vector-Jacobian product (VJP) operation, also called Lop, is abbreviated as ${\rm Lop}(f,\bm{x},\bm{v})$ with direction $\bm{v}\in\mathbb{R}^{K}$ and computes the adjoint directional derivative

For a thorough review of those operations, we refer the reader to Paszke et al. (2017); Baydin et al. (2017); Elliott (2018); throughout this paper we will use the terms VJP and Lop interchangeably.

In practice, evaluating the Rop given a differentiable operator $f$ can be done naturally through forward-mode automatic differentiation which builds ${\rm Rop}(f,\bm{x},\bm{u})$ while forwarding $\bm{x}$ through the computational graph of $f$. However, it can also be done by employing the backward-mode automatic differentiation (Linnainmaa, 1976) i.e. traversing the computational graph of $f$ from output to input, via multiple Lop calls via

where $K$ is the length of the output vector of the mapping $f$. Practical implementations of the above often resort to parallel computations of each $Lop$ forming the rows of the Jacobian matrix reducing computation time at the cost of greater memory requirements. Alternatively, an interesting computational trick recently brought to light in Jamie et al. (2017) allows to compute the $Rop$ from two $Lop$ calls only (instead of K) as follows

detailed derivation can be found in Appendix B. In short, the inner $Lop$ represents the mapping of $\bm{v}$ onto the linear operator ${\bm{J}}_{f}(\bm{x})^{T}$. Hence, ${\rm Lop}(f,\bm{x},\bm{v})^{T}$ being linear in $\bm{v}$, can itself be differentiated w.r.t. $\bm{v}$ and evaluated with direction $\bm{u}$ leading to the above equality. While this alternative strategy might not be of great interest for theoretical studies since the asymptotic time-complexity of ${\rm Lop}$ and ${\rm Rop}$ is similar (proportional to the number of operations $K$ (Griewank & Walther, 2008)). Which of the above strategies will provide the fastest computation in practice depends on the type of operations (input/output shapes) that are present in the computational graph. And while Rop and Lop can be combined, the general problem of optimal Jacobian accumulation (finding the most efficient computation) is NP-complete (Naumann, 2008). Note that current software/hardware in deep learning heavily focus on Lop as backpropagation (gradient based DN training strategy) is a special case of the latter (Rumelhart et al., 1986).

Deep Networks (DNs). A DN is an operator $f$ that maps an observation ¹¹1We consider vectors, matrices and tensors as flattened vectors $\bm{z}\in{\mathbb{R}}^{D}$ to a prediction $\bm{y}\in{\mathbb{R}}^{K}$. This prediction is a nonlinear transformation of the input $\bm{x}$ built through a cascade of linear and nonlinear operators (LeCun et al., 2015). Those operators include affine operators such as the fully connected operator (simply an affine transformation defined by weight matrix $\bm{W}^{\ell}$ and bias vector $\bm{v}^{\ell}$), convolution operator (with circulant $\bm{W}^{\ell}$), and, nonlinear operators such as the activation operator (applying a scalar nonlinearity such as the ubiquitous ReLU), or the max-pooling operator; precise definitions of these operators can be found in Goodfellow et al. (2016).

It will become convenient for our development to group the succession of operators forming a DN into layers in order to express the entire DN as a composition of $L$ intermediate layers $f^{\ell}$, $\ell=1,\dots,L$ as in

where each $f^{\ell}$ maps an input vector of length $D^{\ell-1}$ to an output vector of length $D^{\ell}$. We formally specify how to perform this grouping to ensure that any DN as a unique layer decomposition.

We also explicitly denote the up-to-layer $\ell$ mapping as

with initialization $\bm{z}^{0}(\bm{x})\triangleq\bm{x}$, those intermediate representations $\bm{z}^{1},\dots,\bm{z}^{L-1}$ are denoted as feature maps.

Continuous Piecewise Affine (CPA) DNs. Whenever a DN $f$ employs CPA nonlinearities throughout its layers, the entire input-output mapping is a (continuous) affine spline with a partition $\Omega$ of its domain (the input space of the DN) and a per-region affine mapping (for more background on splines we refer the reader to De Boor et al. (1978); Schumaker (2007); Bojanov et al. (2013); Pascanu et al. (2013)) such that

where the per-region slope and bias parameters are a function of the DN architecture and parameters. Given an observation $\bm{x}\in\omega$, the affine parameters of region $\omega$ from (6) can be obtained explicitly as follows.
1. For each layer $f^{\ell}$ (recall Def 1) compose the possibly multiple linear operator into a single one with slope matrix $\bm{W}^{\ell}$ and bias $\bm{b}^{\ell}$ (any composition of linear operators can be expressed this way by composition and distributive law) and denote the layer’s pre-activation as

in the absence of linear operators in a layer, set $\bm{W}^{\ell}$ as the identity matrix and $\bm{b}^{\ell}$ as the zero-vector.
2. For each layer, encode the states of the layer’s activation operator into the following diagonal matrix

where $\eta>0$ (leaky-ReLU), $\eta=0$ (ReLU), or $\eta=-1$ (absolute value); for the max-pooling operator, see Appendix C; $[.]$ is used to access a specific entry of a matrix/vector, for clarity we will often omit the double superscript in (8) and use $\bm{Q}\left(\bm{h}^{\ell}(\bm{x})\right)$.
3. Combining (7) and (8) to obtain

with $\bm{A}_{\omega}\in\mathbb{R}^{K\times D}$ and $\bm{b}_{\omega}\in\mathbb{R}^{D}$.
Recently, many studies have successfully exploited the CPA property of current DNs to derive novel insights and results such as bounding the number of partition regions (Montufar et al., 2014; Arora et al., 2016; Rister & Rubin, 2017; Serra et al., 2018; Hanin & Rolnick, 2019), characterizing the geometry of the partitioning (Balestriero et al., 2019; Robinson et al., 2019; Gamba et al., 2020); and quantifying the expressive power of CPA DNs as function approximators (Unser, 2019; Bartlett et al., 2019; Nadav et al., 2020; Grigsby & Lindsey, 2020). We propose to follow the same path and exploit the CPA property of DNs to obtain a fast and architecture independent method to evaluate Jacobian-vector products with DNs that do not require any flavor of automatic differentiation. Our method will allow to speed-up recent methods employing JVPs and to ease their deployment.

As we mentioned, the core of our solution resides on the assumptions that the employed DN’s nonlinearities are CPA, we first describe our solution, its implementation and then conclude with thorough empirical validations and solutions to extend our technique to smooth DNs.

We now demonstrate how (fast) JVPs can be used to compute some important quantities that have been shown to characterize DNs e.g. the largest eigenvalues and eigenvectors of $\bm{A}_{\omega}$, the per-region slope matrix that produces the DN output (recall (6)). Recall that for large models, extracting $\bm{A}_{\omega(\bm{x})}$ entirely, and then performing any sort of analysis is not feasible ($196608$ by $196608$ matrix for datasets a la ImageNet (Deng et al., 2009)). It is thus crucial to evaluate the desired quantities only through JVP calls.

Automatic differentiation and backpropagation are behind most if not all recent successes in deep learning, yielding not only state-of-the-art results but also fueling the discoveries of new architectures, normalizations, and the likes. This first wave of research almost exclusively relied on backward differentiation (gradient computations) but recently, many techniques further pushing DN performances have been relying on Jacobian-vector products (JVPs). Current JVPs evaluations rely on automatic differentiation making them slower and less memory efficient than direct evaluation which in turn requires to be implemented by hand based on the architecture employed. We demonstrated in this paper that current DNs, which employ nonlinearities such as (leaky-)ReLU and max-pooling, can have their JVPs evaluated directly without the need for automatic differentiation and without the need to hard-code each architecture only by changing a few lines of code in their implementation. We validated our proposed method against alternative implementations and demonstrated significant speed-ups of a factor of almost $2$ on average over $13$ different architectures and hardware. We also demonstrated that fast JVPs will help in the development of novel analysis and theoretical techniques for CPA DNs that require evaluations of quantities such as eigenvectors/eigenvalues of the per-region slope matrices that benefit from faster JVP evaluations.

This work was supported by NSF grants CCF-1911094, IIS-1838177, and IIS-1730574; ONR grants N00014-18-12571, N00014-20-1-2787, and N00014-20-1-2534; AFOSR grant FA9550-18-1-0478; and a Vannevar Bush Faculty Fellowship, ONR grant N00014-18-1-2047.