CHESSFAD: Chunked Hessian AD

• CHESSFAD is an automatic differentiation methodology that computes Hessians and Hessian–vector products using a chunked forward-mode approach.
• It employs a dual-number design with operator overloading to enable fine-grained parallelism on both CPUs and GPUs while minimizing memory usage.
• Benchmarks on functions like Rosenbrock and Ackley show significant speedups over traditional AD libraries, especially on NVIDIA GPUs.

CHESSFAD (Chunked HESSian using Forward-mode AD) is an automatic differentiation (AD) methodology and software library for efficient, parallel computation of Hessians and Hessian–vector products, specifically targeting applications where large numbers of such products must be computed across many data points simultaneously on modern accelerators such as NVIDIA GPUs. The core design principles of CHESSFAD leverage forward-mode AD with a chunked approach to enable fine-grained parallelism at multiple levels and minimize memory usage by avoiding explicit materialization of the full Hessian. As a lightweight, header-based C++ implementation, CHESSFAD is portable across both CPUs and GPUs (Ranjan et al., 2024).

1. Mathematical Foundations and Dual-Number Design

Let $f:\mathbb{R}^n\to\mathbb{R}$, with Hessian $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$. The principal task is to evaluate $H$ or compute products $H v$ for arbitrary $v \in \mathbb{R}^n$.

CHESSFAD employs forward-mode AD extended to second-order by using a custom data structure, $h$Dual, which encapsulates the function value, relevant first-order derivatives, and a chunk of second-order mixed partials: $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$ with $\{j_1,\ldots,j_c\}$ denoting a contiguous chunk of columns in row $i$, and $c$ the chunk size. By operator overloading, arithmetic on $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$0Duals propagates all required higher-order derivatives. For example, for $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$1, $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$2 as above: $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$3 By initializing appropriate slots to Kronecker deltas, a single invocation of the dual-number-templated $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$4 yields multiple entries per row of $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$5 in a single pass.

Hessian–vector products $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$6 are formed incrementally: for each computed chunk $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$7, the partial dot product $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$8 is accumulated into $$H(x) = [\partial_i \partial_j f(x)]_{i,j = 1,\dots,n}$$9, with results streamed and discarded immediately to avoid full Hessian materialization.

2. Parallelization Strategies: Row and Chunk Concurrency

CHESSFAD exposes two primary, independent levels of parallelism:

• Row-parallelism ($H$0): Each row $H$1 of $H$2 can be computed independently, since $H$3 and $H$4 are mutually independent for $H$5.
• Chunk-parallelism ($H$6): Each row $H$7 is divided into $H$8 contiguous chunks of size $H$9, enabling concurrent computation of these chunks.

The GPU implementation assigns each (instance, row, chunk) tuple to a GPU thread, as in the following representative CUDA code snippet: c++ __global__ void chessVecKernel(double* x, double* v, double* out, int n, int csize) { int eid = blockIdx.x; int tid = threadIdx.x; int nchunk = n / csize; int i = tid / nchunk; int j = tid % nchunk; CHUNK_INIT(y, &x[eid*n], i, j, n, csize); auto temp = f<hDual<csize>>(y); double partial = 0; for(int l=0; l < csize; ++l) partial += temp.v[csize+2+l] * v[eid*n + j*csize + l]; __shared__ double sprod[MAX_N] [MAX_NCHUNKS]; sprod[i] [j] = partial; __syncthreads(); if(j==0){ double sum=0; for(int k=0; k<nchunk; ++k) sum += sprod[i] [k]; out[eid*n+i]=sum; } } $\{j_1,\ldots,j_c\}$5c++ template<int csize> struct hDual { double v[2*csize + 2]; // operator overloads }; Initialization macros (CHUNK_INIT) configure each $H v$0Dual so that the required derivative slots (for a given row $H v$1, chunk $H v$2) are set to $H v$3.

On CUDA, thread-local or register-resident $H v$4Dual arrays store input variables for each thread. Intermediate results are accumulated into a shared memory array per block for efficient in-block reductions, and dot products are streamed to output buffers to minimize global memory traffic and avoid constructing the full Hessian tensor.

4. Empirical Performance and Benchmarking

Performance was evaluated against the CPU automatic differentiation library autodiff using three standard benchmark functions generalized to $H v$5 variables: Rosenbrock, Ackley, and Fletcher–Powell. Each experiment consisted of $H v$6 independent Hessian–vector products on CPU and $H v$7 million on GPU. CHESSFAD delivered speedups over autodiff as follows:

Function	CPU Speedup	A100 GPU Speedup ($H v$8)	A100 GPU Speedup ($H v$9)
Rosenbrock	~20%	$v \in \mathbb{R}^n$0	$v \in \mathbb{R}^n$1
Ackley	~5%	$v \in \mathbb{R}^n$2	$v \in \mathbb{R}^n$3
Fletcher–Powell	~49%	$v \in \mathbb{R}^n$4	$v \in \mathbb{R}^n$5

Kernel-only speedup drops sublinearly with respect to $v \in \mathbb{R}^n$6 due to quadratic scaling, yet remains significant up to $v \in \mathbb{R}^n$7. The time to perform $v \in \mathbb{R}^n$8 sequential computations is roughly equivalent to $v \in \mathbb{R}^n$9 in parallel on the A100 GPU. This suggests strong GPU scaling headroom before parallel efficiency saturates (Ranjan et al., 2024).

5. Computational Complexity and Chunk Size Trade-Offs

Let $h$0 denote the number of multiplications and $h$1 the number of additions in the scalar code for $h$2, with $h$3 variables and chunk size $h$4.

• Without symmetry (CHUNK-HESS):
  • Number of function calls: $h$5
  • Each $h$6Dual–multiply: $h$7 multiplies, $h$8 adds
  • Each $h$9Dual–add: $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$0 adds
  • Total scalar multiplies: $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$1
  • Total scalar adds: $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$2
• With symmetry (SCHUNK-HESS):
  • Only $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$3 chunks are computed
  • Scalar-multiply count: $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$4
  • The optimal chunk size for fixed $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$5 is $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$6

By streaming out each chunk's result directly into the dot product, memory bandwidth is reduced and the need to store the full Hessian is removed. The chunk size $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$7 introduces a trade-off: larger $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$8 reduces the number of AD passes at the expense of increased per-pass memory, while smaller $$hDual(u) = \langle u,\, \partial_i u,\, \partial_{j_1}u,\,\ldots,\,\partial_{j_c}u,\,\partial^2_{i j_1}u,\,\ldots,\,\partial^2_{i j_c}u \rangle$$9 lightens per-kernel memory but demands more passes. Empirically, chunk sizes in the range $\{j_1,\ldots,j_c\}$0–$\{j_1,\ldots,j_c\}$1 offer good balance for $\{j_1,\ldots,j_c\}$2 up to $\{j_1,\ldots,j_c\}$3 on GPU (Ranjan et al., 2024).

6. Applicability and Scope

CHESSFAD is applicable in scientific and engineering settings requiring high-throughput, batched computation of Hessian–vector products, such as large-scale optimization, inverse problems, and high-order sensitivity analyses in finite element modeling. The ability to exploit both row- and chunk-level parallelism enables CHESSFAD to efficiently harness modern multi-core CPUs and GPUs.

A plausible implication is that CHESSFAD's chunked forward-mode scheme could be adapted for larger $\{j_1,\ldots,j_c\}$4 by dynamically selecting chunk sizes and leveraging symmetry when available. Its low memory overhead and advantage in not constructing the full Hessian matrix mark it as particularly suitable for sparse or structured problems where only directional second derivatives are required.