🤖 machine learning

On the Convergence of Jacobian-Free Backpropagation for Optimal Control Problems with Implicit Hamiltonians

This paper establishes convergence guarantees for Jacobian-Free Backpropagation (JFB) in stochastic minibatch settings for optimal control problems with implicit Hamiltonians and demonstrates its scalability to high-dimensional multi-agent and swarm control tasks.

Original authors: Eric Gelphman, Deepanshu Verma, Nicole Tianjiao Yang, Stanley Osher, Samy Wu Fung

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Eric Gelphman, Deepanshu Verma, Nicole Tianjiao Yang, Stanley Osher, Samy Wu Fung

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a fleet of 100 autonomous delivery drones how to navigate a crowded city to deliver packages as quickly and safely as possible.

This paper is essentially a "math manual" that proves a new, highly efficient way to train these drones—even when the rules of the world are so complex that you can't write down a simple formula for the "perfect" flight path.

Here is the breakdown of the paper using everyday analogies.

1. The Problem: The "Invisible Maze" (Implicit Hamiltonians)

In traditional robotics, we usually have a "cheat sheet." If a drone needs to turn, we have a math formula that says: "If you are at point A and moving at speed B, turn exactly 15 degrees to be optimal." This is a closed-form solution.

However, in many real-world scenarios (like the "consumption-savings" or "swarm" problems mentioned in the paper), the rules are implicit. It’s like trying to navigate a maze where the walls move based on how you move. There is no "cheat sheet" formula. To find the best move, you have to constantly solve a mini-puzzle: "If I do this, what happens to the environment, and how does the environment then change my future options?"

Because you can't just look up the answer, training a computer to learn this is incredibly slow and memory-intensive. It’s like trying to learn to play chess by calculating every possible move for the next ten turns every single time you touch a piece.

2. The Old Way: The "Heavy Backpack" (Automatic Differentiation)

The old way to teach these drones was called Automatic Differentiation (AD). Imagine the drone is learning by replaying its entire flight from start to finish. To learn from a mistake at the end of the flight, the computer has to remember every single tiny muscle twitch and wind gust from the beginning of the trip.

As the number of drones increases (from 1 to 100), the "memory" required becomes massive. It’s like a student trying to take a final exam while carrying a backpack that gets heavier with every single page they read. Eventually, the backpack becomes too heavy to move (the computer runs out of memory).

3. The New Way: "Jacobian-Free Backpropagation" (JFB)

The authors use a technique called JFB. Instead of carrying that heavy backpack of "every single detail," JFB is like a student who uses smart intuition instead of perfect memory.

Instead of calculating the exact, heavy mathematical "weight" of every single decision (the Jacobian), JFB uses a "good enough" approximation. It’s like a pilot who doesn't calculate the exact air pressure of every molecule hitting the wing, but instead feels the vibration of the plane and says, "I think I need to tilt left." It is much faster and uses way less "brainpower" (memory).

4. The Paper’s Big Achievement: "The Guarantee"

The problem with "good enough" methods is that they can be unreliable. In the past, people knew JFB worked in small tests, but they couldn't prove it would work when things got big and messy. It was like saying, "I've driven this car on a flat road, so I'm sure it will work in a hurricane."

This paper provides the mathematical proof that the "hurricane" won't break the system.

The authors proved three things:

The "Compass" Works: Even though JFB is using an approximation (a "biased" estimate), they proved that on average, the direction it points is still the right way to go. It’s like a compass that is slightly off, but as long as it always points generally North, you will eventually reach your destination.
It Scales: They tested it on massive problems—like 100 drones or 100 people managing their money—where the old "heavy backpack" methods completely crashed.
It’s Fast: They showed that JFB reaches the "perfect" behavior much faster in real-world time than the old, precise methods.

Summary

In short: We found a way to teach complex, high-speed systems how to behave optimally without needing a "cheat sheet" or a "supercomputer's memory," and we mathematically proved that this shortcut won't lead the system astray.

Technical Summary: On the Convergence of Jacobian-Free Backpropagation for Optimal Control Problems with Implicit Hamiltonians

