Research map / Probabilistic and scientific modeling / Molecules and proteins
Automatic Functional Differentiation in JAX
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Extends JAX to automatically differentiate functionals and operators by treating functions as generalized arrays and deriving JVP and transpose rules for a set of primitive operators (compose, nabla, linearize, linear transpose, integrate).
Problem
Automatic differentiation tools handle maps between real tensors, but functional derivatives of functionals and operators are usually derived by hand or via symbolic Euler-Lagrange forms, with no general computational tool.
Approach
Represent functions as generalized arrays and implement higher-order functions as JAX primitives, each with a JVP (Fréchet-style) rule and a transpose (adjoint) rule, so JAX's existing forward and reverse mode machinery yields functional gradients that are themselves callable Python functions.
Key results
- On the brachistochrone problem, minimizing the functional via functional derivative (minimize FD) and via FD chain rule yields curves close to ground truth, while direct parameter-space minimization overfits the numerical grid.
- Function-call caching turns exponential cost of nested composition f(h)+g(h) into linear cost in nesting depth.
- A two-layer nonlocal neural functional trained with functional gradient descent (lr 0.1, 4 steps) shows steadily decreasing L2 loss, though the learned kernel grows prohibitively large in the computation graph.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use parameter-space optimization of F instead of a functional-derivative-based objective | Brachistochrone MLP · Brachistochrone | curve fidelity | Hurt |
| Use functional derivatives in the optimization objective | Brachistochrone MLP · Brachistochrone | curve fidelity | Helped |
| Minimize the magnitude of the functional gradient | Brachistochrone MLP · Brachistochrone | functional gradient magnitude | Helped |
| Use function-call caching for repeated calls with the same argument | AutoFD · Nested function composition | execution time | Helped |
| Apply functional gradient descent to the neural functional parameters | Nonlocal neural functional · Nonlocal neural functional | L2 loss | Helped |
Surprising findings
- The learned kernel expands in the computation graph so quickly that the step-4 JAXPR graph could not be rendered, limiting practical functional gradient descent to few steps.
Limitations
- Transpose rules for compose and linear transpose require function inversion, which JAX does not support, so those rules are only defined case by case.
- Integration is numerical on a fixed grid, so accuracy is bounded by the chosen quadrature.
- Mixed-order differentiation graphs cannot be jitted, and static shape annotation via jaxtyping is required.
Experimental setup
Benchmarks: brachistochrone problem exchange-correlation functional (gga_x_pbe) usage example nonlocal neural functional training
Metrics: loss functional gradient magnitude execution time with and without caching
Baselines: direct minimization in parameter space (minimize F) Euler-Lagrange functional derivative (minimize F via FD)
Builds on: JAX Fréchet derivative Euler-Lagrange equation Fourier neural operator
Closest papers by meaning
- FunctionEvolve: Structure-Guided Symbolic Regression with LLMs (NeurIPS 2026 · similarity 0.80)
- APIGen: Automated PIpeline for Generating Verifiable and Diverse Function-Calling Datasets (NeurIPS 2024 · similarity 0.79)
- Symbolic Functional Decomposition: A Reconfiguration Approach (AAAI 2025 · similarity 0.79)
- Minimal Macro-Based Rewritings of Formal Languages: Theory and Applications in Ontology Engineering (and Beyond) (AAAI 2024 · similarity 0.78)
- JaxMARL: Multi-Agent RL Environments and Algorithms in JAX (NeurIPS 2024 · similarity 0.78)
- The Berkeley Function Calling Leaderboard (BFCL): From Tool Use to Agentic Evaluation of Large Language Models (ICML 2025 · similarity 0.78)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.