Lagrangian Proximal Gradient Descent for Learning Convex Optimization Models

Anselm Paulus 1 Vít Musil 2 Georg Martius 1

Abstract

relaxations (Wang et al., 2019; Wilder et al., 2019a; Mandi& Guns, 2020; Djolonga & Krause, 2017) and stochastic
smoothing (Berthet et al., 2020; Dalle et al., 2022), overproxy losses (Paulus et al., 2021), to finite difference-based
techniques (Vlastelica et al., 2020). See Appendix A for amore detailed overview of the related work.

We propose Lagrangian Proximal Gradient Descent (LPGD), a flexible framework for learning convex optimization models. Like traditional
proximal gradient methods, LPGD can be interpreted as optimizing a smoothed envelope of thepossibly non-differentiable loss. The smoothening allows training models that do not provide
informative gradients, such as discrete optimization models. We show that the LPGD update canefficiently be computed by rerunning the forward
solver on a perturbed input. Moreover, we provethat the LPGD update converges to the gradientas the smoothening parameter approaches zero.
Finally, we experimentally investigate the potential benefits of applying LPGD even in a fully
differentiable setting.

In this work, we unify and generalize previous finitedifference-based methods (McAllester et al., 2010; Domke,
2010; Vlastelica et al., 2020) in a general framework called
Lagrangian Proximal Gradient Descent (LPGD). It is motivated by traditional proximal optimization techniques(Moreau, 1962; Parikh & Boyd, 2014; Boyd & Vandenberghe, 2014), fostering deep links between traditional andcontemporary methods. Practically, in the differentiable setting, we show that gradients can be computed as the limit of
the LPGD update, which provides an alternative to previous
methods based on differentiating the optimality conditions(Amos & Kolter, 2017; Agrawal et al., 2019b). In the nondifferentiable setting, LPGD allows learning the parameters
even when GD fails, generalizing (Vlastelica et al., 2020) to
non-linear objectives and learnable constraints. Further, we
introduce regularization to stabilize training with LPGD. Finally, we empirically investigate to which extent LPGD can
improve upon GD even in the fully differentiable setting.

1. Introduction
Optimization at inference is inherent to many prediction
tasks, including autonomous driving (Paden et al., 2016),
modelling physical systems (Cranmer et al., 2020), or
robotic control (Kumar et al., 2016). Therefore, embedding optimization algorithms in machine learning models
is a powerful inductive bias. A recent trend has been toembed constrained convex optimization problems that canefficiently be solved to optimality (Amos & Kolter, 2017;
Agrawal et al., 2019a;b; Sun et al., 2022).

2. Background & Notation
2.1. Problem Setup
We consider a model that contains an embedded constrained
convex optimization procedure:

Training a convex optimization model is an instance of bilevel optimization (Gould et al., 2016), which is generally
challenging. When it is possible to propagate informative
gradients through the optimization problem, the task is typically approached with standard stochastic gradient descent(GD) (Amos & Kolter, 2017; Agrawal et al., 2019b). In
some problems, such as discrete optimization layers, gradients are not informative. Previous works have proposed several methods to overcome this, ranging from differentiable

W

z∗

ℓ

θ
e ∈ Rp −−→
w ∈ Rk −→ (x∗ , y ∗ ) ∈ Rn+m →− ℓ(x∗ ) ∈ R

Given an input e ∈ Rp , the model Wθ predicts parametersof the optimization problem w ∈ Rk . A solver then finds an
optimal primal-dual solution pair z ∗ = (x∗ , y ∗ ) ∈ Rn ×Rmof the optimization problem, and the primal solution x∗ ispassed to a loss function ℓ(x∗ ).1 The loss can further consistof additional layers that could be trained standardly. We focus on minimizing the loss with respect to the parameters θ.

1

Max Planck Institute for Intelligent Systems, Tübingen, Germany 2 Masaryk University, Brno, Czech Republic. Correspondence to: Anselm Paulus <anselm.paulus@tue.mpg.de>.

We consider constrained optimization problems of the formz ∗ (w) = (x∗ (w), y ∗ (w)) := arg min max L(x, y, w), (1)x∈X

Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

y1
A more general setup of a loss depending on both primal anddual solutions is discussed in Appendix C.3.

1

where X ⊆ Rn is convex and L is a continuously differentiable Lagrangian that is convex in the primal variables x
and affine in the dual variables y. The affinity assumption is
not restrictive as inequalities can be enforced by introducing
slack variables, see Appendix C.1 for details and examples
covering conic and quadratic programs. Next, we assume
strong duality

3.1. Lagrangian DivergenceFor x ∈ X , we have the inequality
sup L(x, y, w) ≥ inf sup L(x, y, w) = L∗ (w)yand therefore we define the Lagrangian divergence3 as
DL (x, x∗ |w) := sup L(w, x, y) − L∗ (w) ≥ 0

L∗ (w) := min max L(x, y, w) = max min L(x, y, w) (2)
x∈X

y

y

(7)

y

x∈X

(8)

y

x∈X

and, for simplicity, we assume L to be strictly convex in x,
so that optimization (1) always attains a unique solution.2

for x ∈ X and w ∈ Rk . Note that x∗ is determined by w,
therefore DL is a function of only x and w and we abbreviate it as DL (x|w). The divergence holds the key property

2.2. Proximal Gradient Descent & Moreau Envelope

DL (x|w) = 0 if and only if x = x∗ (w) for x ∈ X , (9)nThe Moreau envelope (Moreau, 1962) envτf : R → R of a
possibly non-smooth f : Rn → R is defined for τ > 0 as
1∥x − xb∥22 .
envτf (bx) = inf x f (x) + 2τwhich makes it a good measure of optimality of x given w.

(3)

3.2. Lower Lagrange-Moreau Envelope

The envelope envτf is a smoothed lower bound approximation of f (Rockafellar & Wets, 1998, Theorem 1.25). The
proximal map proxτf : Rn → Rn is given by1∥x − xb∥22
proxτf (bx) = arg inf x f (x) + 2τ=xb − τ ∇ envτf (bx).

Given τ > 0, we call the lower Lagrange-Moreau envelope(L-envelope) the function ℓτ : Rk → R defined asℓτ (w) := inf ℓ(x) + τ1 DL (x|w)
x∈X

(4)

= inf sup ℓ(x) + τ1 [L(x, y, w) − L∗ (w)].

Intuitively, the proximal map searches for an x close to xbwith a lower value of f , and the Moreau envelope is thesum of f and a Euclidean distance penalty at this x. The
proximal point method (Parikh & Boyd, 2014; Güler, 1992)
aims to minimize f by iteratively updating x
e 7→ proxτf (ex).

The lower L-envelope ℓτ is a smoothed lower bound approximation of the function w 7→ ℓ(x∗ (w)). The smoothnessof ℓτ is inherited from the smoothness of the Lagrangian.
A function zτℓ : Rk → Rn+m representing the corresponding lower L-proximal map is given by

Now, assume that f decomposes as f = g + h with gdifferentiable and h potentially non-smooth and consider
ge(x) := g(bx) + ⟨x − xb, ∇g(bx)⟩,

zτℓ (w) := arg inf sup ℓ(x) + τ1 [L(x, y, w) − L∗ (w)]yx∈X= arg inf sup L(x, y, w) + τℓ(x).

(5)

y

x∈X

a linearization of g around x
b. This yields(11)

The objectives in (10) and (11) are strictly convex in x;
hence the optimal points are uniquely attained.

1
proxτ (eg+h) (bx) = arg inf x ge(x) + h(x) + 2τ∥x − xb∥22= proxτ h (bx − τ ∇g(b
x))

(10)

y

x∈X

(6)

3.3. Upper Lagrange-Moreau Envelope

and iterating x
b 7→ proxτ h (bx − τ ∇g(bx)) is called proximal
gradient descent (Parikh & Boyd, 2014). For h = IX
this reduces to projected gradient descent (Parikh & Boyd,2014).

The upper L-envelope ℓτ : Rk → R is defined with maximization instead of minimization asℓτ (w) := sup ℓ(x) − τ1 DL (x|w)

(12)

x∈X

= − inf sup τ1 [L(x, y, w) − L∗ (w)] − ℓ(x).

3. Methodx∈X

