L2O via Deep Unfolded Flows

-->

Learning-to-Optimize via Deep Unfolded Flows

Augustinos D. Saravanos1,

Oswin So1,

H. M. Sabbir Ahmad2,

Chuchu Fan1

1Massachusetts Institute of Technology &nbsp&nbsp&nbsp&nbsp2Boston University

Paper

arXiv

Supplementary Video

Presentation Video

-->

Code

Optimization

Gradient-based

Sampling-based

Machine Learning

Learning-to-<br>Optimize (L2O)

Generative<br>Modeling

Flow<br>Optimizer

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 (&times;10) ones.

-->

How to learn-to-optimize effectively?

Goal:<br>Learn an optimizer that reaches high-quality solutions fast for a given problem class

Optimization point of view:

Gradient-based optimization<br>✓ Uses local derivative information<br>✗ May get stuck in local minima

Sampling-based optimization<br>✓ Exploration, can avoid local minima<br>✗ Struggles to scale to high dimensions

Machine Learning point of view:

Learning-to-optimize<br>Learn efficient optimization algorithm steps from data

Generative modeling<br>Learn to generate samples from multimodal complex distributions

Key Idea:<br>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

Optimization problem<br>Solve unconstrained optimization problems from a given problem class

\[ \min_{x\in\mathbb{R}^n} f(x) \]

Learn a population update<br>We wish to learn a model that transforms populations of points

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

Velocity field<br>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<br>Velocity field NN is conditioned on population objectives, ranking, gradient information, population statistics, etc.

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

Initial training phase<br>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

Unroll FlowOptimizer for \(K\) iterations as sequential layers.

Extended context including history of best points, population statistics history, etc.

Fine-tuning phase<br>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<br>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<br>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\) &rarr; recovers first-order and learning-to-optimize methods

No base sampler, \(N>1\) &rarr; resembles particle swarm optimization (PSO)

Gaussian base sampler &rarr; 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

Real-World Problems

FlowOptimizer also outperforms baselines in real-world non-convex...