Differentiable Simulation & Automatic Differentiation

• Differentiable simulation is a modeling approach that makes physical and engineering simulators fully differentiable with respect to parameters using automatic differentiation.
• It leverages both forward-mode and reverse-mode AD to compute exact derivatives through complex numerical solvers and computational graphs.
• This method accelerates gradient-based optimization in diverse applications such as optics calibration, PDE-constrained inversion, and agent-based control.

Differentiable simulation is the strategy of constructing physical, engineering, or agent-based simulations so that outputs (and derived objectives) are fully differentiable with respect to model, control, or design parameters. This enables the use of gradient-based optimization and inference for high-dimensional parameter spaces where classical black-box or finite-difference methods are inefficient or insufficiently precise. The core enabler is automatic differentiation (AD): a suite of algorithms and software frameworks that compute exact derivatives of arbitrary programs by embedding the chain rule directly into the computation graph, thereby propagating sensitivities from inputs through all algorithmic layers to the outputs. Modern differentiable simulators are now widespread, utilizing AD to calibrate optics and detectors, solve PDE-constrained inverse problems, optimize materials and quantum transport, design neural controllers, and enable scientific machine learning workflows across continuum, rarefied, and agent-based physics. This article reviews the mathematical and architectural foundations, key algorithmic building blocks, technical choices in AD mode and implementation, representative application domains, and performance considerations in differentiable simulation, with authoritative coverage of approaches from physical optics and molecular dynamics to photonic shape optimization, agent-based models, and high-performance computing.

1. Mathematical and Computational Principles

Differentiable simulation augments physics-based modeling with a computational graph structure that enables gradients to propagate through every elementary operation, including arithmetic on state variables, numerical solvers, and control flow. The general paradigm is to formulate a simulator as a sequence of composable, differentiable operators acting on an initial state and parameter vector $\theta$:

$$x_{1} = S_{1}(x_0; \theta), \quad x_2 = S_2(x_1; \theta), \ldots, \quad x_N = S_N(x_{N-1}; \theta)$$

where $S_j(\cdot)$ may be time-stepping, non-linear solvers, or physical transformations. The simulator yields observable(s) $y = G(x_N)$ and an objective $L(y, \theta)$, with $L$ scalar in classical inference/optimization tasks. AD systematically applies the chain rule to compute derivatives $\frac{\partial L}{\partial \theta}$ by propagating local sensitivities either in forward (tangent) or reverse (adjoint) mode.

Representative instances include:

• Physical optics: parameterizing phase, pixel gain, and source positions, then simulating layered physical transformations such as Fourier and non-linear optics (Desdoigts et al., 2024).
• PDE-constrained optimization: embedding discretized forward and adjoint solves, with sensitivities propagated through the full discretized system (Xu et al., 2019, Xue, 19 May 2025).
• Rigid-body and soft-body mechanics: differentiating through ODE/PDE integrators, explicit or implicit, with respect to masses, stiffnesses, or neural policy parameters (Millard et al., 2020, Rojas et al., 2021).
• Agent-based or control-loop models: making traditionally discontinuous simulation elements (branching, min/max, event logic) differentiable by smooth surrogates or smoothing abstractions (Andelfinger, 2021, Kreikemeyer et al., 2023).

The construction of a differentiable simulator mandates that all operators, including numerical solvers, geometric transforms, and even preprocessing/postprocessing, are implemented using AD-compatible primitives in frameworks such as JAX, PyTorch, TensorFlow, or C++ operator-overloading modules (Qianga et al., 26 Nov 2025). This "end-to-end" property is vital for ensuring gradients can be computed with respect to arbitrarily high-dimensional parameter vectors, including millions of detector gains, phase coefficients, or agent control weights.

2. Automatic Differentiation Modes and Implementation

Two principal modes of automatic differentiation are leveraged in differentiable simulation:

