Automatic Differentiation

Follow this topic Ask about this topic

Momentum

6 papers in the last four weeks, up 100% on the four weeks before. 0.1% of all new papers.

Jul 13Week of Sep 28

Latest papers 26

Sep 30, 2026cs.LG

Same Loss, Different Gradients

Differentiable learning typically assumes that the scalar objective evaluated in the forward pass and the gradient supplied to the optimizer in the backward pass describe the same mathematical object. We show that this correspondence can fail when probabilistic objectives rely on finite special-function recurrences, custom backward rules, and numerical clipping. In high-dimensional von Mises-Fisher learning, real numerical implementations can produce identical forward scores and losses at the same learning state while supplying different gradients and following different optimization trajectories. We characterize the structure of this mismatch in finite-start Bessel recurrence and show that classwise radial mismatch can compose through probabilities into a locally nonconservative update field. Evaluating the accuracy of special-function values and derivatives separately is therefore insufficient to characterize the realized learning objective. Motivated by this observation, we introduce AR/FR, a fixed-depth analytic realization that constructs a potential and its derivative jointly, ensuring forward-backward coherence by construction. We establish a uniform cubic-order error bound relative to the exact Bessel ratio over the entire nonnegative concentration axis and propagate this guarantee to learning scores and objectives. As representation dimension increases, the original finite recurrence becomes sequentially deeper, whereas the worst-case AR/FR error guarantee tightens cubically, jointly providing coherence, certified fidelity, and fixed-depth computation. These results suggest that a differentiable numerical primitive is defined by both the values it realizes and the derivatives it actually supplies to the optimizer; together, they constitute the numerical realization of the learning algorithm.

Ningkang Peng, Xiaoqian Peng, Yifan He +4

Nanjing Normal University · Nanjing University of Chinese Medicine

Automatic Differentiation Differentiable Optimization

Sep 27, 2026cs.LG

The limits of exactness: On the failure of automatic differentiation in physics-informed machine learning

Automatic differentiation (AD) lets neural networks compute derivatives of governing equations to machine precision, and this precision has made it the computational backbone of physics-informed machine learning. Yet exactness in the mathematical sense is not the same as fidelity to the physics. Here I argue that a derivative can be numerically perfect and still be the wrong derivative for the problem at hand, because AD, by construction, has no notion of the physical structure a solution must obey. Convection and its associated directionality, diffusion, and dispersion are only the most visible instances of a much longer list that spans all branches of computational science and engineering, including conservation, thermodynamic consistency, symmetry, symplectic structure, positivity, monotonicity, and boundedness. Recognizing this broader gap reframes how the field should build the next generation of PDE-driven neural surrogates.

Ameya D. Jagtap

Aerospace Engineering Department, Worcester Polytechnic Institute, Worcester, MA 01609, USA

Automatic Differentiation PDE Surrogate Modeling Scientific ML Physics-Informed ML

Sep 21, 2026cs.LG

FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention

Transformer models built on the attention mechanism have become a central building block in modern deep learning, yet softmax attention remains a major bottleneck for long-context workloads. While FlashAttention makes the forward and first backward passes I/O-efficient, it does not support backward-over-backward (BoB), which enables exact differentiation through the backward pass for applications such as second-order optimization, test-time training, gradient-based memory, and meta-learning. Existing BoB implementations either materialize large intermediate tensors or exhaust GPU memory at long sequence lengths. We present FlashBoB, an exact, I/O-efficient algorithm for BoB in softmax attention that keeps computation within on-chip tiles and avoids all

N \times N

intermediate tensors, where

N

is the sequence length. The key insight is a hierarchical affine structure in the softmax double backward: two row-wise scalars determine all outputs through affine transformations. This yields a two-pass schedule with bounded on-chip static random-access memory (SRAM) usage and minimal off-chip high-bandwidth memory (HBM) traffic. FlashBoB achieves

Θ(N^2 d^2/M)

HBM traffic (

d

is the head dimension and

M

is the memory size) and, within the standard FlashAttention-style score-recomputation model, matches the inherited large-cache lower bound for exact forward attention. Empirically, it scales exact attention BoB to

N=262\text{K}

on a single A100 80GB GPU, where prior PyTorch exact baselines fail by

N=16\text{K}

, and is up to

6.3\times

faster than FlashBack. These results make exact second-order attention practical at long-context sequence lengths where prior implementations cannot run efficiently.

Anthony Givans, Michael Crawshaw, Mingrui Liu

Softmax Attention Automatic Differentiation Efficient Attention

Sep 17, 2026physics.chem-ph

Truncated automatic sparse differentiation for machine learning interatomic potentials

Machine learning interatomic potentials (MLIPs) learn the mapping from atomic positions to potential energy. The forces, the negative gradient of this energy, drive molecular dynamics and are readily obtained using automatic differentiation. Higher-order derivatives, most notably the Hessian, describe collective motion and allow the direct prediction of experimental observables, but are considered computationally inaccessible for large systems. We suggest a solution: in physical systems, interactions decay with distance, and most MLIPs build on this locality through message passing up to a finite receptive field. This implies both sparsity of higher-order derivatives and their decay with distance. This structure can be exploited using automatic sparse differentiation (ASD). We explain how to compute the sparsity pattern for MLIP derivatives and demonstrate that, for multiple foundation MLIPs, ASD computes full Hessians of large porous materials exactly, but with modest speedups at best. The larger gains come from truncated ASD: discarding small, but nonzero, Hessian entries between distant atoms yields order-of-magnitude speedups with negligible impact on predicted observables.

