🔢 mathematics

Scalable Deep Unfolding of Conic Optimizers

This paper introduces a scalable deep unfolding framework for large-scale semidefinite programs that overcomes memory and numerical stability barriers through matrix-free implicit differentiation and a robust eigenvalue-aware backward rule, enabling learned policies that achieve up to 50 $\times$ speedups over state-of-the-art conic solvers.

Original authors: Alex Oshin, Rahul Vodeb Ghosh, Evangelos A. Theodorou

Published 2026-06-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Alex Oshin, Rahul Vodeb Ghosh, Evangelos A. Theodorou

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 solve a massive, complex puzzle. In the world of robotics and engineering, these puzzles are called optimization problems. They are used to figure out the best way for a robot to move, how to steer a car safely, or how to manage a power grid.

For a long time, computers have used "iterative optimizers" to solve these puzzles. Think of these optimizers as a very methodical, but slow, hiker trying to find the bottom of a valley. They take one step, check if they are lower, take another step, and repeat this thousands of times until they reach the bottom.

Deep Unfolding is a new way to teach this hiker to run instead of walk. Instead of just following a rigid set of rules, the hiker is given a "coach" (a neural network) that learns from experience. The coach tells the hiker exactly how big of a step to take and when to change direction, based on what worked best in previous puzzles. This paper is about teaching this coach to handle the biggest, hardest puzzles of all.

Here is the breakdown of the paper's story, using simple analogies:

The Problem: The "Memory Wall" and the "Sticky Floor"

The researchers tried to apply this "coach" system to a specific type of solver called COSMO, which is great for large-scale problems. However, they hit two massive roadblocks that stopped them from teaching the coach effectively:

The Memory Wall (The Linear System):
To take a step, the solver has to solve a giant math equation involving a huge grid of numbers (a matrix). To teach the coach, the computer needs to remember how it solved that equation so it can learn from its mistakes later.
- The Old Way: It was like trying to remember every single grain of sand on a beach to figure out how to walk on it. As the puzzle got bigger, the computer's memory (RAM) exploded and crashed. It was an $O(n^2)$ problem—doubling the size of the puzzle quadrupled the memory needed.
- The Paper's Fix: They invented a "Matrix-Free" trick. Instead of writing down the whole grid of numbers, they realized they only needed to know how the grid reacts to a single push (a matrix-vector product). It's like learning to walk on the beach by feeling the sand under your feet as you step, rather than trying to memorize the entire beach map. This reduced the memory needed from a massive warehouse to a small backpack ( $O(n)$ ), allowing the system to handle puzzles that were previously impossible.
The Sticky Floor (The Eigenvalue Problem):
Some puzzles involve a special shape called a "PSD cone." To solve them, the computer has to look at the puzzle's "eigenvalues" (think of these as the unique frequencies or tones of the puzzle).
- The Old Way: When two of these tones are exactly the same (repeated eigenvalues), the math used to teach the coach breaks down. It's like trying to calculate the slope of a floor that is perfectly flat; the math says "divide by zero," and the computer crashes or gives nonsense answers. This happened all the time in their specific robotics problems.
- The Paper's Fix: They used a sophisticated mathematical tool called the Daleckii–Krein formula. Think of this as a special "smoothie blender" for math. Instead of getting stuck on the flat spots, this formula knows exactly how to handle the situation where two tones are identical, keeping the math stable and the learning process going.

The Result: The Super-Runner

Once they fixed these two roadblocks, they trained their "coach" to guide the COSMO solver.

The Speedup: The learned solver became incredibly fast. In some tests, it solved problems 50 times faster than the standard, untrained solver.
Real-World Test: They tested this on a "Covariance Steering" problem. Imagine a robot trying to steer a cloud of uncertainty (like a swarm of bees) from point A to point B without hitting anything. When this new solver was used as a helper inside a larger planning system, it made the whole process 30 times faster.
Comparison: It even competed with the "gold standard" solvers (like Clarabel) that are usually considered the best, but did so much faster for the specific types of problems robots face in real-time.

Summary

This paper didn't invent a new robot or a new type of math problem. Instead, it fixed the "engine" that solves these problems.

They removed the memory bottleneck so the engine can run on huge puzzles without running out of gas.
They fixed the mathematical instability so the engine doesn't stall when the road gets tricky.

The result is a "learned" optimizer that acts like a seasoned veteran hiker who knows exactly how to navigate the terrain, solving complex robotics problems in a fraction of the time it used to take.

Technical Summary: Scalable Deep Unfolding of Conic Optimizers

Problem Statement

Deep Unfolding (DU) has proven effective in accelerating iterative optimizers for various domains, including compressed sensing and robotics. However, extending DU to large-scale semidefinite programs (SDPs)—common in robotics applications like covariance steering—has been hindered by two specific obstacles when unrolling full-update ADMM-based conic solvers (e.g., COSMO):

Memory Bottleneck in Linear-System Solves: Backpropagating through the linear-system solve in each iteration typically requires forming the coefficient matrix explicitly. For ADMM, this involves the term $A^\top A$ . Even if the constraint matrix $A$ is sparse, $A^\top A$ is often dense, leading to $O(n^2)$ memory complexity. This prevents scaling to large problem sizes where direct factorization runs out of memory.
Numerical Instability in PSD Projections: The projection onto the positive semidefinite (PSD) cone requires an eigendecomposition. Standard backpropagation through this operation involves terms of the form $1/(\lambda_i - \lambda_j)$ in the gradient. When eigenvalues coincide (a frequent occurrence during training, especially with isotropic covariance targets), these terms cause numerical instability or undefined gradients (NaNs).

