Overview
Science progresses by fitting models to observations and adjusting parameters until predictions match reality. When that fitting process can propagate gradients from a loss function back through every computational step, including numerical solvers, physics engines, and molecular dynamics integrators, the entire simulation becomes a differentiable program. Parameters that once required expensive grid searches or finite-difference approximations can now be optimized with the same gradient-based methods that train neural networks.
This chapter builds the foundations of differentiable programming for scientific discovery. We begin with the mechanics of automatic differentiation (forward mode, reverse mode, and their composition), then introduce JAX's functional transformation model that makes differentiation a first-class operation alongside compilation, vectorization, and parallelism. We extend these ideas to differentiable physics simulations and molecular dynamics, culminating in a hands-on recipe that calibrates a Lennard-Jones force field by backpropagating through a molecular dynamics trajectory.
The gradient-based optimization machinery developed here feeds directly into Chapter 43: Scientific Simulation, where we build larger simulation systems, and Chapter 45: Optimization for Discovery, where we combine gradient information with Bayesian and evolutionary strategies. The concept of differentiating through computational graphs connects back to the Scientific Machine Learning ideas in Chapter 33, where neural networks embed physical constraints directly into their architectures.
Prerequisites
This chapter assumes familiarity with multivariable calculus (partial derivatives, chain rule, Jacobians), linear algebra (matrix-vector products, transposes), and Python programming with NumPy. Prior exposure to gradient-based optimization (gradient descent, learning rates) from Chapter 26 is helpful. No prior JAX experience is required; we build up from first principles. For a calculus refresher, see Appendix A: Mathematical Foundations.
Learning Outcomes
- Explain forward-mode and reverse-mode automatic differentiation, their computational costs, and when to use each.
- Compose JAX's functional transforms (
grad,jit,vmap,pmap) to build efficient differentiable programs. - Build differentiable physics simulations using adjoint methods and ODE solvers.
- Implement a differentiable Lennard-Jones force field and calibrate its parameters against target observables by backpropagating through a molecular dynamics trajectory.
- Integrate differentiable programming components into the Discovery Workbench for automated parameter optimization.
Sections
42.1 Automatic Differentiation
Forward mode, reverse mode, and mixed-mode AD. Dual numbers, computational graphs, and the chain rule mechanized. Why reverse mode costs $O(1)$ per output regardless of input dimension, and how to derive it as transposed Jacobian-vector products.
42.2 JAX Functional Transforms
JAX's composable transformation model: grad, jit, vmap, and pmap. Pytrees, pure functions, and the functional programming discipline. Equinox, Flax, and Optax for building and training differentiable models.
42.3 Differentiable Physics and Molecular Dynamics
Adjoint methods for ODE-constrained optimization. Differentiable physics engines with JAX-MD and diffrax. The adjoint derivation for backpropagating through time-stepped simulations without storing intermediate states.
42.4 Building a Differentiable Force Field
Recipe: calibrate a Lennard-Jones force field in JAX. Backpropagate through a molecular dynamics trajectory to match radial distribution functions, using reverse-mode AD, vmap, and jit for performance.
Bibliography
Foundational Papers
One of the earliest descriptions of automatic differentiation, introducing the concept of mechanically applying the chain rule to computer programs.
The first description of reverse-mode automatic differentiation, later rediscovered as backpropagation for neural network training.
Introduced the adjoint method for backpropagating through ODE solvers with constant memory, enabling deep learning with continuous-depth models.
The JAX-MD paper, demonstrating how JAX's automatic differentiation enables end-to-end differentiable molecular dynamics simulations.
Books and Surveys
The definitive survey of AD for machine learning, covering forward mode, reverse mode, higher-order derivatives, and implementation strategies.
The standard textbook on automatic differentiation, with rigorous treatment of forward and reverse modes, checkpointing, and complexity analysis.
A comprehensive treatment of neural ODEs, SDEs, and CDEs, including the diffrax library for differentiable equation solving in JAX.
Tools and Libraries
Google's composable function transformation library: automatic differentiation, JIT compilation, vectorization, and parallelism on CPU/GPU/TPU.
A JAX library for neural networks and parameterized models that embraces JAX's functional paradigm using Python classes as pytrees.
Google's neural network library for JAX, providing the nnx module for building and training models with automatic state management.
A gradient processing and optimization library for JAX, offering composable gradient transformations (Adam, SGD, clipping, scheduling).
Numerical differential equation solvers in JAX with full autodiff support, enabling differentiable ODE/SDE simulation.
End-to-end differentiable molecular dynamics in JAX, supporting custom potentials, neighbor lists, and GPU-accelerated simulations.