Marcel F. Langer, Adrian Hill, Michele Ceriotti

Laboratory of Computational Science and Modeling (COSMO), EPFL, Lausanne, Switzerland · Machine Learning Group, Technical University of Berlin, Berlin, Germany · BIFOLD – Berlin Institute for the Foundations of Learning and Data, Berlin, Germany

Automatic Differentiation Machine Learning Interatomic Potentials Computational Materials Science

Sep 16, 2026stat.ML

Rank and computation of the pathlifting Jacobian of a DAG ReLU network

This paper provides a self-contained proof of the rank of the pathlifting Jacobian of a DAG ReLU network by performing an induction on the network's number of hidden nodes. In fact, the induction is elementary, and the key recipe is to consider the skeleton matrix of the network, a sparse matrix encoding the network paths, and transform the representation of one of its hidden neurons into an output node. The proof relies on intermediate propositions which link the pathlifting, its Jacobian, the network parameters, and its skeleton matrix, which, on top of permitting to conclude on the rank of the pathlifting Jacobian, also provide a way to compute it without backpropagation and whose computation cost is super efficient in practice compare to usual backpropagation. The paper is provided with a Python module that implements the different propositions of the paper for feed forward networks and is used to experimentally quantifies the computational gain of computing the pathlifting Jacobian with the proposed theory.

Manon Verbockhaven

OCKHAM

Automatic Differentiation ReLU Neural Networks

Sep 7, 2026math.NA

A Systematic Analysis of Automatic Differentiation versus Discretization-based Constraints for Physics-Informed PDE Solvers

Physics-informed neural networks (PINNs) represent a growing frontier in using artificial intelligence to solve partial differential equations (PDEs). Automatic differentiation (AD) plays a central role in this paradigm, which is mesh-free and replaces traditional iterative solvers with gradient-based optimization in continuous space. However, the inherent limitations of AD, particularly in handling higher-order derivatives and discontinuous solutions, pose significant challenges for complex problems. This has motivated a growing number of researchers to explore discretization-based constraints as an alternative path. Yet, the respective applicability of these two paradigms remains largely unexplored. In this work, we conduct systematic experiments across a wide spectrum of problems, from simple linear Poisson to high-Mach hypersonic flows with strong discontinuities. Through a rigorous decomposition of approximation, optimization, and truncation errors, we systematically elucidate the fundamental trade-offs and error-governing mechanisms of both paradigms, as well as two representative network architectures: multi-layer perceptron (MLP) and graph neural network (GNN). Our results reveal a consistent trend: as nonlinearity strengthens, the accuracy advantage of discretization-based constraints becomes increasingly pronounced, with smaller optimization errors compensating for the truncation errors. Moreover, the more complex the nonlinearity and boundary conditions, the greater the advantage of GNN over MLP. These insights offer a robust practical guideline for configuring neural PDE solvers in demanding engineering applications. Our source data and code are available at https://github.com/guoxing0809/neuropde_analysis.

Xing Guo, Hongwei Tang, Zewei Meng +3

School of Systems Science and Engineering, Sun Yat-sen University, Guangzhou, China

Automatic Differentiation Neural PDE Solvers PDE Solving Physics-Informed ML

Sep 2, 2026cs.LG

A Computational Comparison of Fourier Spectral Differentiation and Spatial Automatic Differentiation in Periodic Physics-Informed Neural Networks

Physics-informed neural networks (PINNs) commonly evaluate the spatial derivatives appearing in partial differential equation residuals using automatic differentiation (AD), whose computational and memory costs can become substantial when multiple or high-order derivatives are required. We perform a controlled comparison of spatial AD and Fourier spectral differentiation in periodic physical-space PINNs. Within each paired experiment, the neural representation, temporal differentiation, optimizer, sampling procedure, and training schedule are held fixed, so that the two cases differ only in the spatial differentiation procedure. For the Fourier variant, network outputs are evaluated on a uniform periodic grid and transformed to Fourier space, where spatial derivatives are obtained through spectral multiplication and the same Fourier coefficients are reused across derivative orders. We compare the two procedures in standard PINNs for the Allen--Cahn and Korteweg--de Vries equations and in Causal PINNs for the Allen--Cahn, Korteweg--de Vries, and Kuramoto--Sivashinsky equations. Across these five equation--framework settings, Fourier differentiation yields mean paired end-to-end training speedups ranging from

2.90\times

to

18.52\times

and reduces peak allocated graphics processing unit (GPU) memory by

68.7\%

--

94.1\%

. The final relative

L_2

errors remain of the same order, with neither differentiation procedure showing a consistent accuracy advantage. For the one-dimensional periodic benchmarks considered here, Fourier spectral differentiation therefore provides substantially lower training time and memory usage than spatial AD while retaining comparable solution error, at the cost of requiring a uniform structured spatial grid.

Xilai Liang, Zhao Zhang