1. Problem Statement

The paper addresses a fundamental challenge in learning-based optimal control: optimal feedback control with implicit Hamiltonians.

In many high-dimensional control problems, the Pontryagin Maximum Principle (PMP) provides the necessary conditions for optimality, but the resulting Hamiltonian maximization problem— $\text{arg max}_u H(t, z, p, u)$ —does not admit a closed-form solution. Traditional value function-based methods (which parameterize the value function to recover feedback laws) rely heavily on having an explicit expression for this maximizer. When the maximizer is implicit, existing methods become computationally prohibitive due to the need for expensive implicit differentiation (solving linear systems involving the Jacobian) at every time step of every trajectory during training.

2. Methodology

The authors build upon a recent framework (Gelphman et al., 2025) that utilizes Implicit Neural Networks (INNs) and Jacobian-Free Backpropagation (JFB).

Implicit Neural Networks (INNs): Instead of standard feedforward layers, the control $u^*$ is defined as the fixed point of a learnable operator $T_\theta$ derived from the Hamiltonian optimality condition: $u^*_\theta = T_\theta(u^*_\theta; t, z)$ .
Jacobian-Free Backpropagation (JFB): To avoid the $O(n^3)$ cost of inverting the Jacobian during backpropagation, JFB uses a zeroth-order Neumann approximation. It replaces the expensive implicit gradient $\frac{du^*_\theta}{d\theta} = (I - \frac{\partial T_\theta}{\partial u})^{-1} \frac{\partial T_\theta}{\partial \theta}$ with a simplified version: $\frac{du^*_\theta}{d\theta} \approx \frac{\partial T_\theta}{\partial \theta}$ .
Stochastic Optimization: The training is formulated as minimizing the expected cost over a distribution of initial conditions $\rho$ using Stochastic Gradient Descent (SGD).

3. Key Contributions

The paper provides three primary advancements over previous work:

Theoretical Convergence Guarantees: The most significant contribution is the first convergence analysis for JFB in a stochastic minibatch setting. Previous work only established "sample-wise descent" (that a single update moves in a descent direction). This paper proves that JFB-based SGD converges to stationary points of the expected objective, despite the gradient estimates being systematically biased.
Theoretical Verification: The authors provide empirical evidence that the mathematical assumptions required for their convergence proofs (such as operator contractivity, matrix conditioning, and variance bounds) actually hold in practical training regimes.
High-Dimensional Scalability: They demonstrate that JFB can scale to problems where standard automatic differentiation (AD) fails due to memory constraints (the "memory wall"), specifically targeting problems with hundreds of agents.

4. Results

Convergence Proofs: Through a series of lemmas, the authors prove that if the learning rate $\alpha_j$ satisfies standard Robbins-Monro conditions ( $\sum \alpha_j = \infty, \sum \alpha_j^2 < \infty$ ), the weighted average of the squared expected gradients converges to zero in expectation.
Efficiency vs. AD and CVXPYLayers: In low-to-moderate dimensional benchmarks (e.g., single and six quadrotors), JFB was shown to be significantly faster in wall-clock time than traditional Automatic Differentiation (AD) and more accurate than optimization-based layers (CVXPYLayers), while maintaining a much lower memory footprint.
Scalability to 100-Agent Problems: The authors successfully trained controllers for:
- 100-agent consumption-savings models (100D state/control).
- Swarm quadrotor control (1200D state/400D control).
- Multi-bicycle dynamics (400D state/200D control).
  In these cases, JFB maintained constant memory usage per epoch, whereas AD would have required storing the entire computational graph, leading to memory exhaustion.

5. Significance

This work bridges the gap between control theory and deep learning by providing a mathematically rigorous foundation for using implicit models in high-dimensional settings. By proving that the "shortcut" taken by JFB (ignoring the Jacobian inverse) does not prevent convergence to an optimal solution, the authors provide a principled tool for engineers to design feedback controllers for complex, large-scale systems (like swarms or multi-agent economic models) where closed-form solutions are impossible to obtain.