Fast Jacobian-Vector Product for Deep Networks

TL;DR

This paper introduces a fast, architecture-agnostic method for computing Jacobian-vector products in deep networks with piecewise affine nonlinearities, significantly reducing computation time and simplifying implementation.

Contribution

The authors propose a novel, automatic differentiation-free technique for rapid JVP computation applicable to any deep network with piecewise affine activations, improving efficiency and ease of use.

Findings

On average 2x faster than existing methods

Applicable to 13 different deep network architectures

Requires minimal code modifications for deployment

Abstract

Jacobian-vector products (JVPs) form the backbone of many recent developments in Deep Networks (DNs), with applications including faster constrained optimization, regularization with generalization guarantees, and adversarial example sensitivity assessments. Unfortunately, JVPs are computationally expensive for real world DN architectures and require the use of automatic differentiation to avoid manually adapting the JVP program when changing the DN architecture. We propose a novel method to quickly compute JVPs for any DN that employ Continuous Piecewise Affine (e.g., leaky-ReLU, max-pooling, maxout, etc.) nonlinearities. We show that our technique is on average $2\times$ faster than the fastest alternative over $13$ DN architectures and across various hardware. In addition, our solution does not require automatic differentiation and is thus easy to deploy in software, requiring only the modification of a few lines of codes that do not depend on the DN architecture.