Guangdong Technion–Israel Institute of Technology, Shantou, Guangdong, China · Research Centre for Mathematics and Interdisciplinary Sciences, Shandong University, Qingdao, Shandong Province, 266237, China

Physics-Informed ML Automatic Differentiation Spectral Methods

Aug 31, 2026cs.LG

Hard-ReLU Gradient Descent Selects an Event-Free Sensitivity Limit

Gradient flow is widely used as a continuous-time surrogate for gradient descent, but state convergence does not imply convergence of differentiated training maps in nonsmooth networks. We characterize the fixed-horizon, vanishing-step limit of exact automatic differentiation through hard-ReLU gradient descent. Under a stable finite itinerary of separated, same-direction transverse activation events, gradient-descent states converge at first order to the corresponding piecewise-smooth gradient flow, while the exact derivative of every nonresonant discrete program converges to an event-free regional propagator. The true flow derivative instead interleaves classical saltation matrices that encode event-time sensitivity. For globally convex objectives, any strict activation event prevents complete cancellation of these missing transfers. Moreover, minimal globally 1-strongly convex residual-ReLU risks can realize arbitrarily large reciprocal sensitivity gaps, subject to an explicit transversality-scale tradeoff, and a coupled strongly convex construction yields an open set on which the largest initialization-gradient coordinate is reversed. In a controlled 17-parameter ReLU MLP, state and regional-AD errors vanish under mesh refinement while AD-to-flow errors remain between 0.18 and 0.39; an event-aware corrected product restores convergence. Resolved smoothing likewise recovers the flow sensitivity when the transition layer is sufficiently resolved. These results show that the gradient-flow limit of hard-ReLU training need not remain valid after differentiation.

Xiaoyang Li, Runni Zhou

College of Medicine and Biological Information Engineering Northeastern University, Shenyang, China

Dynamical Systems Automatic Differentiation Differentiable Programming Gradient Descent Dynamics

Aug 11, 2026cs.LG

Derivative Computation in PINNs: Automatic Differentiation, Finite Differences and Beyond

We systematically investigate finite-difference (FD) derivative computation in Physics-Informed Neural Networks (PINNs) as an alternative to automatic differentiation (AD). On three benchmark PDEs we show that, with a properly calibrated step size, FD matches AD in accuracy on every problem while running faster across the full tested batch-size range and using substantially less GPU memory, and that a stochastic variant we propose outperforms AD on a stationary problem. We further show that for neural architectures with inter-sample dependencies (e.g. BatchNorm, self-attention) the standard PyTorch autograd idiom is silently incorrect; the correct per-sample alternative is computationally infeasible at PINN-relevant batch sizes, while FD provides a forward-only approximation that is empirically an order of magnitude closer to the true per-sample derivative.

Maciej J. Mikulski, Tadeusz Uhl

AGH University of Krakow

Automatic Differentiation Neural PDE Solvers PDE Solving Physics-Informed ML

Aug 9, 2026stat.ML

LazyHMC: Hamiltonian Monte Carlo Simulation for Lazy, Infinite Dimensional Probabilistic Programs

Hamiltonian Monte Carlo (HMC) is a successful generic inference method in probabilistic programming, but in its ordinary formulation it needs gradients and finite-dimensional parameter spaces. In Haskell, lazy evaluation lets probabilistic programs express stochastic processes and other non-parametric Bayesian models over implicit infinite-dimensional spaces. This paper develops new formulations of gradient-based HMC for this infinite-dimensional setting, via lazy evaluation. For automatic differentiation, we provide an analysis based on a new notion of "piecewise analytic under cylindrical analytic partition" (PACAP), to show that even if a program is infinite-dimensional and defined lazily, the gradient of the likelihood function is finitely supported. For the Monte Carlo method itself, we develop several HMC variants and a No-U-Turn Sampler that operate over the infinite-dimensional parameter space but are still productive because of lazy evaluation. Experiments cover Gaussian mixture clustering, random walks, and piecewise-constant regression with Poisson-process changepoints.

Maria-Nicoleta Crăciun, C. -H. Luke Ong, Tom Schrijvers +1

University of Oxford, UK · Nanyang Technological University, Singapore, Singapore · KU Leuven, Leuven, Belgium

Markov Chain Monte Carlo Automatic Differentiation Bayesian Inference Probabilistic Programming

Aug 9, 2026cs.GR

Differentiate the Solver, Not the Equation: Reverse-Sweep Adjoints for Block Implicit Simulation

Differentiable simulation is a key component in learning, control, and inverse problems, where gradients through nonlinear implicit solvers are required. Existing approaches either rely on unrolled automatic differentiation, whose memory grows with solver depth, or on equation-level implicit differentiation, which assembles global Jacobians and solves large sparse adjoint systems, discarding the locality of the forward solver -- and differentiating the converged equation rather than the finite computation that actually ran. We propose solver-level differentiation, which differentiates the executed solver itself. When a solver is composed of block implicit updates, its discrete adjoint is obtained by applying the corresponding adjoint updates in reverse order, yielding a reverse-sweep formulation whose backward pass mirrors the forward solver. From an operator perspective, the forward pass realizes an approximate inverse through ordered local solves, and the backward applies its transpose through reverse local adjoint solves, constructing no global system. We instantiate this idea on Vertex Block Descent, yielding a differentiable solver whose reverse colored Gauss-Seidel sweeps are composed entirely of local

