Scheduled Jacobian Chaining

TL;DR

This paper introduces a new approach for efficiently computing Jacobian matrices in programs by optimizing the chain of matrix multiplications, leveraging parallelism and heuristic algorithms to improve performance.

Contribution

It formulates Jacobian computation as a matrix chain product problem and develops a dynamic programming heuristic that incorporates scheduling and parallelism considerations.

Findings

The heuristic performs close to the optimal solution in benchmarks.

Parallelism in Jacobian chain computations can be effectively exploited.

The approach reduces computational complexity for large Jacobian matrices.

Abstract

This paper addresses the efficient computation of Jacobian matrices for programs composed of sequential differentiable subprograms. By representing the overall Jacobian as a chain product of the Jacobians of these subprograms, we reduce the problem to optimizing the sequence of matrix multiplications, known as the Jacobian Matrix Chain Product problem. Solutions to this problem yield "optimal bracketings", which induce a precedence-constraint scheduling problem. We investigate the inherent parallelism in the solutions and develop a new dynamic programming algorithm as a heuristic that incorporates the scheduling. To assess its performance, we benchmark it against the global optimum, which is computed via a branch-and-bound algorithm.