Our goal is to translate the idea of proximal methods tohybrid models (2.1) by defining a smoothed envelope ofthe loss w 7→ ℓ(x∗ (w)) on which we can perform gradientdescent. Given w and optimal solution x∗ , this envelope
should select an x in the proximity of x∗ with a lower lossℓ. The key concept is to replace the Euclidean distance termwith a Lagrangian divergence indicating how close x is to
optimality given w.
2yand the upper L-proximal map z τℓ : Rk → Rn+m asz τℓ (w) := arg sup ℓ(x) − τ1 DL (x|w)
x∈X= arg inf sup L(x, y, w) − τℓ(x).
x∈X

(13)

y

3
In some cases, the Lagrangian divergence coincides with the
Bregman divergence, a proximity measure generalizing the Euclidean distance, opening connections to Mirror descent, see Appendix C.2.

For set-valued solution mapping see Appendix C.5.

2

Analogously to ℓτ , the upper L-envelope ℓτ is a smoothed
upper bound of w 7→ ℓ(x∗ (w)).4 We have the simple relationsℓτ = −(−ℓ)τ

and z τℓ = z−τℓlike in (5), which allows computing the L-proximal map
with the forward pass solver.6 W.l.o.g., exposing the linear
parameters of the Lagrangian as
L(x, y, w) := ⟨x, c⟩ + H(x, y, v)(14)with w = (c, v) and abbreviating ∇ℓ = ∇ℓ(x∗ ) we get

between the L-envelopes and L-proximal maps.

zτ ℓe = arg inf sup⟨x, c⟩ + H(x, y, v) + τ ⟨x, ∇ℓ⟩As both the lower and upper envelope have desirable properties, we will also work with the average L-envelope5ℓτ (w) := 21 [ℓτ (w) + ℓτ (w)].

x∈X

(15)

The upper L-proximal map is computed analogously. This
enables efficient gradient computation for the L-envelopeof the linearized loss using the forward solver7 as
1∇ℓ̃τ (w) = 2τ∇w L w, zτ ℓ̃ − ∇w L w, z τ ℓ̃ . (22)When using this gradient, we refer to Lagrangian Proximal
Gradient Descent (LPGD), or more specifically, to LPGDτ ,LPGDτ and LPGDτ for gradient descent on ℓ̃τ , ℓ̃τ and ℓ̃τ .

(16)

Theorem 3.1. It holds thatlim ∇ℓeτ (c, v) = ∇(c,v) ℓ(x∗ (c, v)) = lim ∇ℓeτ (c, v). (23)where we abbreviate zτℓ = zτℓ (w) and z τℓ = z τℓ (w).

τ →0Example (Direct Loss Minimization). For a loss ℓ(x, xtrue ),feature function Ψ and an optimization problem of the form

τ →0

The proof is given in Appendix C.4. The extension for
set-valued solution mappings is provided in Appendix C.6.

(17)

Example (Blackbox Backpropagation). For a linear program (LP)8x∈X

the gradient of the lower L-envelope update is∇w ℓτ (w) = τ1 [Ψ(x∗ , ξ) − Ψ(xτℓ , ξ)](21)= z (c + τ ∇ℓ, v).

We perform gradient descent on the L-envelope. By Danskin’s theorem (Danskin, 1966), the gradient of the envelopereadsx∗ (w, ξ) = arg min −⟨w, Ψ(x, ξ)⟩y∗3.4. Differentiating the Lagrange-Moreau Envelopes∇ℓτ (w) = τ1 [∇w L(w, zτℓ ) − ∇w L(w, z ∗ )],∇ℓτ (w) = τ1 [∇w L(w, z ∗ ) − ∇w L(w, z τℓ )],(20)x∗ (c) = arg min⟨x, c⟩,
(18)

(24)

x∈X

the LPGDτ update reduces towith x∗ = x∗ (w, ξ) and∇ℓeτ (c) = τ1 [x∗ (c + τ ∇ℓ) − x∗ (c)],xτℓ (w, ξ) = arg min −⟨w, Ψ(x, ξ)⟩ + τ ℓ(x, xtrue ). (19)
x∈X

(25)

the update rule in Blackbox Backpropagation (Vlastelica
et al., 2020). Moreover, their piecewise affine interpolation of the loss c 7→ ℓ̃(x∗ (c)) corresponds to the lower
L-envelope ℓ̃τ .

This recovers the “towards-better” Direct Loss Minimization(DLM) update (McAllester et al., 2010), while the “awayfrom-worse” update corresponds to the upper L-envelope
gradient. The DLM framework has also been generalized to
non-linear objective functions (Song et al., 2016; Lorberbom
et al., 2019), always considering τ → 0.

Example (Implicit Differentiation by Perturbation). For aregularized linear programx∗ (c) = arg min⟨x, c⟩ + H(x)

(26)

x∈X

3.5. Lagrangian Proximal Gradient Descentwith a strongly convex regularizer H, the LPGDτ update is

In practice, due to the loss term, the L-proximal map canbe difficult to evaluate. As in the proximal gradient descentalgorithm, we consider a linearization ℓ̃ of the loss ℓ at x∗

1[x∗ (c + τ ∇ℓ) − x∗ (c − τ ∇ℓ)],∇ℓeτ (c) = 2τ(27)

recovering the update in (Domke, 2010), where only thelimit case τ → 0 is considered.

4Note the link to the Proximal hull (Rockafellar & Wets, 1998,
Example 1.44) corresponding to taking an upper envelope of alower envelope. Note also that for a general loss, the optimization (12) and (13) can diverge. This will not be problematic for the
linearized loss we will consider in Sec. 3.5.
5Lower, upper and average envelopes are closely related to
the right-, left- and double-sided directional derivatives, see Appendix C.4.

6

Note that the loss linearization is only applied after the solver
and does not approximate/linearize the solution mapping.
7
Note that warm-starting the solver with z ∗ strongly accelerates
the computation of the L-proximal map.
8
LP does not satisfy the strict convexity of the Lagrangian, andthe solution map is set-valued in general, cf. Appendix C.5.

3

LPGD_avg Train
LPGD_avg Test

LPGD_upper Train
LPGD_upper Test

LPGD_lower Train
LPGD_lower Test

GD Train
GD Test

100

Loss

10 1
10 2
10 3
10 4
10 5

0

2

4

6

Epoch

8

10

0

2500 5000 7500 10000 12500

Backward Pass Time (s)Figure 2. Comparison of LPGD and gradient descent (GD) on the
Sudoku experiment. Reported are train and test MSE over epochs
and wall-clock time spent in the backward pass. Statistics are over
5 restarts. Additional results can be found in Appendix B.

We consider the Sudoku experiment proposed by Amos &Kolter (2017). The task is to learn the rules of Sudoku inthe form of linear programming constraints from pairs ofincomplete and solved Sudoku puzzles. See Appendix B
for detailed information on the experimental setup. We
compare LPGD to gradient descent (GD), in which we use
the CVXPY (Diamond & Boyd, 2016) implementation of(Agrawal et al., 2019b) to compute the gradients.

Figure 1. Visualization of the upper ℓ̃τ , average ℓ̃τ , and lowerℓ̃τ Lagrange-Moreau envelope for different temperatures τ andaugmentation strengths ρ. The envelopes are smoothed approximations of the linearized loss c 7→ ℓ̃τ (x∗ (c)), illustrated in black.
In Lagrangian Proximal Gradient Descent (LPGD) we optimize
the loss by gradient descent on the Lagrange-Moreau envelope.

The results are reported in Fig. 4.2. LPGDτ reaches a lower
final loss than GD, which shows that LPGDτ produces better
update steps than standard gradients. LPGDτ also outperforms LPGDτ and LPGDτ , which highlights that both lower
and upper envelope carry relevant information for the optimization. Considering backward pass computation time,
all variants of LPGD converge much faster than GD. Thisis due to the efficient computation of LPGD via warmstarting and parallelization. Additional results, including othermetrics and wallclock time convergence, are reported inAppendix B.

3.6. Augmented Lagrangian
The L-envelope inherits its smoothness from the Lagrangian.
This motivates augmenting the Lagrangian with a strongly
convex smoothing term as
Lρ (x, y, w) := L(x, y, w) + ρ2 ||x − x∗ (w)||b 22 .

(28)

The augmentation smoothens the Lagrangian envelope while
not changing the current optimal solution.9

5. Conclusion

4. Experiments

We propose Lagrangian Proximal Gradient Descent (LPGD),
a flexible framework for learning convex optimization models. It unifies and generalizes contemporary optimization
methods while providing deep links to traditional ones.

