A C++ implementation of the discrete adjoint sensitivity analysis method for explicit adaptive Runge-Kutta methods enabled by automatic adjoint differentiation and SIMD vectorization

TL;DR

This work presents a C++ library implementing the discrete adjoint sensitivity analysis for adaptive explicit Runge-Kutta methods, enhanced by automatic adjoint differentiation and SIMD vectorization. By avoiding a black-box AD approach to the entire solver and employing a targeted discrete adjoint formulation, it achieves efficient vector-Jacobian product evaluations via AAD and parallelization through AVX. The authors compare against established tools, showing competitive or superior performance in large-parameter regimes, while adopting a memory-efficient checkpointing strategy. The approach enables efficient gradient computation for optimization and parameter estimation in ODE-based models, with potential extensions to more integrators and second-order adjoints. Overall, the work advances scalable sensitivity analysis for complex, parameter-rich dynamical systems using modern CPU architectures.

Abstract

A C++ library for sensitivity analysis of optimisation problems involving ordinary differential equations (ODEs) enabled by automatic differentiation (AD) and SIMD (Single Instruction, Multiple data) vectorization is presented. The discrete adjoint sensitivity analysis method is implemented for adaptive explicit Runge-Kutta (ERK) methods. Automatic adjoint differentiation (AAD) is employed for efficient evaluations of products of vectors and the Jacobian matrix of the right hand side of the ODE system. This approach avoids the low-level drawbacks of the black box approach of employing AAD on the entire ODE solver and opens the possibility to leverage parallelization. SIMD vectorization is employed to compute the vector-Jacobian products concurrently. We study the performance of other methods and implementations of sensitivity analysis and we find that our algorithm presents a small advantage compared to equivalent existing software.

Figures (5)

• Figure 1: Simplified computational graphs for the intermediate states and parameters. Each node corresponds to an intermediate state, and a directed edge signifies that the child intermediate state depends explicitly on the parent intermediate state. To compute the adjoint variables of the parent nodes, we require the adjoint variables of the children nodes.
• Figure 2: Plot of the solution error ($x$-axis) against the sensitivity error ($y$-axis) for the Van der Pol equations for multiple implementations of adjoint sensitivity analysis.
• Figure 3: Average execution clock time in milliseconds plotted against the global sensitivity error for several methods of sensitivity analysis. The average was taken over $30$ executions. The vertical lines represent a standard deviation. The dashed lines correspond to a linear regression of the data to the power law $t \sim \epsilon^{m}$ ignoring the data corresponding to the largest three error values.
• Figure 4: The average execution time (in milliseconds) is plotted against the global sensitivity error for five different implementations of adjoint sensitivity analysis applied to GLV systems of varying sizes. Vertical error bars indicate one standard deviation. Dashed lines represent linear regressions of the data following the power law $t \sim \epsilon^{m}$, with the slope of each line provided in the plot legend. Data for DENSERKS at $N=200$ is absent due to excessive execution time.
• Figure 5: Plot of execution clock time in milliseconds against the number of parameters in log-log scale. The dashed lines correspond to linear regressions to the power law $t \sim P^{m}$.