3\times 3

adjoint solves. The backward matches automatic differentiation through the identical executed forward to machine precision at every solver depth, where the equation-level adjoint is off by 37% after one sweep; in a controlled same-codebase, same-GPU comparison it is 33x faster and uses 71x less memory than unrolled automatic differentiation; and the same construction is exact on projective dynamics and extended position-based dynamics. We scale differentiable elastodynamics to

10^6

contact-coupled soft bodies (8M vertices) on one GPU. Overall, this work highlights solver structure as a practical organizing principle for efficient differentiable simulation.

Lei Shu, Ying Jiang, Kui Wu +3

Stanford University · University of California, Los Angeles · LightSpeed Studios, Tencent America +1

Automatic Differentiation Differentiable Programming Differentiable Physics Physics-Based Simulation

Jul 31, 2026cs.PL

A Fortran General-Purpose Transpiler: Proof of Concept

Fortran has been the cornerstone of high-performance computing for decades and remains unmatched in many domains. Yet the language faces an expertise gap: a new generation of scientists is barely familiar with it, while many experienced Fortran developers are only now transitioning to modern ecosystems such as JAX. This gap often results in "Fython" - Python code written with a Fortran mindset - that fails to leverage modern frameworks. We present FGPT, a transpiler framework designed to bridge this gap. It provides a systematic pipeline that transpiles Fortran into GPU-adapted Fortran, auto-differentiable Fortran via Tapenade, or NumPy and JAX scripts. Its architecture comprises three stages: (i) a frontend that parses Fortran and extracts target procedures along with all their dependencies; (ii) a middle-end that lowers the code into an intermediate representation, then into GPU-adapted or auto-differentiable Fortran, or a NumPy class; and (iii) a backend that transforms NumPy scripts into JAX modules optimized for GPU acceleration and automatic differentiation. Large language models fail when applied to the scale of community scientific codes-often spanning hundreds of thousands of lines-where consistent transformations, strict numerical fidelity, and validation against production tests are non-negotiable. FGPT addresses these challenges by preserving program semantics throughout the entire translation. We verified the framework on representative climate modeling kernels and demonstrated that it produces correct, differentiable Python implementations without requiring manual intervention. By combining rigorous compiler techniques with modern accelerator support, FGPT offers a scalable, trustworthy path for modernizing legacy Fortran code.

Shivamshan Sivanesan, Kazem Ardaneh

1Climate Modeling Center, Sorbonne University, CNRS, IPSL, Paris, France

Code Translation Automatic Differentiation High-Performance Computing

Jul 27, 2026cs.RO

Amortising Trajectory Optimisation for Residual MPC via Implicit Contact Differentiation

Differentiable simulation can accelerate contact-rich trajectory optimisation by exposing local sensitivities of task outcomes to controls. Existing approaches either use finite differences, which are expensive and step-size sensitive; differentiate iterative contact solvers by unrolling automatic differentiation (AD), which stores a growing computation trace; or require intricate, solver-specific KKT sensitivity derivations. We introduce an AD-assisted implicit derivative for regularised smooth contacts and apply it to Mujoco MJX, based on the Implicit Function Theorem (IFT). The method differentiates the stationarity residual at the tolerance-converged solution, avoiding both solver unrolling and hand-assembled KKT systems. IFT keeps compiled temporary memory nearly constant with solver effort, changing by less than 4

\%

from one to ten iterations versus 10.6

\times

growth for unrolled AD. IFT memory grows slower with active contacts and model dimension, using 20

\times

less memory at 256 contacts and 6

\times

less at 16 contacts and 96 DoF. We further introduce optimiser distillation for residual MPC, amortising batched full-horizon iLQR into a policy that guides short-horizon residual iLQR. Across Finger, Franka, and Unitree, this raises six-step success by 28-98 percentage points over standard iLQR.

Daniel Layeghi, Thomas Corbères, Calum Arnott +4

School of Informatics, University of Edinburgh, Scotland.

Trajectory Optimization Automatic Differentiation Nonlinear MPC Model Predictive Control Differentiable Optimization

Jul 20, 2026cs.LG

FlashPDE: A Drop-In Fused Triton Operator Library for Neural PDE Solvers

Physics-Informed Neural Networks (PINNs) solve PDEs by incorporating physical constraints into neural-network training, but large-scale problems are limited by automatic-differentiation memory overhead and inefficient execution of grid-based PDE operators. We present FlashPDE, a drop-in fused operator library for grid-based scientific machine learning. FlashPDE replaces fragmented PyTorch finite-difference execution with differentiable Triton kernels. Each operator integrates fused stencil evaluation, an analytic discrete-adjoint backward pass, and boundary-gradient correction within a unified PyTorch autograd Function interface. The library provides 14 differentiable PDE operators covering 17 configurations across 1D--3D elliptic, parabolic, and Navier--Stokes systems, while remaining independent of neural architectures and training strategies. Experiments on an NVIDIA A100 GPU show that FlashPDE reduces peak memory usage by up to 37.0x compared with coordinate-based automatic differentiation and reduces CUDA kernel launches by up to 3.5x compared with eager PyTorch finite-difference implementations. Across six representative PDE benchmarks, FlashPDE achieves up to 2.30x end-to-end time-to-solution speedup and up to 19.2x kernel-level acceleration while maintaining numerical agreement with PyTorch finite-difference references. FlashPDE provides a hardware-efficient execution layer that bridges differentiable PDE solvers and GPU-optimized numerical computation within the PyTorch ecosystem.

