  Embedding parameterized optimization problems as layers into machine learning
architectures serves as a powerful inductive bias. Training such architectures
with stochastic gradient descent requires care, as degenerate derivatives of
the embedded optimization problem often render the gradients uninformative. We
propose Lagrangian Proximal Gradient Descent (LPGD) a flexible framework for
training architectures with embedded optimization layers that seamlessly
integrates into automatic differentiation libraries. LPGD efficiently computes
meaningful replacements of the degenerate optimization layer derivatives by
re-running the forward solver oracle on a perturbed input. LPGD captures
various previously proposed methods as special cases, while fostering deep
links to traditional optimization methods. We theoretically analyze our method
and demonstrate on historical and synthetic data that LPGD converges faster
than gradient descent even in a differentiable setup.