4.1. Visualizations
We visualize the different L-envelopes of the linearizede ∗ (c)) in Fig. 4.1, for a quadratic loss on theloss c 7→ ℓ(xsolution to the linear program (24) with X = [0, 1]n and a
one-dimensional random cut through the cost space.

LPGD approximates gradients as finite differences and only
requires accessing the forward solver as a black-box oracle.
Formulated as Gradient Descent (GD) on a loss functionsmoothening, LPGD allows learning general objective andconstraint parameters even in the non-differentiable setting.

4.2. Sudoku
The main practical use-case of LPGD is the case when gradients of the loss are uninformative. However, we focus our
experiments on fully differentiable problems, to investigate
whether going beyond standard gradient-based methods isbeneficial even in these cases.

Empirically, we explore the potential benefits of LPGD overGD even in a fully differentiable setting. We find that in
our synthetic setup, LPGD achieves faster convergence andbetter final results when compared to GD, which motivatesapplying LPGD in other fully differentiable experimental
setups. Finally, an exciting direction for future work is
to explore further the connections to Mirror descent (seeAppendix C.2) and experimentally compare it to LPGD.

9This also affects the adjoint derivative obtained from implicit
differentiation, in which it introduces a regularizing term in the
Jacobian. We discuss this in Appendix C.1.

4

References

Blondel, M.
Structured prediction with projection oracles. In Wallach, H. M., Larochelle, H., Beygelzimer,A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.),
Advances in Neural Information Processing Systems 32:
Annual Conference on Neural Information Processing
Systems 2019, NeurIPS 2019, December 8-14, 2019,
Vancouver, BC, Canada, pp. 12145–12156, 2019. URL
https://proceedings.neurips.cc/paper/2019/
hash/7990ec44fcf3d7a0e5a2add28362213cAbstract.html.

Agrawal, A., Verschueren, R., Diamond, S., and Boyd, S. A
rewriting system for convex optimization problems. Journal ofControl and Decision, 5(1):42–60, 2018.
Agrawal, A., Amos, B., Barratt, S. T., Boyd, S. P., Diamond,S., and Kolter, J. Z. Differentiable convex optimization
layers. In Wallach, H. M., Larochelle, H., Beygelzimer,A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.),
Advances in Neural Information Processing Systems 32:
Annual Conference on Neural Information Processing
Systems 2019, NeurIPS 2019, December 8-14, 2019,
Vancouver, BC, Canada, pp. 9558–9570, 2019a. URL
https://proceedings.neurips.cc/paper/2019/
hash/9ce3c52fc54362e22053399d3181c638Abstract.html.

Blondel, M., Martins, A. F. T., and Niculae, V. Learningwith fenchel-young losses. J. Mach. Learn. Res., 21:35:1–
35:69, 2020. URL http://jmlr.org/papers/v21/19021.html.
Blondel, M., Llinares-López, F., Dadashi, R., Hussenot, L.,
and Geist, M. Learning energy networks with generalized
fenchel-young losses. In NeurIPS, 2022. URL http://papers.nips.cc/paper_files/paper/2022/
hash/510cfd9945f8bde6f0cf9b27ff1f8a76Abstract-Conference.html.

Agrawal, A., Barratt, S., Boyd, S., Busseti, E., and Moursi,
W. M. Differentiating through a cone program. J. Appl. Numer. Optim, 1(2):107–115, 2019b. URL https://doi.org/
10.23952/jano.1.2019.2.02.

Boyd, S. P. and Vandenberghe, L. Convex Optimization. Cambridge University Press, 2014. ISBN 978-0-521-833783. doi: 10.1017/CBO9780511804441. URL https://web.stanford.edu/%7Eboyd/cvxbook/.

Amos, B. and Kolter, J. Z. Optnet: Differentiable optimization as a layer in neural networks. In Precup, D. and Teh,
Y. W. (eds.), Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine
Learning Research, pp. 136–145. PMLR, 2017. URL http://proceedings.mlr.press/v70/amos17a.html.

Chen, T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud,
D. Neural ordinary differential equations. In Bengio, S.,
Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi,N., and Garnett, R. (eds.), Advances in Neural Information
Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December
3-8, 2018, Montréal, Canada, pp. 6572–6583, 2018. URL
https://proceedings.neurips.cc/paper/2018/
hash/69386f6bb1dfed68692a24c8686939b9Abstract.html.

Amos, B., Koltun, V., and Kolter, J. Z. The limited multi-label
projection layer. CoRR, abs/1906.08707, 2019. URL http://arxiv.org/abs/1906.08707.
Bai, S., Kolter, J. Z., and Koltun, V. Deep equilibrium
models.
In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R.
(eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information
Processing Systems 2019, NeurIPS 2019, December 8-14,
2019, Vancouver, BC, Canada, pp. 688–699, 2019. URL
https://proceedings.neurips.cc/paper/2019/
hash/01386bd6d8e091c2ab4c7c7de644d37bAbstract.html.

Cranmer, M. D., Greydanus, S., Hoyer, S., Battaglia, P. W.,
Spergel, D. N., and Ho, S. Lagrangian neural networks. CoRR,
abs/2003.04630, 2020. URL https://arxiv.org/abs/
2003.04630.
Dalle, G., Baty, L., Bouvier, L., and Parmentier, A. Learning with combinatorial optimization layers: a probabilistic
approach. CoRR, abs/2207.13513, 2022. doi: 10.48550/
arXiv.2207.13513. URL https://doi.org/10.48550/
arXiv.2207.13513.

Bai, S., Koltun, V., and Kolter, J. Z.
Multiscale deep
equilibrium models.
In Larochelle, H., Ranzato, M.,
Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in
Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020,
NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL
https://proceedings.neurips.cc/paper/2020/
hash/3812f9a59b634c2a9c574610eaba5bedAbstract.html.

Danskin, J. M. The theory of max-min, with applications.
SIAM Journal on Applied Mathematics, 14(4):641–664, 1966.
ISSN 00361399. URL http://www.jstor.org/stable/
2946123.
Diamond, S. and Boyd, S. CVXPY: A Python-embedded modelinglanguage for convex optimization. Journal of Machine Learning
Research, 17(83):1–5, 2016.

Bauschke, H. H., Dao, M. N., and Lindstrom, S. B. Regularizing
with bregman-moreau envelopes. SIAM J. Optim., 28(4):3208–
3228, 2018. doi: 10.1137/17M1130745. URL https://doi.org/10.1137/17M1130745.

Djolonga, J. and Krause, A. Differentiable learning of submodular
models. In Advances in Neural Information Processing Systems,
pp. 1013–1023, 2017.

Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J., and Bach,
F. R. Learning with differentiable perturbed optimizers. CoRR,
abs/2002.08676, 2020. URL https://arxiv.org/abs/
2002.08676.

Domke, J. Implicit differentiation by perturbation. In Lafferty,
J. D., Williams, C. K. I., Shawe-Taylor, J., Zemel, R. S.,
and Culotta, A. (eds.), Advances in Neural Information
Processing Systems 23: 24th Annual Conference on Neural5

Güler, O. New proximal point algorithms for convex minimization.
SIAM J. Optim., 2(4):649–664, 1992. doi: 10.1137/0802032.
URL https://doi.org/10.1137/0802032.

Information Processing Systems 2010. Proceedings of ameeting held 6-9 December 2010, Vancouver, British Columbia,
Canada, pp. 523–531. Curran Associates, Inc., 2010. URL
https://proceedings.neurips.cc/paper/2010/
hash/6ecbdd6ec859d284dc13885a37ce8d81Abstract.html.

Jang, E., Gu, S., and Poole, B. Categorical reparameterization withgumbel-softmax. In 5th International Conference on Learning
Representations, ICLR 2017, Toulon, France, April 24-26, 2017,
Conference Track Proceedings. OpenReview.net, 2017. URL
https://openreview.net/forum?id=rkE3y85ee.

Dupont, E., Doucet, A., and Teh, Y. W. Augmented neuralodes. In Wallach, H. M., Larochelle, H., Beygelzimer,A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.),
Advances in Neural Information Processing Systems 32:
Annual Conference on Neural Information Processing
Systems 2019, NeurIPS 2019, December 8-14, 2019,
Vancouver, BC, Canada, pp. 3134–3144, 2019.
URL
https://proceedings.neurips.cc/paper/2019/
hash/21be9a4bd4f81549a9d1d241981cec3cAbstract.html.

