Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Scheduled Jacobian Chaining

Simon Märtens, Uwe Naumann

TL;DR

The paper reframes the Jacobian accumulation problem for differentiable programs as a Jacobian Chain Product with the goal of minimizing the number of scalar $\text{fma}$ operations through optimal bracketing. It then introduces Scheduling-aware Dynamic Programming to solve a precedence-constrained, in-tree variant of the problem, producing near-optimal elimination sequences that exploit parallelism under a given machine count $m$ and memory constraints. Although the underlying scheduling problem is NP-hard, the DP approach remains a practical heuristic, and its effectiveness is validated against a Branch-and-Bound global optimum using large-scale statistical benchmarks. The work provides a reference implementation, case-study guidance, and avenues for future work in runtime tape partitioning, communication costs, and sparsity-aware extensions, underscoring the practical impact for high-performance AD in computational engineering.

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.

Scheduled Jacobian Chaining

TL;DR

The paper reframes the Jacobian accumulation problem for differentiable programs as a Jacobian Chain Product with the goal of minimizing the number of scalar operations through optimal bracketing. It then introduces Scheduling-aware Dynamic Programming to solve a precedence-constrained, in-tree variant of the problem, producing near-optimal elimination sequences that exploit parallelism under a given machine count and memory constraints. Although the underlying scheduling problem is NP-hard, the DP approach remains a practical heuristic, and its effectiveness is validated against a Branch-and-Bound global optimum using large-scale statistical benchmarks. The work provides a reference implementation, case-study guidance, and avenues for future work in runtime tape partitioning, communication costs, and sparsity-aware extensions, underscoring the practical impact for high-performance AD in computational engineering.

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.
Paper Structure (22 sections, 2 theorems, 32 equations, 13 figures, 3 tables)

This paper contains 22 sections, 2 theorems, 32 equations, 13 figures, 3 tables.

Table of Contents

  1. Introduction.
  2. Algorithmic differentiation (AD).
  3. Jacobian Chaining.
  4. Elimination Sequences.
  5. Scheduled Jacobian Chaining.
  6. The Scheduling Problem.
  7. DP with Scheduling.
  8. Optimality.
  9. Correctness and Determinism.
  10. Statistical Assessment.
  11. Implementation.
  12. Conclusion and Future Work.
  13. Case Studies.
  14. Communication Costs.
  15. Sparse Jacobian Computations.
  16. ...and 7 more sections

Key Result

Theorem 2.1

Figure eqn:scheduled_chaining with $t=\infty$ solves the Scheduled Limited-Memory Matrix-Free Dense Jacobian Chain Product Bracketing problem.

Figures (13)

  • Figure 1.1: Example of a DAG labeled with the local derivatives and the corresponding primal single assignment code of the function $f: \mathbb{R}^2 \rightarrow \mathbb{R}$ with $\bm{x} \mapsto y \equiv f(\bm{x}) = \sin(x_0 \cdot x_1)$.
  • Figure 1.2: DP recurrence relation for Dense Jacobian Chain Product Bracketing
  • Figure 2.1: Task dependencies for example elimination sequence at the end of Section \ref{['ssec:elimination_sequences']}.
  • Figure 2.2: DP relation for Scheduled Limited-Memory Matrix-Free Dense Jacobian Chain Product Bracketing.
  • Figure 2.3: Task dependencies for example elimination sequence with 3 available machines.
  • ...and 8 more figures

Theorems & Definitions (4)

  • Theorem 2.1
  • Proof 1
  • Theorem 2.2
  • Proof 2