Automatic Differentiation

Computes derivatives by propagating tangent values alongside the primal computation from inputs to outputs. For a function with n inputs and m outputs, forward mode requires n passes to compute a full Jacobian.

• Mechanism: Dual numbers carry both value and derivative component through each operation
• Efficiency: Ideal when n < m (few inputs, many outputs)
• Use case: Computing directional derivatives or Jacobian-vector products for sensitivity analysis in molecular property prediction
• Memory cost: Proportional to the number of operations, not the number of inputs

Computes derivatives by propagating adjoint values from outputs backward through the computation graph to inputs. This is the algorithm underlying backpropagation in deep learning frameworks like PyTorch and JAX.

• Mechanism: Records a computational graph (tape) during forward pass, then traverses it in reverse
• Efficiency: Ideal when n >> m (many inputs, few outputs) — exactly the case for neural network potentials where millions of parameters map to a single scalar energy
• Memory cost: Must store all intermediate activations, leading to the checkpointing trade-off in large molecular dynamics simulations
• Gradient of energy: Computes forces on all atoms in a single backward pass

AD operates on a directed acyclic graph (DAG) where nodes represent variables and edges represent functional dependencies. Each elementary operation (addition, multiplication, sin, exp) has a known derivative rule.

• Dynamic graphs: PyTorch builds the graph on-the-fly during execution, enabling flexible control flow in equivariant neural networks
• Static graphs: TensorFlow compiles the graph before execution, enabling aggressive operator fusion and memory optimization
• Tracing vs. source transformation: Operator overloading intercepts each operation at runtime, while source-to-source AD rewrites the program before execution
• Custom gradients: Users can register hand-derived derivatives for operations like spherical harmonics where automatic rules may be suboptimal

Unlike numerical differentiation (finite differences) which introduces truncation and round-off errors, AD computes derivatives to machine precision by applying the chain rule to elementary operations.

• Numerical differentiation error: f'(x) ≈ [f(x+h) - f(x)]/h suffers from O(h) truncation error and catastrophic cancellation as h → 0
• Symbolic differentiation: Produces exact expressions but suffers from expression swell — the derivative formula can be exponentially larger than the original function
• AD advantage: Delivers exact derivatives with only a small constant factor overhead (typically 2-5x the original function cost)
• Critical for force matching: Training neural network potentials requires accurate forces; numerical noise in gradients would corrupt the learned potential energy surface

Reverse-mode AD stores all intermediate values during the forward pass, creating a memory bottleneck for deep networks and long molecular dynamics trajectories. Gradient checkpointing trades computation for memory.

• Checkpointing strategy: Discard intermediate activations, then recompute them on-demand during the backward pass from stored checkpoints
• Optimal schedule: The binomial checkpointing algorithm achieves O(log n) memory with O(n log n) recomputation
• Application: Training equivariant neural networks on large molecular systems where storing per-atom features for thousands of atoms exceeds GPU memory
• JAX's approach: Functional transforms like jax.checkpoint automatically apply checkpointing to arbitrary functions

AD can be applied recursively to compute Hessians (second derivatives) and higher-order tensors required for vibrational analysis and transition state optimization in quantum chemistry.

• Forward-over-reverse: Combine forward and reverse modes to efficiently compute Hessian-vector products without materializing the full Hessian matrix
• Application: Computing harmonic vibrational frequencies from a neural network potential by differentiating the energy gradient
• Force constants: The Hessian of the potential energy with respect to atomic coordinates yields the force constant matrix for normal mode analysis
• Framework support: JAX provides jax.hessian, PyTorch provides torch.autograd.functional.hessian

Computes derivatives by propagating tangent values forward through the computational graph alongside the primal calculation. Dual numbers are used to carry both value and derivative, making this mode efficient for functions with few inputs and many outputs.

• Mechanism: Applies the chain rule from inputs to outputs
• Complexity: O(N) in number of inputs
• Use case: Computing Jacobian-vector products for directional derivatives
• Limitation: Inefficient for gradient computation in neural networks with millions of parameters

The workhorse of deep learning, also known as backpropagation. Computes gradients by first executing a forward pass to record all intermediate values, then propagating adjoint sensitivities backward from the output to all inputs simultaneously.

• Mechanism: Applies the chain rule from outputs to inputs
• Complexity: O(N) in number of outputs, independent of input count
• Use case: Training neural network potentials on forces and energies
• Key advantage: Computes full gradient of a scalar loss with respect to millions of parameters in a single backward pass

A directed acyclic graph representing the sequence of elementary operations in a function. Each node stores an operation and its intermediate value, while edges define data dependencies. AD traverses this graph to apply the chain rule.

• Static graphs: Defined once before execution, enabling aggressive compiler optimizations (TensorFlow 1.x)
• Dynamic graphs: Constructed on-the-fly during execution, allowing native control flow (PyTorch, JAX)
• Wengert tape: The linearized record of operations generated during the forward pass that reverse-mode AD traverses backward

The fundamental building block of reverse-mode AD. Instead of computing the full Jacobian matrix explicitly, a VJP efficiently computes the product of a row vector with the Jacobian. This avoids the prohibitive O(n²) memory cost of materializing the full Jacobian.

• Pullback function: The VJP operator associated with each primitive operation
• Gradient accumulation: VJPs from multiple downstream nodes are summed at branch points
• Efficiency: Reduces gradient computation from O(n²) to O(n) for scalar-output functions

AD can be applied recursively to compute second derivatives like Hessians and Laplacians, essential for quantum chemistry applications such as vibrational frequency analysis and transition state optimization.

• Forward-over-reverse: Computes Hessian-vector products efficiently without materializing the full Hessian matrix
• Double backpropagation: Applying reverse-mode AD twice to obtain second-order derivatives
• Use in NNPs: Computing force constants and phonon spectra from neural network potentials
• Frameworks: JAX's grad, jacfwd, and jacrev compose naturally for arbitrary-order differentiation

An AD implementation strategy that parses source code into an intermediate representation and generates new code that computes derivatives. This approach enables compile-time optimization and produces efficient derivative code without runtime overhead.

• Tools: Tapenade, Zygote (Julia), Enzyme (LLVM-level)
• Advantage: Can differentiate through control flow, loops, and recursion
• Differentiable programming: The paradigm where entire programs are written to be end-to-end differentiable
• Application: Differentiating through entire quantum chemistry workflows, including SCF convergence