Research map / Reinforcement learning / Visual and representation RL
Optimizing Automatic Differentiation with Deep Reinforcement Learning
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.
Proposes AlphaGrad, which casts cross-country elimination for automatic differentiation as a single-player RL game (VertexGame) solved with AlphaZero and a transformer, minimizing multiplications, plus a JAX interpreter Graphax for runtime gains.
Problem
Finding the elimination order in cross-country elimination that minimizes the number of multiplications needed to compute exact Jacobians via automatic differentiation; existing heuristics (Markowitz, forward/reverse mode) and prior search methods are suboptimal.
Approach
Formulate vertex elimination ordering as a single-player game (VertexGame) with reward equal to negative multiplication count, and train an AlphaZero agent with a transformer policy/value network over compressed adjacency-tensor tokens to discover new AD algorithms tailored to each function.
Key results
- AlphaGrad reduces multiplications versus minimal Markowitz on all tasks, e.g. random function f 6374 vs 12083 (about 47% fewer) and RoeFlux_3d 811 vs 938, with improvements up to 33% reported.
- On RoeFlux_3d, Graphax runtime with AlphaGrad orders is 63.92 ms vs 91.27 ms for Markowitz at batch size 512 on CPU.
- For 2-layer MLP and Transformer encoder with 50 MCTS simulations, AlphaGrad failed to beat reverse-mode AD; with 250 simulations it marginally outperformed it (389 vs 392 multiplications).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use AlphaGrad's learned elimination order instead of the strongest baseline | Jacobian computation · RoeFlux_1d | multiplications: 407 -> 320 | Helped |
| Use AlphaGrad's learned elimination order instead of the strongest baseline | Jacobian computation · RobotArm_6DOF | multiplications: 288 -> 231 | Helped |
| Use AlphaGrad's learned elimination order instead of the strongest baseline | Jacobian computation · HumanHeartDipole | multiplications: 172 -> 149 | Helped |
| Use AlphaGrad's learned elimination order instead of the strongest baseline | Jacobian computation · PropaneCombustion | multiplications: 90 -> 88 | Helped |
| Use AlphaGrad's learned elimination order instead of the strongest baseline | Jacobian computation · Random function g | multiplications: 451 -> 417 | Helped |
| Use AlphaGrad's learned elimination order instead of the strongest baseline | Jacobian computation · BlackScholes | multiplications: 350 -> 312 | Helped |
| Use AlphaGrad's learned elimination order instead of the strongest baseline | Jacobian computation · RoeFlux_3d | multiplications: 938 -> 811 | Helped |
| Use AlphaGrad's learned elimination order instead of the strongest baseline | Jacobian computation · Random function f | multiplications: 9333 -> 6374 | Helped |
| Use AlphaGrad with 50 MCTS simulations instead of reverse-mode AD | 2-layer MLP · 2-layer MLP | multiplications: 392 -> 398 | Hurt |
| Increase AlphaGrad search from 50 to 250 MCTS simulations | 2-layer MLP · 2-layer MLP | multiplications: 398 -> 389 | Helped |
| Use AlphaGrad with 50 MCTS simulations instead of reverse-mode AD | Transformer · Transformer | multiplications: 4688 -> 4831 | Hurt |
| Increase AlphaGrad search from 50 to 250 MCTS simulations | Transformer · Transformer | multiplications: 4831 -> 4656 | No clear effect |
Surprising findings
- With only 50 MCTS simulations AlphaGrad did not beat reverse-mode AD on the deep learning tasks, which the authors attribute to backpropagation being well suited to funnel-like graphs.
- Multi-task training yielded new best results for random function f (5884) and BlackScholes (307), though single-graph training was generally better.
Limitations
- Optimizes only multiplication count, which does not capture memory access, fusion or actual runtime fully.
- Requires retraining for every computational graph; no support for dynamic control flow.
- Input tensors limited to two dimensions in VertexGame, so the transformer task could only be run at batch size 1.
Experimental setup
Datasets: RoeFlux_1d RoeFlux_3d RobotArm_6DOF HumanHeartDipole PropaneCombustion BlackScholes_Jacobian 2-layer MLP Transformer encoder random functions f and g
Benchmarks: RoeFlux_1d RoeFlux_3d RobotArm_6DOF HumanHeartDipole PropaneCombustion BlackScholes_Jacobian 2-layer MLP Transformer encoder random function f random function g
Metrics: number of multiplications median runtime (ms)
Baselines: forward-mode AD reverse-mode AD minimal Markowitz degree JAX forward-mode AD JAX reverse-mode AD
Builds on: AlphaZero Gumbel AlphaZero AlphaTensor JAX PPO
Closest papers by meaning
- Betty: An Automatic Differentiation Library for Multilevel Optimization (ICLR 2023 · similarity 0.85)
- Learning What to Defer for Maximum Independent Sets (ICML 2020 · similarity 0.85)
- Learning to Stop: Deep Learning for Mean Field Optimal Stopping (ICML 2025 · similarity 0.83)
- DIME: Diffusion-Based Maximum Entropy Reinforcement Learning (ICML 2025 · similarity 0.83)
- Revisiting Scalable Hessian Diagonal Approximations for Applications in Reinforcement Learning (ICML 2024 · similarity 0.83)
- SFT-then-RL Outperforms Mixed-Policy Methods for LLM Reasoning (NeurIPS 2026 · similarity 0.83)
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.