Peiyu Zang, Bosen Xie, Ruoxiang Xu +1

School of Mathematical Sciences, Beijing Normal University · 2Beijing University of Posts and Telecommunications

Automatic Differentiation Neural PDE Solvers GPU Kernel Optimization PDE Solving Differentiable Programming Physics-Informed ML

Jul 3, 2026cs.LG

Differentiate the Evaluator, Not the Program: An Efficient Runtime Representation for Neuro-Symbolic Learning

AI systems increasingly propose executable scientific models whose value depends on both their symbolic structure and their fitted continuous parameters. This makes parameter calibration the bottleneck of program-and-parameter co-search: an outer loop can generate thousands of candidate programs, but each needs an inner gradient-based optimization before it can be assessed. Staging each candidate into its own differentiable graph makes individual models fast but sacrifices the program-as-data property that keeps search fluid; interpreter-based approaches preserve programs as runtime data but pay interpreter overhead that dominates the numerical work. We present the Native Differentiable Virtual Machine (NDVM), a runtime representation that differentiates executable programs without compiling each candidate into a separate graph. NDVM separates symbolic structure from differentiable numeric state: tags, symbols, environments, and control remain native runtime data, while numeric payloads live in dense batched buffers with exact reverse-mode gradients recorded along the realized execution trace, so one evaluator walk is amortized across large populations of parameter vectors. A locked cost model of a real differentiable self-hosted Scheme interpreter motivates the design. We realize NDVM as a native runtime with forward and gradient equivalence to the reference backend, about 60x per-lane batch amortization, near-linear multicore scaling, and two independent front ends. In fixed-budget co-search over LLM-proposed programs, NDVM reaches high-quality solutions about 24x sooner in wall-clock time, suggesting runtime differentiation as a practical systems foundation for scientific discovery workflows.

Lucas Sheneman

Institute for Interdisciplinary Data Sciences University of Idaho

Automatic Differentiation Differentiable Programming Differentiable Optimization Neurosymbolic Learning

Jun 26, 2026physics.comp-ph

Mosaic: A Benchmark Suite for Differentiable Physics Solvers

Differentiable partial differential equation (PDE) solvers underpin solver-in-the-loop ML training, gradient-based optimal control, and inverse problems, yet the practical cost of obtaining correct, usable gradients from a given solver on a given problem is largely undocumented. Integration effort, computational cost, gradient accuracy, and numerical conditioning vary widely across solvers and are discoverable only by trial and error. We introduce Mosaic, an extensible benchmarking framework for differentiable PDE solvers that standardizes access to solver gradients. Each solver is packaged as a containerized component (Tesseract) exposing a uniform gradient API regardless of language or automatic differentiation (AD) strategy, enabling researchers to evaluate, compare, and build on non-trivial physical solvers. Our evaluation of 14 solvers across fluid dynamics, structural mechanics, and heat transfer demonstrates that the benchmark surfaces practically relevant differences: order-of-magnitude variation in computational cost and Jacobian conditioning, alongside structural incompatibilities that eliminate solvers from realistic tasks entirely. Despite this variation, all solvers that produce gradients converge to similar optima, indicating that the practical barriers are memory limits, numerical stability, and setup compatibility rather than gradient accuracy alone. Mosaic is open-source and available at https://github.com/pasteurlabs/mosaic.

Andrin Rehmann, Heiko Zimmermann, Dion Häfner

Pasteur Labs, Brooklyn, NY, USA

Automatic Differentiation PDE Solving Benchmark Design Differentiable Programming Differentiable Physics

Jun 22, 2026cs.LG

FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training

Reverse-mode differentiation computes every weight gradient, writes it to memory, and only then lets the optimizer read it back. This two-phase schedule sets the memory ceiling of modern training: at the seam between the phases, every layer's gradient is live at once. We argue that this materialized gradient is an artifact of how differentiation is staged, not a quantity that learning requires -- and we eliminate it. FORGE folds the optimizer step into the backward pass and applies it one tile at a time, entirely in registers, so each gradient tile is consumed the instant it is produced and never becomes a tensor. The fusion changes only when the update happens, not what it computes: in full precision the fused step is provably exact -- the identical optimizer update, for every element-wise rule -- and that exactness survives tensor- and sequence-parallel sharding; in the bf16 and 8-bit regimes used in practice it is faithful rather than bit-identical, its deviation bounded and, for the weight store, rendered unbiased by stochastic rounding. Because each gradient tile is born and consumed in the same registers, it is never converted down to bf16 to be stored and read back; FORGE thus preserves the full-precision fidelity that both bf16 and 8-bit optimizers lose to that conversion. Nor is the method tied to one architecture or one optimizer: linear layers are ubiquitous, and FORGE reclaims the gradient memory of any of them under any element-wise rule. Empirically FORGE more than halves the memory of an optimizer step and, at the small batch sizes typical of fine-tuning and continued pretraining, runs about 1.5x faster; integrated into tensor-parallel Megatron-LM it fits 8B training at four times the micro-batch a standard optimizer allows on the same GPUs.

