A Fully First-Order Layer for Differentiable Optimization

TL;DR

This paper introduces a first-order method for differentiable optimization layers that avoids Hessian computations, significantly reducing complexity and enabling efficient gradient calculation in bilevel optimization problems.

Contribution

The authors propose a novel first-order algorithm for differentiable optimization that eliminates Hessian evaluations, with theoretical guarantees and an open-source implementation.

Findings

Achieves finite-time, non-asymptotic approximation guarantees.

Computes approximate hypergradients using only first-order information.

Matches the best known rate for non-smooth non-convex optimization.

Abstract

Differentiable optimization layers enable learning systems to make decisions by solving embedded optimization problems. However, computing gradients via implicit differentiation requires solving a linear system with Hessian terms, which is both compute- and memory-intensive. To address this challenge, we propose a novel algorithm that computes the gradient using only first-order information. The key insight is to rewrite the differentiable optimization as a bilevel optimization problem and leverage recent advances in bilevel methods. Specifically, we introduce an active-set Lagrangian hypergradient oracle that avoids Hessian evaluations and provides finite-time, non-asymptotic approximation guarantees. We show that an approximate hypergradient can be computed using only first-order information in $\tilde{\oo}(1)$ time, leading to an overall complexity of $\tilde{\oo}(\delta^{-1}\epsilon^{-3})$ for constrained bilevel optimization, which matches the best known rate for non-smooth non-convex optimization. Furthermore, we release an open-source Python library that can be easily adapted from existing solvers. Our code is available here: https://github.com/guaguakai/FFOLayer.