Fast Jacobian-Vector Product for Deep Networks
This addresses a bottleneck for researchers and practitioners in deep learning who rely on JVPs for applications like optimization and adversarial analysis, offering a significant speed improvement.
The paper tackles the computational expense of Jacobian-vector products (JVPs) in deep networks by proposing a novel method for networks with Continuous Piecewise Affine nonlinearities, achieving on average 2x faster computation than alternatives across 13 architectures and hardware.
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.