Kumar, V., Todorov, E., and Levine, S. Optimal control withlearned local models: Application to dexterous manipulation.
In Kragic, D., Bicchi, A., and Luca, A. D. (eds.), 2016 IEEE
International Conference on Robotics and Automation, ICRA
2016, Stockholm, Sweden, May 16-21, 2016, pp. 378–383. IEEE,
2016. doi: 10.1109/ICRA.2016.7487156. URL https://doi.org/10.1109/ICRA.2016.7487156.

Elmachtoub, A. N. and Grigas, P. Smart “predict, then optimize”.
Manag. Sci., 68(1):9–26, 2022. doi: 10.1287/mnsc.2020.3922.
URL https://doi.org/10.1287/mnsc.2020.3922.

LeCun, Y. and Huang, F. J. Loss functions for discriminative training of energy-based models. In Cowell, R. G.
and Ghahramani, Z. (eds.), Proceedings of the Tenth International Workshop on Artificial Intelligence and Statistics, AISTATS 2005, Bridgetown, Barbados, January 68, 2005. Society for Artificial Intelligence and Statistics,
2005. URL http://www.gatsby.ucl.ac.uk/aistats/fullpapers/207.pdf.

Ferber, A. M., Wilder, B., Dilkina, B., and Tambe, M. Mipaal: Mixed integer program as a layer. In The Thirty-Fourth
AAAI Conference on Artificial Intelligence, AAAI 2020, The
Thirty-Second Innovative Applications of Artificial Intelligence
Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New
York, NY, USA, February 7-12, 2020, pp. 1504–1511. AAAI
Press, 2020. URL https://ojs.aaai.org/index.php/
AAAI/article/view/5509.

Lorberbom, G., Jaakkola, T. S., Gane, A., and Hazan, T.
Direct optimization through arg max for discrete variational auto-encoder. In Wallach, H. M., Larochelle, H.,
Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett,
R. (eds.), Advances in Neural Information Processing
Systems 32: Annual Conference on Neural Information
Processing Systems 2019, NeurIPS 2019, December 8-14,
2019, Vancouver, BC, Canada, pp. 6200–6211, 2019. URL
https://proceedings.neurips.cc/paper/2019/
hash/1a04f965818a8533f5613003c7db243dAbstract.html.

Fredrikson, M., Lu, K., Vijayakumar, S., Jha, S., Ganesh, V., andWang, Z. Learning modulo theories. CoRR, abs/2301.11435,
2023. doi: 10.48550/arXiv.2301.11435. URL https://doi.org/10.48550/arXiv.2301.11435.
Fung, S. W., Heaton, H., Li, Q., McKenzie, D., Osher, S. J., andYin, W. Fixed point networks: Implicit depth models withjacobian-free backprop. CoRR, abs/2103.12803, 2021. URL
https://arxiv.org/abs/2103.12803.

Mandi, J. and Guns, T. Interior point solving for lp-based
prediction+optimisation. In Larochelle, H., Ranzato, M.,
Hadsell, R., Balcan, M., and Lin, H. (eds.), Advancesin Neural Information Processing Systems, volume 33,
pp. 7272–7282. Curran Associates, Inc., 2020.
URL
https://proceedings.neurips.cc/paper/2020/
file/51311013e51adebc3c34d2cc591fefeePaper.pdf.

Geng, Z., Zhang, X., Bai, S., Wang, Y., and Lin, Z. On training
implicit models. CoRR, abs/2111.05177, 2021. URL https://arxiv.org/abs/2111.05177.
Ghaoui, L. E., Gu, F., Travacca, B., Askari, A., and Tsai, A. Y.
Implicit deep learning. SIAM J. Math. Data Sci., 3(3):930–958,
2021. doi: 10.1137/20M1358517. URL https://doi.org/
10.1137/20M1358517.

McAllester, D. A., Hazan, T., and Keshet, J. Direct loss
minimization for structured prediction. In Lafferty, J. D.,
Williams, C. K. I., Shawe-Taylor, J., Zemel, R. S., andCulotta, A. (eds.), Advances in Neural Information Processing
Systems 23: 24th Annual Conference on Neural Information Processing Systems 2010. Proceedings of a meeting
held 6-9 December 2010, Vancouver, British Columbia,
Canada, pp. 1594–1602. Curran Associates, Inc., 2010. URL
https://proceedings.neurips.cc/paper/2010/
hash/ca8155f4d27f205953f9d3d7974bdd70Abstract.html.

Gould, S., Fernando, B., Cherian, A., Anderson, P., Cruz, R. S., andGuo, E. On differentiating parameterized argmin and argmax
problems with application to bi-level optimization. arXiv
preprint arXiv:1607.05447, 2016.
Gu, F., Chang, H., Zhu, W., Sojoudi, S., and Ghaoui, L. E.
Implicit graph neural networks. In Larochelle, H., Ranzato,
M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advancesin Neural Information Processing Systems 33: Annual
Conference on Neural Information Processing Systems 2020,
NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL
https://proceedings.neurips.cc/paper/2020/
hash/8b5c8441a8ff8e151b191c53c1842a38Abstract.html.

Moreau, J. J. Fonctions convexes duales et points proximaux
dans un espace hilbertien. Comptes rendus hebdomadaires
des séances de l’Académie des sciences, 255:2897–2899, 1962.
URL https://hal.science/hal-01867195.

6

Niepert, M., Minervini, P., and Franceschi, L. Implicit MLE:
backpropagating through discrete exponential family distributions. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N.,
Liang, P., and Vaughan, J. W. (eds.), Advances in Neural
Information Processing Systems 34: Annual Conference on
Neural Information Processing Systems 2021, NeurIPS 2021,
December 6-14, 2021, virtual, pp. 14567–14579, 2021. URL
https://proceedings.neurips.cc/paper/2021/
hash/7a430339c10c642c4b2251756fd1b484Abstract.html.

Sun, H., Shi, Y., Wang, J., Tuan, H. D., Poor, H. V., and Tao,
D. Alternating differentiation for optimization layers. CoRR,
abs/2210.01802, 2022. doi: 10.48550/arXiv.2210.01802. URL
https://doi.org/10.48550/arXiv.2210.01802.
Vlastelica, M., Paulus, A., Musil, V., Martius, G., and Rolínek, M.
Differentiation of blackbox combinatorial solvers. In 8th International Conference on Learning Representations, ICLR 2020,
Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net,
2020. URL https://openreview.net/forum?id=
BkevoJSYPB.

O’Donoghue, B., Chu, E., Parikh, N., and Boyd, S. P. Conic
optimization via operator splitting and homogeneous selfdual embedding. J. Optim. Theory Appl., 169(3):1042–1068,
2016. doi: 10.1007/s10957-016-0892-3. URL https://doi.org/10.1007/s10957-016-0892-3.

Vlastelica, M., Rolínek, M., and Martius, G. Neuro-algorithmic
policies enable fast combinatorial generalization. In Meila,
M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 1824 July 2021, Virtual Event, volume 139 of Proceedingsof Machine Learning Research, pp. 10575–10585. PMLR,
2021. URL http://proceedings.mlr.press/v139/
vlastelica21a.html.

Paden, B., Cáp, M., Yong, S. Z., Yershov, D. S., and Frazzoli,
E. A survey of motion planning and control techniques for
self-driving urban vehicles. IEEE Trans. Intell. Veh., 1(1):33–
55, 2016. doi: 10.1109/TIV.2016.2578706. URL https://doi.org/10.1109/TIV.2016.2578706.

Wang, P., Donti, P. L., Wilder, B., and Kolter, J. Z. Satnet: Bridging deep learning and logical reasoning using a differentiablesatisfiability solver. In Chaudhuri, K. and Salakhutdinov, R.
(eds.), Proceedings of the 36th International Conference on
Machine Learning, ICML 2019, 9-15 June 2019, Long Beach,
California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 6545–6554. PMLR, 2019. URL http://proceedings.mlr.press/v97/wang19e.html.

Parikh, N. and Boyd, S. P. Proximal algorithms. Found. Trends
Optim., 1(3):127–239, 2014. doi: 10.1561/2400000003. URL
https://doi.org/10.1561/2400000003.
Paulus, A., Rolínek, M., Musil, V., Amos, B., and Martius,
G. Comboptnet: Fit the right np-hard problem by learning
integer programming constraints. In Meila, M. and Zhang,
T. (eds.), Proceedings of the 38th International Conference
on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 8443–8453. PMLR, 2021. URL http://
proceedings.mlr.press/v139/paulus21a.html.