1. Forward-mode AD propagates directional derivatives alongside primal computations, making it the method of choice when the number of parameters is small relative to the number of outputs. It is efficiently implemented via dual numbers or truncated power series algebras ("TPSA"), as in multi-language ADVar modules for particle-in-cell and beam physics simulations (Qianga et al., 26 Nov 2025). The chain rule is encoded at the operator level by overloading arithmetic and mathematical functions, enabling transparent embedding into existing codes.

2. Reverse-mode AD (adjoint mode) computes the gradient of a scalar objective with respect to many parameters at a cost similar to one forward evaluation plus one backward pass. This is the dominant strategy for high-dimensional inverse problems and learning, realized via tape-based systems (recording execution traces), source code transformations (JAX/XLA, PyTorch's computation graphs), or explicit adjoint equations (adjoint ODE/PDE integration, as in modular physical simulations (Xu et al., 2019, Millard et al., 2020, Li et al., 2022, Xue, 19 May 2025)).

A key technical distinction arises in implicit solvers and time integration. When internal steps are defined via minimization or implicit equations (e.g., backward-Euler, equilibrium, or finite-elements), gradients require implicit differentiation. This is typically addressed by matrix-free solutions using the implicit function theorem, i.e., solving linearized systems or CG iterations for the gradient, facilitated by second-order AD primitives (jax.jvp, jax.vjp) (Rojas et al., 2021, Xue, 19 May 2025).

Memory and performance considerations are central, especially for long time horizons or large systems. Strategies to mitigate tape growth include checkpointing, reversible integration, custom adjoint implementations, and specialized analytic backward passes (e.g., PROFESS-AD's approach for density functional theory minimization) (Tan et al., 2022).

3. Domain-Specific Architectures and Layered Simulators

Several research groups have constructed high-performance, domain-specialized differentiable simulators and workflows, illustrating the architectural idioms and best practices in the field.

Optics and astronomical imaging: dLux composes layered transformations (pupil, phase, Fourier, detector gain layers) as JAX functions, wrapped as @jax.jit-compiled kernels. High-dimensional calibration, including millions of pixel gains and phase modes, is feasible using Optax optimizers (Desdoigts et al., 2024). This construction generalizes to any imaging pipeline expressible as a composition of differentiable layers.

Agent-based and event-driven models: To realize differentiable agent-based simulations, discrete control elements (if/else branches, state transitions, hard min/max) are replaced with parameterized smooth surrogates such as logistic sigmoids, softmin/softmax operations, and differentiable timers (Andelfinger, 2021). Recent advances involve smooth interpretation (convolutional smoothing via the SI operator) and AD-powered Monte Carlo gradient estimators (DGO), making complex control flow amenable to gradient-based methods (Kreikemeyer et al., 2023).

High-performance and distributed settings: In parallel particle-in-cell and accelerator codes, ADVars implemented via operator-overloading propagate sensitivities through all computation and communication. These modules are compatible with MPI and emerge in Fortran, C++, Python, Java, and Julia, allowing seamless deployment in legacy and modern codes alike (Qianga et al., 26 Nov 2025). Explicit formulas, operator overloading, and multi-language APIs are standard.

Physical chemistry and materials science: Automatic differentiation frameworks enable concise, error-free computation of gradients and higher derivatives in phase-equilibrium, OFDFT, and atomistic simulations. In phase equilibrium, AD allows for uniform, stable gradients in Newton–Raphson solvers, halving required iterations and eliminating convergence anomalies due to noisy Jacobians (Yang, 2023). In atomistic force field optimization (Gangan et al., 2024), complete MD/energy-minimization pipelines are built as differentiable programs, with loss functions that can target energies, forces, elastic moduli, phonons, or radial distribution functions.

Photonic shape optimization: The AutoDiffGeo paradigm defines shape primitives and their unions/intersections as compositions of smooth functions, making the mapping from shape parameters to grid-based permittivity masks fully differentiable. Adjoint field simulations composed with these mappings enable end-to-end photonic device optimization at orders-of-magnitude lower wall time than finite-difference approaches (Hooten et al., 2023).

Differentiable climate and weather modeling: Coupling learnable coordinate transforms (e.g., neural-network-based vertical coordinates) with fully differentiable dynamical cores, including exact computation of geometric metric terms via AD, optimizes grid layouts to reduce spurious modes and improve physical fidelity in predictive atmospheric models (Whittaker et al., 19 Dec 2025).

4. Optimization, Inference, and Calibration Workflows

Differentiable simulation shifts model calibration and design from hand-derived adjoint models and brute-force finite-differences to large-scale gradient-based optimization (Xu et al., 2019, Hooten et al., 2023). Typical objectives include:

• Maximum-likelihood or negative log-posterior for inference (e.g., phase retrieval, detector calibration, PDE parameter identification, quantum transport inversion).
• Policy optimization and reinforcement learning (neural agent controllers for locomotion or traffic signal control), enabling direct backpropagation through the environment (Rojas et al., 2021, Andelfinger, 2021).
• Multiobjective design, where loss functions amalgamate property errors for elasticity, vibration, structure, or photonic transmission, with per-property weights (Gangan et al., 2024).
• System identification and adaptive control (robot dynamic parameter estimation, real-time re-synthesis of model parameters informed by data streams) (Millard et al., 2020).

Algorithmic solvers include first-order methods (Adam, L-BFGS-B) and, in more nonlinear settings, second-order Newton-CG methods with exact Hessian-vector products for quadratic convergence (Xue, 19 May 2025). In agent-based or event-driven optimization, additional smoothing (or stochastic) techniques are required to ensure meaningful gradients across conditional control flows (Kreikemeyer et al., 2023).

Computational statistics for these approaches consistently demonstrate dramatic speedup and efficacy compared to gradient-free or finite-difference methods. In photonic optimization, AD-based pipelines accelerate gradient evaluation by $>50\times$, reducing wall-clock optimization time $4\times$ or more even on identically specified hardware (Hooten et al., 2023). In agent-based traffic signal optimization (with 2,500 variables), each gradient-based optimization batch gains $\sim$100 km average improvement, compared to $<$20 km for gradient-free methods in equivalent runtime (Andelfinger, 2021). In nonlinear diffusion and elasticity inversions, access to implicit Hessians via AD reduces optimization iterations from tens to a handful for Newton-CG, with commensurate wall time reduction (Xue, 19 May 2025).

5. Challenges, Limitations, and Extensions

While differentiable simulation is broadly enabling, several limitations and tradeoffs remain:

• Memory Overhead: Naïve reverse-mode AD requires storing all intermediates, which becomes prohibitive for long time horizons or large spatial grids. Strategies such as checkpointing, semi-analytic adjoints, or custom backward passes can alleviate but not fully eliminate memory growth (Li et al., 2022, Xue, 19 May 2025, Tan et al., 2022).
• Discontinuities and Non-Smooth Logic: Classical branching, event-driven updates, and threshold logic impede gradients. Smoothing replacements (sigmoid, softmin), smooth interpretation (SI), or Monte Carlo gradient estimators (DGO) provide partial solutions at an accuracy–cost tradeoff (Kreikemeyer et al., 2023, Andelfinger, 2021).
• Implicit and Second-Order Differentiation: Implicit time-stepping and equilibrium constraints require robust implicit differentiation. Matrix-free and analytic strategies for Hessian-vector products are established for finite-element PDEs but remain an open area for more complicated solver structures (Rojas et al., 2021, Xue, 19 May 2025).
• Stiffness and Scale-Sensitivity: Stiff ODEs, PDEs, or highly nonlinear systems may require specialized integrators and adjoint schemes to ensure stable gradients (Millard et al., 2020, Xu et al., 2019).
• Computational Overhead: Multiple language support, operator overloading, or source-to-source transformations introduce some performance penalty compared to hand-coded adjoints but generally remain within a small factor of optimized codes for realistic parameter counts (Qianga et al., 26 Nov 2025). High-order or symbolic approaches (SDA) can greatly reduce evaluation time for higher-order derivatives in select applications (Zhang, 1 Jun 2025).

Extensions and future directions highlighted include:

• Distributed and multi-GPU AD frameworks, to enable scaling to petascale simulations (Xu et al., 2019).
• Symbolic Differential Algebra (SDA): explicit generation and simplification of high-order derivatives, yielding orders-of-magnitude speedup for higher-order sensitivity analysis, code generation, and verified integration (Zhang, 1 Jun 2025).

6. Applications and Impact Across Disciplines

Differentiable simulation via automatic differentiation now underpins research and applications across a spectrum of fields:

Domain	Key Usage/Problem	Notable Approach/Paper
Physical optics	PSF calibration, phase retrieval, hardware design	dLux JAX-based pipeline (Desdoigts et al., 2024)
Atomistic physics	Force field fitting to elasticity, phonons, RDFs	End-to-end AD MD (Gangan et al., 2024)
Fluids & atmosphere	Learnable vertical coordinates, closure models	Solver-in-the-loop + AD (Whittaker et al., 19 Dec 2025)
Photonics	Shape optimization with differentiable geometry	AutoDiffGeo (Hooten et al., 2023)
Quantum transport	Differentiable TB/Schrödinger solver for device I–V	JAX-based auto-diff pipeline (Williams et al., 2023)
Cosmology	Field-level adjoint in PM N-body simulation	Reverse-time adjoint (Li et al., 2022)
Multi-scale flows	End-to-end differentiable hydrodynamics/kinetics	Scientific ML + diffprog (Xiao, 23 Jan 2025)
Rigid/soft dynamics	Optimal control, system identification	Reverse mode/CSA/implicit diff. (Millard et al., 2020, Rojas et al., 2021)
Agent-based systems	Gradient-based traffic, epidemic, and NN policies	Smooth primitives + AD (Andelfinger, 2021, Kreikemeyer et al., 2023)
Finite-element PDEs	Inverse design with 2nd-order implicit differentiation	Mat-free Hessian/adjoint AD (Xue, 19 May 2025)
DFT materials	Energies, stresses, phonons, elasticity	PROFESS-AD: PyTorch-based AD (Tan et al., 2022)
Beam physics	Sensitivity analysis, optimization in PIC codes	Multi-language ADVar (Qianga et al., 26 Nov 2025)

These advances have shifted the landscape for calibration, design, control, and scientific discovery, making high-dimensional, data-rich, and accuracy-sensitive optimization tasks tractable and reproducible.

7. Future Prospects and Trends

The trajectory of differentiable simulation and automatic differentiation in computational science is toward broader generality, greater scalability, and deeper integration with machine learning. Standout themes include:

• Differentiable physics engines adaptable to arbitrary modular components (nonlinear, hybrid, data-driven).
• Solver-in-the-loop optimization, backpropagating through full time-integrators and grid transforms for algorithmic design optimization (Whittaker et al., 19 Dec 2025).
• Distributed and acceleratored frameworks, mainstreamed by JAX, PyTorch, emerging C++ DSLs, and multi-language operator-overloading APIs (Desdoigts et al., 2024, Qianga et al., 26 Nov 2025).
• Higher-order sensitivity and uncertainty quantification via explicit, symbolic, and matrix-free adjoints (Xue, 19 May 2025, Zhang, 1 Jun 2025).
• Hybrid analytic-AD adjoints for differential-algebraic, stiff, and complex multi-physics systems.
• Large-scale optimization and Bayesian inference, leveraging exact gradients for parameter inference and design in models with $10^6$–$10^8$ degrees of freedom (Li et al., 2022, Desdoigts et al., 2024).

This anticipate future work on improving memory efficiency, handling discontinuities, composability with learning architectures, and end-to-end differentiation in distributed, multi-physics, and agent-based environments.