Dikshant Kukreja, Kritarth Prasad, Avinash Anand +6

Puch AI · Singapore Institute of Technology · Nanyang Technological University +2

Automatic Differentiation Backpropagation Memory-Efficient Optimization Efficient Language Model Training LLM Training Memory-Efficient Training

Jun 13, 2026cs.LG

Automatic Differentiation from Scratch: How PyTorch Computes Gradients in Physics-Informed Neural Networks

This paper traces, with explicit numerical values, how PyTorch's automatic differentiation (AD) engine computes gradients for Physics-Informed Neural Network (PINN) training -- a setting that requires two levels of differentiation: computing the physics derivative

\hat{y}'(t)=d\hat{y}/dt

through the network, and computing parameter gradients

\nabla_θL

of a loss that itself depends on

\hat{y}'(t)

. Using a 1-3-3-1 multilayer perceptron and the initial value problem

y'(t)+y(t)=0

,

y(0)=1

, we trace the complete pipeline at every node: the computational graph built during the forward pass, the reverse-mode backward traversal that computes all 22 parameter gradients in a single pass, and the graph-on-graph mechanism by which \texttt{create_graph=True} enables correct differentiation through the physics-informed residual. Every adjoint value is verified against the hand derivations of Tahimi (2026), connecting the

P/Q

sensitivity framework to the vector--Jacobian products used by PyTorch's autograd engine.

Abdeladhim Tahimi

CECA, Universidade Federal de Alagoas (UFAL), Rio Largo, Alagoas, 57100-000, Brazil

Automatic Differentiation Backpropagation

Jun 7, 2026cs.PL

Compile Once, Differentiate Everywhere: A Differentiable Meta-Circular Interpreter

The boundary between program execution and gradient-based optimization has long limited the use of code itself as a learnable scientific model. We present a compiler that translates a self-hosting subset of Scheme into differentiable computation graphs for autograd backends. Because the subset can compile its own evaluator, this yields differentiable meta-circular interpretation (DMCI): a compiled Scheme interpreter executes programs supplied as data, while reverse-mode autodiff propagates gradients to continuous constants embedded in those programs. The interpreter is compiled once, so new programs inherit differentiability without recompilation or custom gradient machinery, while retaining closures, recursion, and data structures. We prove that gradients through the compiled interpreter are correct almost everywhere and show that they match direct compilation to numerical precision across 171 recursive and higher-order program-seed pairs. We then use DMCI for program-and-parameter co-search, where a large language model proposes Scheme programs and exact gradients calibrate their continuous parameters through a single frozen interpreter. This enables OpenEvolve-style program search in which an outer loop proposes discrete program structures and DMCI supplies exact gradient-based calibration of each candidate's continuous parameters. On battery capacity-fade data, the search recovers a knee-like degradation structure and improves held-out extrapolation over hand-crafted baselines on the harder early-extrapolation split, matching them on the later split. On a high-dimensional El Nino inverse problem, DMCI optimizes an interpreted Kalman-filter likelihood where gradient-free search fails. These results extend symbolic regression and neurosymbolic search from closed-form expressions to executable, stateful programs, making model-generated code directly optimizable against data.

Lucas Sheneman

Institute for Interdisciplinary Data Sciences University of Idaho

Automatic Differentiation Program Synthesis Differentiable Programming Differentiable Optimization Symbolic Regression

May 22, 2026physics.comp-ph

Finite Element-Based Material Learning via Automatic Differentiation: Learning constitutive neural network models from full-field deformation data

The identification of constitutive neural network models from heterogeneous full-field deformation data provides a robust alternative to traditional calibration methods based on homogeneous stress-strain experiments, particularly given the high dimensionality of trainable parameters. Existing approaches must balance generality, robustness, and computational efficiency: Conventional finite element model updating is broadly applicable but computationally demanding; weak-form methods offer efficiency but are sensitive to noise and data scarcity; neural operator models are highly expressive but require extensive training datasets. This work presents FE-MAD (Finite Element-Based Material learning via Automatic Differentiation), an end-to-end differentiable framework that integrates a constitutive neural network model within a JAX-FEM nonlinear solver and identifies its parameters through gradient-based minimization of a measurement-mismatch loss. Newton tangent stiffness and loss gradients are computed automatically using forward- and reverse-mode automatic differentiation throughout the entire pipeline, thereby removing the need for analytic adjoints or offline surrogate models. FE-MAD is demonstrated for two architectures: a grey-box Constitutive Artificial Neural Network (CANN), a polyconvex, fully connected model with high flexibility, and a white-box CANN, an expert-system network with phenomenologically interpretable strain-energy terms. Focusing on incompressible isotropic hyperelasticity, FE-MAD is evaluated on three open experimental datasets: (1) full digital image correlation (DIC) of a perforated tensile specimen, (2) a reduced-data scenario with a one-dimensional stretch profile and global force-displacement curve, and (3) a heterogeneous matrix-inclusion system in which both phases constitutive laws are identified and generalized to twenty-two previously unseen samples.

