Learning-to-Optimize via Deep Unfolded Flows

Augustinos D. Saravanos1, Oswin So1, H. M. Sabbir Ahmad2, Chuchu Fan1
1Massachusetts Institute of Technology     2Boston University
Paper arXiv Code
Optimization Gradient-based Sampling-based Machine Learning Learning-to- Optimize (L2O) Generative Modeling Flow Optimizer
Deep-unfolded FlowOptimizer: K flow-model iterations, each conditioned on an extended context, unrolled and trained end-to-end against a training loss.
FlowOptimizer transports a population of candidate solutions down a rugged non-convex landscape, guided by a learned velocity field, converging to the global minimum over iterations.

Abstract

We introduce FlowOptimizer, a deep unfolded, flow-based framework for learned iterative optimization. Motivated by the expressiveness of flow models, we represent each optimization iteration via a velocity field that operates on a population of candidate solutions, i.e., a set of parallel iterates, conditioned on contextual information including their objective values and gradients, as well as population-level statistics. The velocity field is initially trained in a simulation-free manner by matching displacements from source populations to improved target ones obtained through sampling the objective. Subsequently, we unfold this velocity field as the internal iteration of an optimization sequence, and fine-tune it in an end-to-end manner by directly optimizing objective values over a targeted class of problems. Notably, FlowOptimizer is a self-supervised framework whose training relies solely on objective evaluations without requiring knowledge of solutions. We evaluate our approach on a series of tasks from standard non-convex optimization benchmarks to real-world problems from supply chain, robotics and power grid applications. FlowOptimizer consistently outperforms well-established sampling-based/gradient-based traditional optimization and learning-to-optimize methods by orders of magnitude in terms of solution quality. We further highlight its ability to be trained on low-dimensional problems and successfully generalize to substantially higher-dimensional (×10) ones.

How to learn-to-optimize effectively?

Goal:

Learn an optimizer that reaches high-quality solutions fast for a given problem class

Optimization point of view:
Gradient-based optimization ✓ Uses local derivative information ✗ May get stuck in local minima
Sampling-based optimization ✓ Exploration, can avoid local minima ✗ Struggles to scale to high dimensions
Machine Learning point of view:
Learning-to-optimize Learn efficient optimization algorithm steps from data
Generative modeling Learn to generate samples from multimodal complex distributions
Key Idea:

Represent optimization steps as a flow model (learned velocity field) over a population of candidate solution points

FlowOptimizer: A flow model as the core optimization step

One FlowOptimizer core iteration: the current population is passed through an optional base sampler to z0, then a context-conditioned flow model integrates the velocity field to z1, producing the next population; the context c^k aggregates objective values, gradients, and population statistics.
Optimization problem

Solve unconstrained optimization problems from a given problem class

\[ \min_{x\in\mathbb{R}^n} f(x) \]
Learn a population update

We wish to learn a model that transforms populations of points

\[ \mathcal{X}^{k+1} = \mathcal{T}(\mathcal{X}^k,\, c^k) \]
Velocity field

Represent via a velocity field ODE that operates on the populations

\[ \dot{z}_t = v_\theta(z_t,\, t,\, c^k), \quad z_{t=0}=\mathcal{X}^k \]
Conditioned on Context

Velocity field NN is conditioned on population objectives, ranking, gradient information, population statistics, etc.

Permutation Invariance

Velocity field is parametrized with a self-attention architecture that operates jointly on the population state and the local context

Initial training phase

Flow matching loss to match displacements to improved populations

\[ \mathcal{L}_{\mathrm{FM}} = \mathbb{E}_{t,z_0,z_1}\, \lVert v_\theta(z_t,t,c) - \Delta z \rVert_2^2 \]

Deep Unfolded FlowOptimizer

Deep-unfolded FlowOptimizer: K flow-model iterations, each conditioned on an extended context, unrolled as sequential layers and trained end-to-end against a training loss.

Unroll FlowOptimizer for \(K\) iterations as sequential layers.
Extended context including history of best points, population statistics history, etc.

Fine-tuning phase

Training loss directly minimizes the weighted sum of all iteration objectives

\[ \mathcal{L}_{\mathrm{FT}} = \sum_{k=1}^{K} w_k\, \ell_k(\mathcal{X}^k) \]
Per-iteration loss

Combines mean and best-sample (softmin) values of populations

\[ \begin{aligned} \ell_k(\mathcal{X}^k) = {}&\alpha\,\mathrm{Best}\big(f(\mathcal{X}_k)\big) \\ &+ (1-\alpha)\,\mathrm{Mean}\big(f(\mathcal{X}_k)\big) \end{aligned} \]
Self-supervised

Training happens in a self-supervised manner. We only rely on evaluating the populations at each iteration without requiring any known solutions!

Connections to classical and learned optimizers
  • No base sampler, \(N=1\) → recovers first-order and learning-to-optimize methods
  • No base sampler, \(N>1\) → resembles particle swarm optimization (PSO)
  • Gaussian base sampler → resembles CEM and evolutionary strategies

Standard Non-Convex Optimization Benchmarks

Train and evaluate on challenging non-convex optimization benchmarks
FlowOptimizer reaches high-quality solutions much faster than gradient/sampling/L2O-based baselines
3D surface plots of the Ackley, Rastrigin, and Levy non-convex benchmark functions.
Results on Ackley, Rastrigin, and Levy (20D): FlowOptimizer reaches orders-of-magnitude lower best-f error in far less wall-clock time than CEM, CMA-ES, MS-GD, MS-NAG, MS-Adam, L2O-GD, PSO, and RS baselines; error histograms show its error distribution concentrated well below all baselines.

Real-World Problems

FlowOptimizer also outperforms baselines in real-world non-convex optimization problems.
Problem Class Dim. \(n\) FlowOpt (ours) RS CEM CMA-ES MSGD MSNAG MS-Adam L2O-GD PSO
Robotic Arm 100.0030.180.090.0080.0140.0210.0090.0120.005
300.0070.310.180.040.180.170.070.120.011
Power Grid 205e-40.530.140.030.050.030.0070.030.02
501.2e-30.830.310.110.230.170.050.080.09
Supply Chain 400.0070.470.090.050.070.030.020.0130.02
800.0130.680.140.060.090.030.020.0270.03

Scalability to high-dimensional problems

Train on 20D problems and then evaluate on 200D ones from the same class
FlowOptimizer maintains its performance on higher dimensions than the ones used for training
Ackley 200D and Levy 200D convergence: trained on 20D and evaluated on 200D, FlowOptimizer still reaches orders-of-magnitude lower best-f error in far less wall-clock time than all baselines, which stagnate.

Key Take-Aways: Why FlowOptimizer works well?

✶ Flows for optimization

Highly expressive for optimization steps over populations of samples

✶ Learning problem geometry

Learns recurrent patterns of the geometry of the underlying problem class

✶ Population interactions

Operates on populations and learns interactions among members

BibTeX

@inproceedings{
saravanos2026learningtooptimize,
title={Learning-to-Optimize via Deep Unfolded Flows},
author={Augustinos D Saravanos and Oswin So and H M Sabbir Ahmad and Chuchu Fan},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=ZOtOq7hxJP}
}