Wilder, B., Dilkina, B., and Tambe, M. Melding the data-decisions
pipeline: Decision-focused learning for combinatorial optimization. In The Thirty-Third AAAI Conference on Artificial Intelligence, AAAI 2019, The Thirty-First Innovative Applications of Artificial Intelligence Conference, IAAI 2019, The
Ninth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2019, Honolulu, Hawaii, USA, January 27 - February 1, 2019, pp. 1658–1665. AAAI Press,
2019a. doi: 10.1609/aaai.v33i01.33011658. URL https://doi.org/10.1609/aaai.v33i01.33011658.

Paulus, M. B., Choi, D., Tarlow, D., Krause, A., and Maddison,
C. J. Gradient estimation with stochastic softmax tricks. In
Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H.
(eds.), Advances in Neural Information Processing Systems 33:
Annual Conference on Neural Information Processing Systems
2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL
https://proceedings.neurips.cc/paper/2020/
hash/3df80af53dce8435cf9ad6c3e7a403fdAbstract.html.

Wilder, B., Ewing, E., Dilkina, B., and Tambe, M. End to end
learning and optimization on graphs. In Wallach, H. M.,
Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B.,
and Garnett, R. (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information
Processing Systems 2019, NeurIPS 2019, December 8-14,
2019, Vancouver, BC, Canada, pp. 4674–4685, 2019b. URL
https://proceedings.neurips.cc/paper/2019/
hash/8bd39eae38511daad6152e84545e504dAbstract.html.

Rockafellar, R. T. and Wets, R. J. Variational Analysis, volume 317 of Grundlehren der mathematischen Wissenschaften.
Springer, 1998. ISBN 978-3-540-62772-2. doi: 10.1007/9783-642-02431-3. URL https://doi.org/10.1007/9783-642-02431-3.

Winston, E. and Kolter, J. Z. Monotone operator equilibrium networks. In Larochelle, H., Ranzato, M., Hadsell,R., Balcan, M., and Lin, H. (eds.), Advances in Neural
Information Processing Systems 33:
Annual Conference on Neural Information Processing Systems 2020,
NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL
https://proceedings.neurips.cc/paper/2020/
hash/798d1c2813cbdf8bcdb388db0e32d496Abstract.html.

Sahoo, S. S., Vlastelica, M., Paulus, A., Musil, V., Kuleshov,V., and Martius, G. Gradient backpropagation through combinatorial algorithms: Identity with projection works. CoRR,
abs/2205.15213, 2022. doi: 10.48550/arXiv.2205.15213. URL
https://doi.org/10.48550/arXiv.2205.15213.
Song, Y., Schwing, A. G., Zemel, R. S., and Urtasun, R. Training
deep neural networks via direct loss minimization. In Balcan, M.
and Weinberger, K. Q. (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York
City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 2169–2177. JMLR.org,
2016. URL http://proceedings.mlr.press/v48/
songb16.html.

Xu, M., Garg, S., Milford, M., and Gould, S. Deep declarative dynamic time warping for end-to-end learning of alignment paths. CoRR, abs/2303.10778, 2023. doi: 10.48550/
arXiv.2303.10778. URL https://doi.org/10.48550/
arXiv.2303.10778.

7

A. Related work
Numerous implicit layers have been proposed in recent years, including neural ODEs (Chen et al., 2018; Dupont et al.,2019) and root-solving-based layers (Bai et al., 2019; 2020; Gu et al., 2020; Winston & Kolter, 2020; Fung et al., 2021;
Ghaoui et al., 2021; Geng et al., 2021).
In this work, we focus on optimization-based layers. A lot of research has been done on obtaining the true informative
gradient of such a layer, either by use of the implicit function theorem to differentiate quadratic programs (Amos & Kolter,2017), conic programs (Agrawal et al., 2019b), ADMM (Sun et al., 2022), dynamic time warping (Xu et al., 2023), or by
finite differences (Domke, 2010; McAllester et al., 2010; Song et al., 2016; Lorberbom et al., 2019).
Another direction of research has investigated problems in which no informative gradients exist. The techniques for stilltraining the model range from continuous relaxations of SAT problems (Wang et al., 2019) and submodular optimization(Djolonga & Krause, 2017), over regularization of linear programs (Amos et al., 2019; Wilder et al., 2019a; Mandi &Guns, 2020; Paulus et al., 2020) to stochastic smoothing (Berthet et al., 2020; Dalle et al., 2022), learnable proxies (Wilder
et al., 2019b) and generalized straight-through-estimators (Jang et al., 2017; Sahoo et al., 2022). Other works have built on
geometric proxy losses (Paulus et al., 2021) and finite differences (Vlastelica et al., 2021; Niepert et al., 2021).
A special case of an optimization layer is to embed an optimization algorithm as the final component of the prediction
pipeline. This encompasses energy-based models (LeCun & Huang, 2005; Blondel et al., 2022), structured prediction(McAllester et al., 2010; Blondel, 2019; Blondel et al., 2020), smart predict-then-optimize (Ferber et al., 2020; Elmachtoub& Grigas, 2022) and symbolic methods such as SMT solvers (Fredrikson et al., 2023).

B. Experiments
We are given pairs of incomplete and solved sudokus. The task is to learn a model that solves incomplete sudokus. Instead oflimiting ourselves to the mini-sudoku case (4 × 4 grid) as in (Amos & Kolter, 2017), we consider the full 9 × 9 sudoku grid.
It is modelled as a one-hot-encoding; hence the partial input e and solved output x have dimension p = k = n = 9 × 9 × 9.
The dataset consists of N = 9000 training and 1000 test instances. The solution procedure is modelled as a genericbox-constrained linear programx∗ (A, b; e) = arg min⟨x, e⟩(29)x∈[0,1]nsubject to Ay + b = 0(30)

with a sufficient number of constraints to represent the rules of mini-sudoku. This differs from the original formulation in(Amos & Kolter, 2017), as we replace the quadratic regularizer with box constraints. Such a formulation is possible as our
method is not limited to quadratic programs.
We follow the training protocol described in (Amos & Kolter, 2017), using a modification of the public codebase for thepaper. The model is trained by minimizing the mean square error loss between the predictions and one-hot encodings of thecorrectly solved sudokus, for which an informative gradient with respect to the constraints exists. For evaluation, we follow(Amos & Kolter, 2017) in refining the prediction by taking an argmax over the one-hot dimension and report the percentageof violated ground-truth sudoku constraints. The hyperparameters α, τ and ρ were selected in a grid search.
We solve the optimization problems using CVXPY (Diamond & Boyd, 2016; Agrawal et al., 2018). The optimization
problem is automatically transformed to a conic program by CVXPY for which we implemented LPGD. We compare
LPGD to the standard implementation of the conic program differentiation (Agrawal et al., 2019b), with the augmentation
as described in Sec. C.1.
We chose the hyperparameters learning rate α, τ and ρ with a grid search. The best hyperparameters for LPGD are τ = 104 ,ρ = 10, α = 0.1, for gradient descent they are ρ = 103 , α = 0.1. We use these hyperparameters in our final evaluation.
The loss and error training curves are presented in Fig. 3 and Fig. 4. LPGD converges much faster than GD in terms oftime spent in the backward pass. This is also visible in the convergence in terms of total wallclock time, although to alesser extent, as the forward pass becomes more and more computationally demanding as training progresses and non-trivial
constraints are discovered. It is also apparent that LPGDτ achieves a lower final loss and error than GD, LPGDτ and LPGDτ .
This suggests that both lower and upper envelope carry relevant information for the optimization.
8

LPGD_upper Train
LPGD_upper Test

LPGD_lower Train
LPGD_lower Test

GD Train
GD Test

100

100

10 1

10 1

10 1

10 2

10 2

10 2

Loss

100

Loss

Loss

LPGD_avg Train
LPGD_avg Test

10 3

10 3

10 3

10 4

10 4

10 4

10 5

0

2

4

6

Epoch

8

10

10 5 02000 4000 6000 8000 1000012000

Backward Pass Time (s)10 5 010000 20000 30000 40000

Total Time (s)Figure 3. Results for the Sudoku experiment. Reported are the batched train and test loss over 11 epochs. We additionally report the loss
over total wallclock time, and total time spent in the backward passes.

LPGD_avg Train
LPGD_avg Test

LPGD_lower Train
LPGD_lower Test

GD Train
GD Test

100

100

10

1

10 1

10 1

10

2

10

3

10

4

10

5

Error (correct solution)

Error (correct solution)

010Error (correct solution)

LPGD_upper Train
LPGD_upper Test

10 2

10 2

10 3

10 3

10 4

10 4

10 5
0

2

4

6

Epoch

8

10

10 5
0

2000 4000 6000 8000 1000012000

Backward Pass Time (s)010000 20000 30000 40000

Total Time (s)Figure 4. Results for the Sudoku experiment. Reported are the batched train and test error over 11 epochs. We additionally report theloss over total wallclock time, and total time spent in the backward passes. The error of an instance is zero if it is a valid solution to thesudoku, otherwise it is one.

9

C. Theoretical Considerations
C.1. Implicit Differentiation with Augmentation
We inspect how the augmentation in (28) affects existing methods for computing the adjoint derivative of an optimization
problem.
Quadratic Program.

For a symmetric matrix H we can write a quadratic program with inequality constraints as(x∗ , s∗ ) = arg min 12 xT Hx + cT y

subject to Ax + b + s = 0.

(31)

x,s≥0

In Lagrangian form, we can write it asz ∗ = (x∗ , s∗ , y ∗ ) = arg min max L(H, c, A, b, x, y)(32)

L(H, c, A, b, x, y) = xT Hx + cT x + (Ax + b + s)T y.

(33)

x,s≥0

y

with the Lagrangian

The augmentation in (28) augments the Lagrangian toLρ (H, c, A, b, x, y) = xT Hx + cT x + (Ax + b + s)T y + ρ2 ∥x − x∗ ∥22(34)zρ∗ (H, c, A, b) = arg min max Lρ (H, c, A, b, x, y).

(35)

and we write
x,s≥0

y

As described in (Amos & Kolter, 2017), the optimization problem can be differentiated by treating it as an implicit layer via
the KKT optimality conditions, which are given as
Hx + AT y + c + ρ(x − x∗ ) = 0

diag(y)s = 0

Ax + b + s = 0

s ≥ 0.

(36)

Assuming strict complementary slackness renders the inequality redundant and the conditions reduce to the set of equations
Hx + AT y + c + ρ(x − x∗ )diag(y)s0 = F (x, s, y, H, c, A, b) = (37)
Ax + b + s
which admits the use of the implicit function theorem. It states that under the regularity condition that ∂F∂z is invertible,zρ∗ (w) can be expressed as an implicit function, and we can compute its Jacobian by linearizing the optimality conditions
around the current solution
0=∂F∂F ∂zρ∗
+
.
∂z ∂w
∂w

(38)

with

H + ρI
∂F
= 0
∂z
A

0
diag( ys )
I


AT
I .
0

(39)

It is now possible to compute the desired Vector-Jacobian-product as∇w ℓ(x∗ (w)) =∂F T ∂F −T∂z ∗ T∇ℓ(x∗ ) = −∇ℓ(x∗ ),
∂w
∂w ∂z

(40)

which involves solving a linear system. The augmentation term, therefore, serves as a regularizer for this linear system.
10

Conic Program.

A conic program (Boyd & Vandenberghe, 2014) is defined as(x∗ , s∗ ) = arg min cT y

subject to Ax + s + b = 0

(41)

x,s∈K

where K is a cone. The Lagrangian of this optimization problem is
Lρ (A, b, c, x, s, y) = cT x + (Ay + s + b)T y(42)

which allows an equivalent saddle point formulation given byz ∗ = (x∗ , s∗ , y ∗ ) = arg min max L(A, b, c, x, s, y).
x,s∈K

(43)

y

The KKT optimality conditions are
AT y + c = 0(s, y) ∈ K × K∗

Ax + s + b = 0

where K∗ is the dual cone of K. The skew-symmetric mapping0
Q = Q(A, b, c) =  −A
−cT

AT
0
−bT


c
b
0

sT y = 0

(44)

(45)

is used in the homogenous self-dual embedding (O’Donoghue et al., 2016), a feasibility problem that embeds the conic
optimization problem. Agrawal et al. (2019b) solve and differentiate the self-dual embedding. We use the CVXPY
implementation of this method as our baseline for computing the true adjoint derivatives of the optimization problem. The
augmentation in (28) changes the stationarity condition and, thereby, the skew-symmetric mapping as


ρI
AT c
0
b .
Qρ = Qρ (A, b, c) =  −A
(46)
−cT −bT 0
We adjust the CVXPY implementation accordingly for our experiments.
C.2. Relation to Mirror DescentStandard Mirror Descent. Classical mirror descent is an algorithm for minimizing a function ℓ(x) over a closed convex
set X ⊆ Rn . The algorithm is defined by the distance-generating function (or mirror map) ϕ : Rn → R, a strictly convex
continuously differentiable function. The mirror descent algorithm also requires the assumption that the dual space of ϕ isall of Rn , i.e. {∇ϕ(x) | x ∈ Rn } = Rn , and that the gradient of ϕ diverges as ∥x∥2 → ∞.
The Bregman divergence of the mirror map is defined as
Dϕ (x, xb) := ϕ(x) − ϕ(bx) − ⟨x − xb, ∇ϕ(b
x)⟩

(47)

b),
envϕτℓ (bx) := min ℓ(x) + τ1 Dϕ (x, x(48)

proxϕτℓ (bx) := arg min ℓ(x) + τ1 Dϕ (x, x
b)

(49)

The (left) Bregman-Moreau envelope (Bauschke et al., 2018) and the (left) Bregman-Moreau proximal map are defined asxx=xb − τ ∇ envϕτℓ (b
x).

(50)

xk+1 = arg min ℓ̃(x) + IX (x) + τ1 Dϕ (x, xk )(51)

Then, the mirror descent algorithm in proximal form is given by iteratively applying the Bregman-Moreau proximal map ofℓ̃ + IX asx= arg min⟨x, ∇ℓ(xk )⟩ + τ1 Dϕ (x, xk )
x∈X

11

(52)

An equivalent form of this algorithm isηk = ∇ϕ(xk )(53)

ηk+1 = ηk − τ ∇ℓ(xk )xek+1 = (∇ϕ)

−1

(54)

(ηk+1 )

(55)

xk+1 = arg min Dϕ (x, xek+1 )(56)

x∈X

which highlights that gradients are applied in the dual space of ϕ, and ∇ϕ serves as the mapping between primal and(“mirrored”) dual space.
Lagrangian Mirror Descent. In this section, we derive Lagrangian Mirror Descent (LMD), an alternative algorithm to
LPGD inspired by mirror descent. We define
Lw (x) := sup L(x, y, w)

(57)

y

and assume that Lw is strongly convex and continuously differentiable in x. As the key step, we define the distancegenerating function (mirror map) as ϕ = Lw . This has the interpretation that distances are measured in terms of the
Lagrangian, similar to the intuition behind the previously defined Lagrangian divergence.
This mirror map leads to the Bregman divergence
DLw (x, xb) = Lw (x) − Lw (bx) − ⟨x − xb, ∇Lw (bx)⟩
DLw (x, x∗ ) ≥ 0,
DLw (x, x∗ ) = 0 ⇔ x = x∗ (w)for x ∈ X .

(58)
(59)
(60)

We define the Lagrange-Bregman-Moreau envelope at w as the Bregman-Moreau envelope at x∗ = x∗ (w), i.e.
Lw∗ℓL
τ (w) := envτ ℓ+IX (x )(61)

= min ℓ(x) + IX (x) + τ1 DLw (x, x∗ )x(62)= min ℓ(x) + τ1 (L(x, w) − L∗ (w) − ⟨x − x∗ , ∇x L(x∗ , w)⟩),

(63)

x∈X

Lw
∗
xL
τ ℓ (w) := proxτ ℓ+IX (x )(64)

= arg min ℓ(x) + IX (x) + τ1 DLw (x, x∗ )x
= arg min ℓ(x) + τ1 (L(x, w) − L∗ (w) − ⟨x − x∗ , ∇x L(x∗ , w)⟩).
x∈X

(65)
(66)

Similar to how we defined LPGD as gradient descent on the Lagrange-Moreau envelope of the linearized loss, we define
Lagrangian Mirror Descent (LMD) as gradient descent on the Lagrange-Bregman-Moreau envelope of the linearized loss,
i.e.
L
∗
L
∗
∗
1
∇w ℓ Lτ (w) = τ ∇w [L(xτ ℓ , w) − L(x , w) − ⟨xτ ℓ − x , ∇x L(x , w)⟩].

(67)

Again, the approximation allows efficiently computing the gradients using the forward solver as
xL(u, v) = arg inf τ ⟨x, ∇ℓ⟩ + ⟨x, u⟩ + H(x, y, v) − ⟨x, ∇x L(x∗ , w)⟩
τ ℓe
x∈X

= x∗ (u + τ ∇ℓ − ∇x L(x∗ , w), v).

(68)
(69)

If X = Rn , then the original optimization problem (1) is unconstrained and we have the optimality condition ∇x L(x∗ , w) =0. Therefore, the Bregman divergence
DLw (x, x∗ ) = Lw (x) − Lw (x∗ ) − ⟨x − x∗ , ∇Lw (x∗ )⟩∗∗= L(x, w) − L(x , w) = L(x, w) − L (w) = DL (x|w)

(70)
(71)

is equal to our Lagrangian divergence. It follows that, in this case, the Lagrange-Moreau envelope/LPGD is equivalent to the
Lagrange-Bregman-Moreau envelope/LMD.
12

C.3. Extension to General Loss Functions
Loss on Dual Variables. In the main text we only considered losses depending only on the primal variables, i.e. ℓ(x). Fora loss on the dual variables ℓ(y) we can reduce the situation to the primal case, asy ∗ (w) = arg sup min L(x, y, w) = arg inf max −L(x, y, w).
x∈X

y

y

x∈X

(72)

This amounts to simply negating the Lagrangian in all equations while swapping x and y.
Loss on Primal and Dual Variables. The situation becomes more involved for a loss function L(x, y) depending on
both primal and dual variables. If it decomposes into a primal and dual component, i.e. L(x, y) = ℓp (x) + ℓd (y), we cancompute the envelopes of the individual losses independently. Note that a linearization of the loss L̃ trivially decomposes
this way. Adding the envelopes together yields a combined lower and upper envelope for the total loss as(ℓp )τ (w) + (ℓd )τ (w) = inf sup[ τ1 L(x, y, w) + ℓp (x)] − sup inf [ τ1 L(x, y, w) − ℓd (y)](73)(ℓd )τ (w) + (ℓp )τ (w) = sup inf [ τ1 L(x, y, w) + ℓd (y)] − inf sup[ τ1 L(x, y, w) − ℓp (x)].

(74)

y

x∈X
y

y

x∈X

x∈X

x∈X

y

Assuming strong duality of these optimization problems, this leads to(ℓp )τ (w) + (ℓd )τ (w) = inf sup[ τ1 L(x, y, w) + ℓp (x)] − inf sup[ τ1 L(x, y, w) − ℓd (y)](75)(ℓd )τ (w) + (ℓp )τ (w) = inf sup[ τ1 L(x, y, w) + ℓd (y)] − inf sup[ τ1 L(x, y, w) − ℓp (x)]

(76)

y

x∈X

x∈X

y

x∈X

x∈X

y

y

Strong duality holds in particular for a linearized loss L̃, as this only amounts to a linear perturbation of the original
optimization problem. Unfortunately, computing the average envelope(ℓp )τ (w) + (ℓd )τ (w) = 12 { inf sup[ τ1 L(x, y, w) + ℓp (x)] − inf sup[ τ1 L(x, y, w) − ℓd (y)](77)

+ inf sup[ τ1 L(x, y, w) + ℓd (y)] − inf sup[ τ1 L(x, y, w) − ℓp (x)]}

(78)

x∈X

x∈X

y

x∈X

y

x∈X

y

y

or its gradient would now require four evaluations of the solver, which can be expensive. To reduce the number of evaluations,
we instead “combine” the perturbations of the primal and dual loss, i.e. we defineLτ (w) := inf sup L(x, y) + τ1 [L(x, y, w) − L∗ (w)]

(79)

Lτ (w) := inf sup L(x, y) + τ1 [L(x, y, w) − L∗ (w)]

(80)

x∈X

y

x∈X y
τ
L (w) := 12 {Lτ (w) + Lτ (w)}
= inf sup[ τ1 L(x, y, w) + L(x, y)] − inf sup[ τ1 L(x, y, w) − L(x, y)].
x∈X y
x∈X y

(81)
(82)

Note that Lτ /Lτ (w) are not necessarily lower/upper bounds of the loss anymore. However, these combined envelopes also
apply to loss functions that do not separate into primal and dual variables, and computing their gradients requires fewer
additional solver evaluations. For L(x, y) = ℓ(x) we haveLτ (w) = ℓτ (w),Lτ (w) = ℓτ (w),Lτ (w) = ℓτ (w).

(83)

For the rest of the appendix, we will work with L instead of ℓ for full generality and reduce the situation to a primal loss ℓwith the relations above.
C.4. Recovery of True Gradient
We aim to show that for a linear loss approximation L̃, the LPGD update recovers the true gradient as τ approaches zero.
We again expose the linear parameters of the LagrangianL(z, u, v) = ⟨z, u⟩ + H(z, v)
13

(84)

and get from Danskin’s envelope theorem (Danskin, 1966)∇u L∗ (u, v) = ∇u L(z ∗ , u, v) = z ∗ (u, v).

(85)



(86)

We define
dw :=


∇z L
.
0

Then it holds thatlim ∇w L̃τ (w) = lim τ1 [∇w L∗ (w + τ dw) − ∇w L∗ (w)]
τ →0

T
T 
∂ 2 L∗∂ 2 L∗∂ 2 L∗∇z L
dw = 2= 2 dw = 2
0
∂ w
∂ w
∂ w

τ →0

T

=

(87)
(88)

T

∂(∇u L∗ )∂z ∗ T∂ 2 L∗∇z L =∇z L =∇z L = ∇w L(z ∗ (w))
∂w∂u
∂w
∂w

(89)

The main step in this derivation was to identify the Jacobian of the solution mapping as a sub-matrix of the Hessian of the
optimal Lagrangian function, which is a symmetric matrix.10 Exploiting the symmetry of the Hessian then allows computing
the gradient, which is a co-derivative (backward-mode), as the finite difference between two solver outputs, which usually
only computes a directional (forward-mode) derivative as
∆z =

∂z ∗
∆w = lim τ1 [z ∗ (w + τ ∆w) − z ∗ (w)]
τ →0∂w
= lim τ1 [∇u L∗ (w + τ ∆w) − ∇u L∗ (w)].
τ →0

(90)
(91)

We observe that computing the gradient of the L-envelope is the backward-mode counterpart of the forward-mode rightsided directional derivative. This observation also gives an interpretation of not taking the limit in theLPGD update: In
forward-mode, checking how the solver reacts to finite perturbation of the parameters intuitively provides higher-orderinformation than linear sensitivities to infinitesimal perturbations. The finite difference in the gradient of the L-envelope
serves the equivalent purpose in backward-mode, by back-propagating higher-order information instead of linear sensitivities
as in standard directional co-derivatives (back-propagation).
Note that for L(x, y) = ℓ(x) this reduces tolim ∇w ℓ̃τ (w) = ∇w ℓ(x∗ (w)).

τ →0

(92)

An analogous proof and discussion also hold for the upper envelope ℓ̃τ and average envelope ℓ̃τ , corresponding to the leftand double-sided directional derivatives.
C.5. Background: Set-valued Solution Mappings
In the main text, we assumed that the Lagrangian is strictly convex, which entails a unique solution to the optimization
problem. We now turn to the more general case, in which the optimization problem does not have a unique solution i.e. thesolution mapping is set-valued.
S : Rm ⇒ Rn : w 7→ S(w) = {z ∗ ∈ Rn | L(z ∗ , w) = L∗ (w)}(93)

We still assume access to a solver oracle that returns a single element of the solution set z ∗ ∈ S(w).
One possible approach to generalizing the concepts of derivatives and co-derivatives to set-valued mappings is the machineryof graphical differentiation described in detail in Rockafellar & Wets (1998).
As the name suggests, this approach relies on a graphical interpretation of the derivative and co-derivative mappings to
generalize them from single-valued functions to set-valued mappings.
10Note that a similar derivation already appeared in (Domke, 2010), but only for primal variables without considering the benefits offinite values of τ over the limit.

14

Before we state the definition of graphical derivatives, we first develop a description of derivatives of single-valued functionsin terms of graphs that can be generalized to set-valued functions. For a single-valued function z ∗ (w) the directional
derivative∗

Dz ∗ (w) : Rm → Rn : ∆w 7→ ∂z∂w (w)∆w(94)

is a function mapping perturbations of the input ∆w to corresponding changes in the output. In a graphical interpretation, it
gives a first-order local approximation of the function graphGph z ∗ = {(w, z) | z = z ∗ (w)}∗Gph Dz ∗ (w) = {(∆w, ∆z) | ∆z = ∂z∂w (w)∆w}.

(95)

(96)

Likewise, the graph of the directional co-derivative∗TD∗ z ∗ (w) : Rn → Rm : dz 7→ ∂z∂w (w)dz(97)

is given by∗T

Gph D∗ z ∗ (w) = {(dw, dz) | dw = ∂z∂w (w)dz}(98)

We observe that up to a sign flip, the vectors in the graph of the derivative and the co-derivative are orthogonal⟨(dw, −dz), (∆w, ∆z)⟩ = ⟨dw, ∆w⟩ − ⟨−dz, ∆z⟩∗∗= ⟨ ∂zb T dz, ∆w⟩ − ⟨dz, ∂zb∂w (w)∂w (w)∆w⟩
∂z ∗
∂z ∗
b
− ⟨dz, ∂w (w)∆w⟩b= ⟨dz, ∂w (w)∆w⟩

=0

(99)
(100)
(101)
(102)

We can, therefore, equivalently define the co-derivative as the function with the graphGph D∗ z ∗ (w) = {(dw, dz) | (dw, −dz) ⊥ Gph Dz ∗ (w)}.

(103)

These properties can now be generalized to graphical derivatives of set-valued mappings (Rockafellar & Wets, 1998,
Definition 8.33). A local approximation to the graph of the set-valued mapping gives the graph of the directional derivative.
In general, a local approximation of a set C at x is given by its tangent cone TS (x) (Rockafellar & Wets, 1998, Definition 6.1),
therefore the graphical derivative is given byDS(w | z) : Rm ⇒ Rn : ∆w 7→ {∆z | (∆w, ∆z) ∈ TGph S (w, z)}(104)

Gph DS(w | z ∗ ) = TGph S (w, z ∗ )(105)

The graphical co-derivative is obtained as all the vectors orthogonal to Gph DS(w | z) after a sign flip in the secondbS (x) (Rockafellar
component. The vectors orthogonal to a tangent cone TS (x) are described by the regular normal cone N& Wets, 1998, Definition 6.3), (Rockafellar & Wets, 1998, Proposition 6.5). Therefore, we can express the graphicalco-derivative as
bGph S (w, z ∗ )}.
D∗ S(w | z ∗ ) : Rn ⇒ Rm : dz 7→ {dw | (dw, −dz) ∈ N(106)

We can back-propagate sub-differentials using the co-derivative mapping chain rule∂w L(S(w | z ∗ )) = D∗ S(w | z ∗ )(∇L(z ∗ ))
15

(107)

C.6. Recovery of True Sub-differential
Our goal is to show that the sub-differential of the Lagrange-Moreau envelope of the linearized loss ∂(u,v) L̃τ (u, v | z ∗ )converges in a set-valued sense to the sub-differential of the loss ∂(u,v) L(S(u, v | z ∗ )). We will show that this only holds inthe form of an inclusion.
Define the gradient of the optimal Lagrangian with respect to the non-linear parameters at the selected optimal solution asg ∗ := ∇v L(z ∗ , u, v).

(108)(z ∗ , g ∗ ) = ∇(u,v) L(z ∗ , u, v) ∈ ∂(u,v) L∗ (u, v) =: S + (u, v)(109)

We can then compactly write

where we defined the extended solution mapping S + : Rk 7→ Rk , which contains the gradient with respect to v in addition
to the optimal solution set.
The lower Lagrange-Moreau envelope of the linearized loss at z ∗ is defined asL̃τ (u, v | z ∗ ) = inf sup⟨z, ∇L(z ∗ )⟩ + τ1 [L(z, u, v) − L(z ∗ , u, v)]

(110)

x∈X y
= τ1 [ inf sup⟨z, u + τ ∇L(z ∗ )⟩ + H(z, v) − L(z ∗ , u, v)]x∈X y= τ1 [L∗ (u + τ ∇L(z ∗ ), v) − L(z ∗ , u, v)]

(111)
(112)

Note that we make the dependence on z ∗ explicit, as z ∗ is not uniquely determined anymore by (u, v). With a slight abuseof set-valued notation, the sub-differential of the Lagrange-Moreau envelope can again be written as the finite difference
between two extended solver solutions∂(u,v) L̃τ (u, v | z ∗ ) = τ1 [∂(u,v) L∗ (u + τ ∇L(z ∗ ), v) − ∇(u,v) L(z ∗ , u, v)](113)= τ1 [S + (u + τ ∇L(z ∗ ), v) − (z ∗ , g ∗ )]−−−→ DS + (u, v | z ∗ , g ∗ )(∇L(z ∗ ), 0).
τ →0

(114)
(115)

In the final step, we used that the finite difference converges to the graphical derivative (Rockafellar & Wets, 1998, Eq.
8(14)).
Next, we relate the graphical derivative to the desired graphical co-derivative in (107). A promising statement in thisdirection is given in (Rockafellar & Wets, 1998, Theorem 13.57), which states that under some regularity conditions, thegraphical derivative of a sub-differential mapping is contained in its graphical co-derivative. We reduce our situation tothis case by phrasing the solution mapping as a sub-differential mapping. To see this, we inspect the graphs of the solution
mapping
Gph S + = {(u, v | z ∗ , g ∗ ) | 0 ∈ ∂z (uT z ∗ + H(z ∗ , v) + Iz (z ∗ )), g ∗ ∈ ∇v H(z ∗ , v) + ∂IRm (v)}∗∗= {(u, v | z , g ) | (−u, g) ∈ ∂(H + IC )(z, v)}

(116)
(117)

and the graph of the sub-differential
Gph ∂(H + IC ) = {(z ∗ , v | u, g ∗ ) | (u, g ∗ ) ∈ ∂(H + IC )(z ∗ , v)}(118)= {(z ∗ , v | −u, g ∗ ) | (−u, g ∗ ) ∈ ∂(H + IC )(z ∗ , v)}.

(119)

We observe that the graphs are equal up to permutation and sign flips of the coordinates. This allows us to write(du, dv) ∈ D∗ S + (u, v | z ∗ , g ∗ )(dz, 0) ⇔ (du, dv, −dz, 0) ∈ NGph S + (u, v | z ∗ , g ∗ )(120)⇔ (−dz, dv, −du, 0) ∈ NGph ∂(H+IC ) (z ∗ , v | −u, g ∗ )⇔ (−dz, dv) ∈ D∗ ∂(H + IC )(z ∗ , v | −u, g ∗ )(du, 0)

(121)
(122)

⇐ (−dz, dv) ∈ D∂(H + IC )(z ∗ , v | −u, g ∗ )(du, 0)

(123)

∗

∗

⇔ (du, 0, −dz, dv) ∈ TGph ∂(H+IC ) (z , v | −u, g )⇔ (dz, 0, du, dv) ∈ TGph S + (u, v | z ∗ , g ∗ )

(124)
(125)

⇔ (du, dv) ∈ DS + (u, v | z ∗ , g ∗ )(dz, 0)

(126)

16

The key step in this derivation (the implication) is precisely the statement of (Rockafellar & Wets, 1998, Theorem 13.57).
It requires H + IC to be prox-regular and subdifferentilly continuous at (z ∗ , u) for (−u, g ∗ ). This is the case for allproper, lower semi-continuous, convex functions (Rockafellar & Wets, 1998, Example 13.30), which also includes indicatorfunctions of convex sets. As C is convex and H is assumed to be convex in z and v and affine in y, these conditions are
fulfilled, and the statement holds.
Finally, this gives us the desired result∂(u,v) L̃τ (u, v | z ∗ ) −−−→ DS + (u, v | z ∗ , g ∗ )(∇L(z ∗ ), 0) ⊆ D∗ S + (u, v | z ∗ , g ∗ )(∇L(z ∗ ), 0)
τ →0

(127)

= D∗ S(u, v | z ∗ )(∇L(z ∗ ))(128)

= ∂(u,v) L(S(u, v | z ∗ )).

(129)

For a primal loss L(x, y) = ℓ(x) this reduces to∂(c,v) ℓ̃τ (c, v | x∗ ) −−−→ DS + (c, v | x∗ , g ∗ )(∇ℓ(x∗ ), 0) ⊆ D∗ S + (c, v | x∗ , g ∗ )(∇ℓ(x∗ ), 0)
τ →0

17

(130)

= D∗ S(c, v | x∗ )(∇ℓ(x∗ ))(131)= ∂(c,v) ℓ(S(c, v | x∗ )).

(132)