Naive vector-Jacobian product vs JAX.vjp | Benchmark comparison

Vector-Jacobian products evaluate the left multiplication of a (cotangent) vector to a Jacobian matrix. This video will show that the intrinsic of the JAX deep learning framework is faster than a naive implementation for sparse Jacobians.

Sparse Jacobians are common in Scientific Computing. Unfortunately, JAX does not support sparse array storage formats or sparse Automatic Differentiation. Hence, obtaining a Jacobian matrix of such functions will allocate a dense array. This quickly becomes intractable/infeasible for higher dimensional problems due to the unnecessary computation and storage of the many zeros.

Luckily, often the full Jacobian is of no relevance, but its effect on vectors that are either multiplied from the right (Jacobian-vector product) or from the left (vector-Jacobian product). That's also what is sufficient for the iterative solution to linear systems of equations that involve these matrices. Such systems for instance arise in the adjoint-mode sensitivity analysis over nonlinear systems of equations.

-------

📝 : Check out the GitHub Repository of the channel, where I upload all the handwritten notes and source-code files (contributions are very welcome): https://github.com/Ceyron/machine-lea...

📢 : Follow me on LinkedIn or Twitter for updates on the channel and other cool Machine Learning & Simulation stuff:   / felix-koehler   and   / felix_m_koehler  

💸 : If you want to support my work on the channel, you can become a Patreon here:   / mlsim  

🪙: Or you can make a one-time donation via PayPal: https://www.paypal.com/paypalme/Felix...

-------

Timestamps:
00:00 Intro
00:27 Vector-Valued Function with small locality
02:24 Investigating the sparsity structure of Jacobian
03:27 Function for naive vJp
04:59 Function for clever vJp with JAX intrinsic
06:44 Kicking off the lazy JIT compilation
07:45 Benchmarking naive and clever approach
08:18 Comparison & Discussion
09:04 vJp as part of adjoint sensitivity methods
09:21 Outro