Matthias Knipper, Chenyi Ji, Malte Brand +1

Computational Mechanics in Medicine, Applied Medical Engineering, RWTH Aachen University Pauwelsstraße 20, 52074 Aachen, Germany · Institute for Continuum and Material Mechanics, Hamburg University of Technology, Eißendorfer Straße 42, 21073 Hamburg, Germany

Automatic Differentiation Hyperelasticity Materials Science Physics-Informed ML Finite Element Method

May 22, 2026cs.LG

Archimedean Copula Inference via Taylor-Mode AD

No existing nested Archimedean copula tool handles all three of (a) arbitrary per-variable (right-)censoring in survival analysis, (b) arbitrary nesting trees, and (c) exact parameter gradients. Existing implementations handle only bivariate problems, low dimensional (i.e.,

d \leq 10

) cases, two layers of nesting, or only hand-derived copula nestings. We present \textsc{acopula}, a JAX-native framework that, given any Archimedean generator -- classical or neural -- evaluates exact nested-copula likelihoods and parameter gradients under arbitrary censoring masks in polynomial time. The mechanism is polynomial powering of Taylor-mode automatic differentiation output, which replaces per-family hand-derived partial Bell polynomial tables with a single differentiable computation that any user-defined generator can drive. We conduct extensive simulations to verify the correctness of \textsc{acopula}. We then demonstrate (a) per-variable censoring on

85{,}229

MIMIC-IV ICU admissions in high dimensions with

d{=}53

, fit by both classical Archimedean families and nested neural Archimedean copulas; (b) an 11-sector hierarchical model on S&P~500 daily returns at

d{=}98

; (c) family-agnostic censored MLE across ten families, five of them with no prior implementation, on a retinopathy study; and (d) a

{\sim}650\times

per-density speedup over R's \texttt{nacLL} at

d{=}35

, scaling quadratically to

d{=}8{,}000

.

Cambridge Yang, Dongdong Li

Independent · Harvard Medical School

Parameter Estimation Automatic Differentiation Copula Models

May 6, 2026cs.LG

Differentiable Parameter Optimization for DAEs with State-Dependent Events

Differential-algebraic equations (DAEs) with state-dependent events arise in systems whose continuous dynamics are constrained by algebraic equations and interrupted by mode changes, switching logic, impacts, or state reinitializations. Gradient-based parameter learning for such systems is challenging because algebraic variables are implicitly defined, event times depend on the parameters, and reset maps introduce discontinuities. This paper studies differentiable parameter optimization for semi-explicit DAEs with events. We formulate the learning problem as a constrained least-squares problem with DAE dynamics, algebraic constraints, guard equations, and reset maps. We then develop two complementary gradient-computation strategies. The first is an automatic-differentiation-through-simulation method that solves algebraic variables inside the vector field, differentiates the algebraic solve using the implicit function theorem, and handles events through segmented differentiable integration. The second is an explicit discrete-adjoint method that represents the forward simulation as an event-split residual system and computes gradients by solving for the Lagrange multipliers of smooth-segment and event residuals. The formulation clarifies that residual terms in the adjoint method are equality constraints, not heuristic penalties. We compare the two approaches in terms of gradient interpretation, event-time handling, implementation complexity, and local validity. Both methods provide gradients for the event path selected by the forward simulation and are valid under fixed event ordering and transversal guard crossings.

Ion Matei, Maksym Zhenirovskyy, Anthony Wong

Fujitsu Research of America

Dynamical Systems Automatic Differentiation Constrained Optimization Differentiable Physics Differentiable Optimization

May 4, 2026cs.RO

Exact Higher-Order Derivatives for SE(3) via Analytical/AD Methods

Fast prototyping of new SE(3) estimation objectives remains awkward in practice. Modern Lie-group frameworks -- GTSAM, manif, Sophus, SymForce, Ceres -- target first-order workloads through different code-generation and automatic-differentiation strategies, each optimized for a particular seam between hand-derived geometry and generic differentiation. The remaining gap is a compact, AD-safe path from these first-order primitives to exact Hessians, observed-information matrices, and higher-order derivative tensors: the quantities needed for exact Newton steps, observed-information covariance estimates, and covariance correction. This paper presents a hybrid analytical/AD recipe for SE(3) negative log-likelihoods. The practitioner writes the NLL gradient once, generic over a scalar type, and places the analytical/AD seam at the point-action interface y = Tx. Closed-form Lie-group Jacobians are used up to this interface; AD is applied only beyond it. The same source is then instantiated with ordinary floating-point scalars for gradients, vector-seeded dual numbers for exact Hessians in a single forward-mode pass, and nested dual numbers for higher-order derivative tensors. On a representative 6-DoF, 5-landmark SE(3) NLL, the advocated seeded-Hessian path is approximately 5x faster than finite-differencing the AD gradient on this benchmark while matching a nested-AD oracle to machine precision. The implementation adds roughly 70 lines of analytical-Jacobian code over an AD-only baseline. We also identify and fix a removable singularity in the standard SO(3)/SE(3) scalar basis that would otherwise produce NaNs at the origin under seeded AD, and we audit which Lie-group derivative tensors require this stabilized basis. The result is a practical path from rapidly written SE(3) objectives to exact higher-order derivatives, with predictable runtime and no finite-difference tuning.

