🔢 mathematics

Learning Gradient Flow: Using Equation Discovery to Accelerate Engineering Optimization

This paper introduces the Learned Gradient Flow (LGF) optimizer, which accelerates engineering optimization by using data-driven equation discovery to learn surrogate models of continuous-time dynamics from trajectory data, thereby avoiding expensive objective and gradient evaluations.

Original authors: Grant Norman, Conor Rowan, Kurt Maute, Alireza Doostan

Published 2026-06-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Grant Norman, Conor Rowan, Kurt Maute, Alireza Doostan

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 find the lowest point in a vast, foggy valley. You can’t see the bottom, so you have to feel the ground with your feet to figure out which way slopes downward. Every time you take a step, you have to stop, carefully measure the slope, and then decide where to step next. This is how traditional computer optimization works: it’s accurate, but it’s incredibly slow because it has to do that heavy "measuring" work over and over again.

This paper introduces a clever shortcut called Learned Gradient Flow (LGF). Instead of measuring the slope every single step, the computer learns to predict the path.

Here is how it works, broken down into simple concepts:

1. The "Ghost" Path

The authors realized that when a computer tries to solve a problem, the variables it changes (like the shape of a bridge or the temperature in a room) move in a smooth, predictable pattern, much like a ball rolling down a hill. This movement is called a "gradient flow."

Instead of calculating the exact physics of the hill at every millisecond, the LGF method watches the computer take the first few steps. It records where the variables go. Then, it uses a mathematical technique called SINDy (which stands for Sparse Identification of Nonlinear Dynamics) to find the simple mathematical rule that describes that movement.

Think of it like watching a dancer take the first three steps of a routine. Once you see the pattern, you don’t need to watch them step-by-step to know where they’ll be in five seconds; you can just predict their position based on the rhythm you’ve learned.

2. The Surrogate Model

Once the computer has learned this "rhythm" or rule, it creates a surrogate model. This is a lightweight, fast version of the problem.

The Old Way: The computer solves the complex physics equation (the expensive part) to find the next step.
The LGF Way: The computer uses the simple rule it learned to "fast-forward" through many steps at once without solving the heavy physics equations.

It’s like switching from walking up a mountain one step at a time to taking a cable car that follows the same path but gets you to the next checkpoint much faster.

3. Checking the Map (Scheduled Retraining)

Of course, the valley might change shape, or the prediction might get slightly off course. If the computer just guessed forever, it might end up in the wrong place.

To fix this, the LGF method uses a strategy called Scheduled Retraining. It works in cycles:

Learn: Take a few real, careful steps to update the prediction rule.
Predict: Use the rule to jump ahead many steps quickly.
Check: Stop, take one real step to see if the prediction was right, and update the rule if needed.

This ensures the computer stays on the right track without wasting time measuring every single inch of the journey.

4. Handling Huge Problems

Sometimes, the problems are so complex that there are thousands of variables moving at once (like designing a bridge with millions of tiny parts). Tracking all of them is too much memory for the computer.

The paper shows that even in these huge cases, the movement often follows a simple pattern in a lower-dimensional space. It’s like watching a flock of birds: even though there are hundreds of birds, they move as a single, cohesive shape. The LGF method finds that simple shape, predicts where the flock will go, and then translates that back to the individual birds. This saves massive amounts of computing power.

What Did They Test It On?

The authors tested this method on five different engineering and science problems:

Finding Material Properties: Figuring out the conductivity of a heated bar by looking at temperature data.
Topology Optimization: Designing the most efficient structure for a bridge or part.
Heat Conduction: Solving complex heat transfer problems.
Full Wave Inversion: Used in geophysics to understand underground structures by analyzing how waves travel through them.
Neural Network Training: Training AI models to solve physics problems (using a method called the Deep Ritz Method).

The Results

The method worked surprisingly well. In some cases, it sped up the optimization process by 200% to 250%. This means the computer reached the correct answer in a third or a quarter of the time it would have taken using traditional methods.

However, the authors warn that you can’t just "fast-forward" too aggressively. If you try to jump too far ahead without checking in, the prediction becomes inaccurate, and the computer might get lost. You have to find the right balance between learning and predicting.

In Summary

This paper shows that by treating optimization problems like dynamic systems (like a ball rolling down a hill) and using data to learn the "rules of motion," we can create fast, lightweight shortcuts. These shortcuts allow computers to solve complex engineering and scientific problems much faster by skipping the tedious, expensive calculations at every single step.

Technical Summary: Learning Gradient Flow: Using Equation Discovery to Accelerate Engineering Optimization

Problem Statement
Gradient-based optimization is a fundamental tool in engineering design and computational science, yet it is often computationally prohibitive. Standard algorithms, such as gradient descent, Newton’s method, and ADAM, require repeated evaluations of the objective function and its gradients. In engineering contexts, these evaluations often necessitate solving large-scale partial differential equation (PDE) systems, leading to high costs per iteration. While surrogate modeling is a common strategy to accelerate optimization, existing approaches typically replace the entire optimization pipeline or approximate the loss landscape directly. This work addresses the specific challenge of accelerating the iterative update process itself by leveraging the continuous-time dynamical systems perspective of optimization.

