An illustrated guide to automatic sparse differentiation

Summary

Contents $$ \newcommand{\colorf}[1]{\textcolor{RoyalBlue}{#1}} \newcommand{\colorh}[1]{\textcolor{RedOrange}{#1}} \newcommand{\colorg}[1]{\textcolor{PineGreen}{#1}} \newcommand{\colorv}[1]{\textcolor{VioletRed}{#1}} \def\sR{\mathbb{R}} \def\vx{\mathbf{x}} \def\vv{\mathbf{v}} \def\vb{\mathbf{e}} \newcommand{\vvc}[1]{\colorv{\vv_{#1}}} \newcommand{\vbc}[1]{\colorv{\vb_{#1}}} \newcommand{\dfdx}[2]{\frac{\partial f_{#1}}{\partial x_{#2}}(\vx)} \newcommand{\J}[2]{J_{#1}(#2)} \def\Jf{\J{f}{\vx}} \def\Jg{\J{g}{\vx}} \def\Jh{\J{h}{g(\vx)}} \def\Jfc{\colorf{\Jf}} \def\Jgc{\colorg{\Jg}} \def\Jhc{\colorh{\Jh}} \newcommand{\D}[2]{D{#1}(#2)} \def\Df{\D{f}{\vx}} \def\Dg{\D{g}{\vx}} \def\Dh{\D{h}{g(\vx)}} \def\Dfc{\colorf{\Df}} \def\Dgc{\colorg{\Dg}} \def\Dhc{\colorh{\Dh}} $$ First-order optimization is ubiquitous in machine learning (ML) but second-order optimization is much less common. The intuitive reason is that high-dimensional vectors (gradients) are cheap, whereas high-dimensional matrices (Hessians) are expensive. Luckily, in numerous applications of ML to science or engineering, Hessians and Jacobians exhibit sparsity: most of their coefficients are known to be zero. Leveraging this sparsity can vastly accelerate automatic differentiation (AD) for Hessians and Jacobians, while decreasing its memory requirements . Yet, while traditional AD is available in many high-level programming languages like Python and Julia , automatic sparse differentiation (ASD) is not as widely used. One reason is that the underlying theory was developed in the AD community, outside of the ML research ecosystem. With this blog post, we aim to shed light on the inner workings of ASD, bridging the gap between the ML and AD communities by presenting well established techniques from the latter field. We start out with a short introduction to traditional AD, covering the computation of Jacobians in both forward and reverse mode. We then dive into the two primary components of ASD: sparsity pattern de...