Frank O. Kuehnel

Sigma Point Labs, Denver, CO, USA

Automatic Differentiation Lie Group Methods Second-Order Optimization Differentiable Optimization

May 3, 2026cs.LG

Floating-Point Networks with Automatic Differentiation Can Represent Almost All Floating-Point Functions and Their Gradients

Theoretical studies show that for any differentiable function on a compact domain, there exists a neural network that approximates both the function values and gradients. However, such a result cannot be used in practice since it assumes real parameters and exact internal operations. In contrast, real implementations only use a finite subset of reals and machine operations with round-off errors. In this work, we investigate whether a similar result holds for neural networks under floating-point arithmetic, when the gradient with respect to the input is computed by the automatic differentiation algorithm

D^\mathtt{AD}

. We first show that given a floating-point function

φ

(e.g., a loss function), arbitrary function values and gradients can be represented by a floating-point network

f

and

D^\mathtt{AD}(φ\circ f)

, respectively. We further extend this result: given

φ_1,\dots,φ_n

,

D^\mathtt{AD}(φ_i\circ f)

can simultaneously represent arbitrary gradients while

f

represents the target values, under mild conditions. Our results hold for practical activation functions, e.g.,

\mathrm{ReLU}

,

\mathrm{ELU}

,

\mathrm{GeLU}

,

\mathrm{Swish}

,

\mathrm{Sigmoid}

, and

\mathrm{tanh}

.

Sejun Park, Yeachan Park, Geonho Hwang

Department of Artificial Intelligence, Korea University, Seoul, 02841, Republic of Korea · Department of Mathematics and Statistics, Sejong University, Seoul, 05006, Republic of Korea · Department of Mathematical Sciences, Gwangju Institute of Science and Technology, Gwangju, 61005, Republic of Korea

Automatic Differentiation Neural Network Approximation Theory

Apr 28, 2026cs.LG

Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge

We study long-context balanced entropic optimal transport (OT) attention on TPU hardware through a stopped-base, fixed-depth tail-refinement surrogate. After a stopped

T

-step Sinkhorn solve, we unroll a short refinement tail and differentiate that surrogate exactly. For the reported

R=2

TPU path, the backward pass contains four staircase plan factors. We prove an exact one-reference-tile schedule: the

R=2

score cotangent is a single reference plan tile times an explicit modifier field built from vector cotangents and dual differences. This yields block-wise cost

O((T+R)LW)

,

O(Ld)

input storage, and

O(L)

additional HBM usage for fixed head dimension

d

and band width

W

on the balanced fixed-support path. We also formalize the current \texttt{dustbin_block} path as the same unit-target surrogate on an augmented support, so the adjoint schedule lifts to the single-active-dustbin path used in our TPU runs; this bridge is algebraic and does not claim a general KL-unbalanced or arbitrary-capacity gap model. We provide a local surrogate-bias bound, an a posteriori bias certificate, and a projective contraction certificate for strictly positive active blocks. On synthetic masked problems, the optimized kernel matches exact autodiff of the same centered surrogate to within

10^{-5}

--

10^{-10}

. On TPU v6e-8, a four-configuration Pfam screen completes end-to-end, and a promoted balanced

R=2

run sustains roughly

8.5

examples per second through a three-hour budget, reaching step

1437

. Held-out Pfam test shards improve reconstruction from

5.57

to

2.05

and sparse CE from

5.53

to

5.30

relative to step

0

, with CE logged diagnostically rather than optimized directly; target-barycenter alignment metrics do not materially improve, and a deterministic diagonal reference remains stronger on those metrics.

Dylan Forde

Independent Researcher

Automatic Differentiation Transformer Attention Differentiable Programming Algorithm Unrolling Attention Mechanisms Optimal Transport

Dec 5, 2025cs.NE

ADSEQ: A delay-aware autograd-compatible framework for spike-event delivery in SNNs

Spiking neural networks (SNNs), central to computational neuroscience and neuromorphic machine learning (ML), require efficient simulation and gradient-based training. While AI accelerators offer promising speedups, gradient-based SNNs typically implement sparse spike events using dense, memory-heavy data-structures. Existing exact gradient methods lack generality, and current simulators often omit or inefficiently handle delayed spikes. We address this by deriving gradient computation through spike event queues, including delays, and implementing this into memory-efficient, autodifferentiable spike event queues (ADSEQ). These are benchmarked across CPU, GPU, TPU, and LPU platforms. We find that queue design strongly shapes performance. CPUs, as expected, perform well with traditional tree-based or FIFO implementations, while GPUs excel with ring buffers for smaller simulations, yet under higher memory pressure prefer more sparse data-structures. TPUs seem to favor an implementation based on sorting intrinsics. Selective spike dropping provides a simple performance-accuracy trade-off, which could be enhanced by future autograd frameworks adapting diverging primal/tangent data-structures.

Lennart P. L. Landsmeer, Amirreza Movahedin, Said Hamdioui +1

Department of Computer Engineering, Delft University of Technology, Netherlands · NeuroComputingLab, Department of Neuroscience, Erasmus Medical Center, Netherlands

Automatic Differentiation Spiking Neural Networks