Methodology

The authors propose a deep unfolding framework for the COSMO solver that addresses these two obstacles through custom backward rules, enabling the learning of hyperparameter policies and warm-starts for large-scale conic problems.

1. Matrix-Free Implicit Differentiation for Linear Solves

To avoid the $O(n^2)$ memory cost of forming the dense $A^\top A$ matrix, the authors employ a matrix-free implicit differentiation rule for the conjugate gradient (CG) based linear-system solve.

Forward Pass: The solver uses the indirect mode, solving the reduced normal equations $(P + \sigma I + A^\top D(\rho)A)\tilde{x} = \xi$ using CG, which relies solely on matrix-vector products (matvecs).
Backward Pass: Instead of unrolling the CG iterations or forming the dense Jacobian, the authors derive a rule based on the adjoint system. By solving the adjoint system $M\tilde{\lambda} = -\nabla_{\tilde{x}}L$ using the same CG method, the gradient with respect to the hyperparameters $\rho$ is computed as $\nabla_\rho L = (A\tilde{\lambda}) \odot (A\tilde{x})$ .
Result: This approach never explicitly forms the dense coefficient matrix, reducing memory complexity from $O(n^2)$ to $O(n + \text{nnz}(A))$ and enabling backpropagation at scales where previous methods fail.

2. Stable Backward Rule for PSD Projection

To resolve the instability of eigendecomposition gradients under repeated eigenvalues, the authors replace standard autograd with a backward rule based on the Dalečii–Krein (DK) representation of the Fréchet derivative.

Mechanism: The PSD projection is treated as a primary matrix function $T = U f(\Lambda) U^\top$ with $f(x) = \max(x, 0)$ . The DK theorem provides a formula for the derivative where the singular term $1/(\lambda_j - \lambda_i)$ is multiplied by a numerator $f(\lambda_j) - f(\lambda_i)$ .
Stability: When eigenvalues coincide ( $\lambda_i = \lambda_j$ ), the numerator vanishes at the same rate as the denominator, resulting in a removable singularity defined by the limit $f'(\lambda_i)$ . This ensures the gradient remains finite and well-defined even with repeated eigenvalues.
Efficiency: This rule reuses the eigendecomposition computed in the forward pass, incurring no additional asymptotic computational cost compared to naive approaches, unlike alternatives such as unrolling Newton-Schulz iterations.

Key Contributions

Scalable Linear-System Differentiation: A matrix-free implicit differentiation rule that reduces memory usage from $O(n^2)$ to $O(n)$ , enabling deep unfolding of full-update conic solvers on large-scale SDPs.
Numerically Stable PSD Projection: A backward rule based on the Dalečii–Krein theorem that remains well-defined under repeated eigenvalues, a critical requirement for training on covariance steering and other SDP problems.
Empirical Validation: A demonstration that these rules allow the unrolling of the COSMO solver to learn lightweight hyperparameter policies and warm-starts. The learned approach is evaluated on standalone SDPs/SOCPs and as a subroutine in Sequential Convex Programming (SCP) for nonlinear covariance steering.

Results

The authors evaluate the learned COSMO solver against the standard COSMO solver (with adaptive heuristics) and the state-of-the-art interior-point solver Clarabel across several problem classes:

Standalone SDPs and SOCPs: The learned policy achieves significant speedups depending on the problem class:
- Lovász $\vartheta$ SDPs: Up to 50.4× speedup in wall-clock time.
- Max-cut SDPs: 2.62× speedup.
- Robust Kalman Filters: 4.91× speedup.
- Lyapunov LMIs: 1.16× speedup.
Nonlinear Covariance Steering (SCP): When used as a subroutine in SCP for nonlinear systems (double integrator and unicycle dynamics), the learned approach delivers a 30× speedup compared to vanilla COSMO.
Memory Scaling: Experiments show that while baseline methods exhaust GPU memory (24 GB) at $n \approx 20\text{--}50k$ variables, the proposed matrix-free method scales linearly, staying under 300 MB even at $n=50k$ .

Significance and Limitations

The paper claims that these contributions make it possible to apply deep unfolding to large-scale conic optimization, a domain previously limited by memory and numerical stability constraints. The learned policies provide substantial acceleration for real-time robotics applications where fixed computational budgets are required.

The authors note specific limitations:

Scalability: While the linear solve is matrix-free, the PSD projection still requires a dense $O(p^3)$ eigendecomposition for each semidefinite block of size $p$ . Scalability gains are realized in the number of variables/constraints, not the block size itself.
Accuracy Regime: The method targets low-to-moderate accuracy ( $\epsilon \approx 10^{-3}$ ), typical for real-time control. It does not claim to outperform interior-point solvers (like Clarabel) at tight tolerances where first-order methods converge slowly.
Training Overhead: A separate policy is trained per problem family, incurring an upfront cost (roughly 100 problems per class) that must be amortized over many solves. The current work does not demonstrate transfer across families or scaling to instances substantially larger than those seen during training.