Scalable Deep Unfolding of Conic Optimizers

원문 초록

Deep unfolding (DU) accelerates iterative optimizers by training learnable components through unrolled iterations, but extending DU to the large-scale semidefinite programs (SDPs) common in robotics has remained limited. Unrolling a full-update conic solver such as COSMO exposes two obstacles: backpropagating through the per-iteration linear-system solve incurs memory quadratic in the problem size when the coefficient matrix is formed explicitly, and backpropagating through the positive semidefinite (PSD) cone projection is numerically unstable under repeated eigenvalues. We address the first with a matrix-free implicit differentiation rule that reduces memory from $O(n^2)$ to $O(n + \mathrm{nnz}(A))$, and the second with a stable backward rule based on the Dale\v{c}kii--Krein representation of the Fr\'echet derivative. Together, these enable learning lightweight hyperparameter policies and warm starts for COSMO. On standalone SDPs and second-order cone programs, ranging from max-cut and Lov\'asz $\vartheta$ to robust estimation and control, the learned policies outperform tuned COSMO on every problem class, reducing iterations to tolerance by up to 60$\times$, and transfer without retraining to instances up to 10$\times$ larger in the size parameter while retaining 2--5$\times$ speedups. Our matrix-free solver implementation also batches naturally on the GPU, solving 64 Lov\'asz $\vartheta$ SDPs 14$\times$ faster than sequential COSMO.jl. When embedded in sequential convex programming for nonlinear covariance steering, our learned optimizers outperform the interior-point solver Clarabel in terms of both success rate and solve time.