Methodology
The authors propose the Learned Gradient Flow (LGF) optimizer, which treats the evolution of optimization variables as a dynamical system. The core premise is that in the limit of vanishing step size, discrete optimization updates correspond to continuous-time ordinary differential equations (ODEs), known as gradient flows. Instead of computing expensive gradients at every step, LGF uses data-driven equation discovery to learn a surrogate model of these dynamics.

The methodology relies on the Sparse Identification of Nonlinear Dynamics (SINDy) framework. The process involves two phases:

Data-Generating Phase: The optimizer runs for $K$ iterations using the true, expensive gradient evaluations. The trajectory of the optimization variables (state vector $\mathbf{a}$ ) is recorded.
Surrogate Prediction Phase: Using the recorded trajectory, SINDy identifies a sparse coefficient matrix $\Xi$ that approximates the time derivative of the state variables $\dot{\mathbf{a}}$ as a linear combination of candidate library functions (typically polynomials) $\phi(\mathbf{a})$ . The resulting surrogate ODE, $\frac{d\mathbf{a}}{dt} = \hat{f}(\mathbf{a}) = \Xi^T \phi(\mathbf{a})$ , is then integrated forward in time for $M$ total steps ( $M > K$ ). This allows the optimizer to skip $M-K$ expensive gradient evaluations.

To handle high-dimensional problems where the state space is too large for direct SINDy application, the authors employ dimensionality reduction. The optimization history is projected into a low-rank latent space using Singular Value Decomposition (SVD). The dynamics are learned in this reduced space and then lifted back to the original space.

The framework is generalized to different optimizers:

Gradient Descent: The surrogate directly models $\dot{\mathbf{a}} \approx -\nabla z(\mathbf{a})$ .
Newton’s Method: The surrogate models the continuous-time Newton flow, $\dot{\mathbf{a}} = -(\nabla^2 z)^{-1} \nabla z$ , using only state variable data.
ADAM: Since ADAM depends on momentum and variance history variables ( $m$ and $v$ ), the system is expanded to include these variables. The authors use a "partially known" ODE approach, where SINDy is used only to model the expensive gradient term $\nabla z(\mathbf{a})$ , while the rest of the ADAM update rules are retained analytically.

To maintain accuracy over long horizons, the authors introduce scheduled retraining, where the surrogate model is discarded and rebuilt every $M$ iterations.

Key Contributions

Proposal of SINDy-based Surrogates: The authors introduce the use of SINDy to model the continuous-time dynamics of unconstrained optimization problems.
Generalization to Multiple Optimizers: The methodology is extended beyond gradient descent to include Newton’s method and the ADAM optimizer.
Development of LGF Optimizer: A PyTorch-based optimizer is developed that builds SINDy models from optimization trajectory data in real-time, supporting both full-dimensional and reduced-dimensional spaces.
Demonstration on Engineering Problems: The efficacy of the method is showcased on five standard problems from engineering mechanics and scientific machine learning.

Results
The LGF optimizer was tested on five numerical examples:

1D Dynamic Inverse Problem: Using gradient descent, LGF achieved a 200% acceleration (saving 2 out of 3 gradient evaluations) while recovering material parameters with accuracy comparable to standard gradient descent.
3D Topology Optimization: For a high-dimensional problem (13,824 variables), dimensionality reduction was used to project the dynamics into a 2D latent space. LGF achieved 100% and 250% acceleration in separate tests, producing density fields nearly identical to the reference solution.
Nonlinear Heat Conduction: Using Newton’s method, LGF achieved a 33% acceleration. The surrogate solution was indistinguishable from the true Newton method, with a normalized difference in temperature fields of $1.7 \times 10^{-13}$ . The authors noted that overly aggressive acceleration (e.g., 100%) led to divergence in this case.
Full Wave Inversion (FWI): Using ADAM, LGF achieved a 50% acceleration. The recovered slowness field had a relative mean-squared error comparable to standard ADAM, demonstrating robustness in a non-convex landscape.
Deep Ritz Method (Physics-Informed ML): LGF was applied to a stochastic optimization problem using Monte Carlo integration for the loss function. Using ADAM with mini-batching, LGF achieved a 43% acceleration, capturing relevant features of the optimization dynamics despite noisy gradients.

Significance and Claims
The paper claims that learned gradient flows can significantly expedite convergence by capturing critical features of the optimization trajectory while avoiding expensive evaluations of the objective and its gradient. The authors emphasize that the key contribution is the integration of a learned dynamical system within the optimization loop, allowing developments in data-driven ODE modeling to be repurposed for optimization.

The authors acknowledge limitations, noting that the method is not universally applicable. It may struggle with constrained optimization problems where non-smooth behaviors (e.g., active-set updates) disrupt the smooth dynamics assumed by the surrogate. They also highlight that the selection of hyperparameters, particularly the retrain interval $M$ , currently relies on trial and error and requires further study to automate. The work positions LGF as a general framework where the model identification step can be replaced by any procedure that infers a surrogate ODE from optimization history.