Differentiating Through Integer Linear Programs with Quadratic Regularization and Davis-Yin Splitting

Drawing upon recent advances in convex optimization (Ryu & Yin, 2022) and implicit neural networks (Fung et al., 2022; McKenzie et al., 2021), we propose a method designed specifically for “differentiating through” large-scale integer linear programs. Our approach is fast, easy to implement using our provided code, and, unlike several prior works (e.g. see Pogančić et al. (2019); Berthet et al. (2020)), trains completely on GPU. Numerical examples herein demonstrate our approach, run using only standard computing resources, easily scales to problems with tens of thousands of variables. Theoretically, we verify our approach computes an informative gradient via a refined analysis of Jacobian-free Backpropagation (JFB) (Fung et al., 2022). In summary, we do the following.

• $\vartriangleright$

  Building upon McKenzie et al. (2021), we use the three operator splitting technique of Davis & Yin (2017) to propose DYS-Net.

• $\vartriangleright$

  We provide theoretical guarantees for differentiating through the fixed point of a certain non-expansive, but not contractive, operator.

• $\vartriangleright$

  We numerically show DYS-Net easily handles combinatorial problems with tens of thousands of variables.

We focus on optimization problems of the form equation 1 where $f(x;d)=w(d)^{\top}x$ and $\mathcal{X}$ is the integer or binary points of a polytope, which we may assume to be expressed in standard form (Ziegler, 2012):

In other words, equation 1 is an Integer Linear Program (ILP). Similar to works by Elmachtoub & Grigas (2022); Wilder et al. (2019); Mandi & Guns (2020) and others we replace the model equation 2 with its continuous relaxation, and redefine

as a step towards making the computation of $\mbox{d}{x}_{\Theta}\big{/}\mbox{d}w_{\Theta}$ feasible.

We assume access to training data in the tuple form $(d,x^{\star}(d))$. Such data may be obtained by observing experienced agents solve equation 1, for example taxi drivers in a shortest path problem Piorkowski et al. (2009). We focus on the $\ell_{2}$ loss:

and $\mathcal{D}$ is the distribution of contextual data, although more sophisticated losses can be used with only superficial changes to our results. In principle we select the optimal weights by solving $\operatorname*{arg\,min}_{\Theta}\mathcal{L}(\Theta)$. In practice, the population risk is inaccessible, and so we minimize empirical risk instead (Vapnik, 1999):

For notational brevity, we temporarily suppress the dependence on $d$. The gradient of $\ell_{2}$ with respect to $\Theta$, evaluated on a single data point is

As discussed in Section 1, ${x}^{\star}_{\Theta}$ is piecewise constant as a function of $w_{\Theta}$, and this remains true for the LP relaxation equation 4. Consequently, for all $w_{\Theta}$, either $\partial{x}_{\Theta}/\partial\Theta=0$ or $\partial{x}_{\Theta}/\partial\Theta$ is undefined—neither case yields an informative gradient. To remedy this, Wilder et al. (2019); Mandi & Guns (2020) propose adding a small amount of regularization to the objective function in equation 4 to make it strongly convex. This ensures ${x}_{\Theta}$ is a continuously differentiable function of $w_{\Theta}$. We add a small quadratic regularizer, modulated by $\gamma\geq 0$, to obtain the objective

and henceforth replace equation 4 by

During training, we aim to solve equation 8 and efficiently compute the derivative $\partial{x}_{\Theta}/\partial\Theta$.

Finally, we offer a note of caution when using the quadratically regularized problem equation 8 as a proxy for the original integer linear program equation 1. Temporarily defining

we observe that two sources of error have arisen; one from replacing $\mathcal{X}$ with $\mathcal{C}$ (i.e. the discrepancy between $\hat{x}_{\Theta}$ and $\tilde{x}_{\Theta}$), and one from replacing the linear objective function with its quadratically regularized counterpart (i.e. the discrepancy between $\tilde{x}_{\Theta}$ and $x_{\Theta}$). When $A$ is totally unimodular, $\tilde{x}_{\Theta}$ is guaranteed to be integral, in which case $\hat{x}_{\Theta}=\tilde{x}_{\Theta}$. Moreover, Wilder et al. (2019, Theorem 2) shows that

whence, as in the proof Berthet et al. (2020, Proposition 2.2), $\lim_{\gamma\to 0^{+}}x_{\Theta}(d)=\tilde{x}_{\Theta}(d)$ (at least when $\mathcal{C}$ is compact). Thus, for totally unimodular ILPs (e.g. the shortest path problem), it is reasonable to expect that $x_{\Theta}(d)$ obtained from equation 8, for well-tuned $\Theta$ and small $\gamma$, are of acceptable quality. For ILPs which are not totally unimodular no such theoretical guarantees exist. However, we observe that this relaxation works reasonably well, at least for the knapsack problem.

Introduced in Amos & Kolter (2017), and studied further in Agrawal et al. (2019b; c); Bolte et al. (2021); Blondel et al. (2022), an optimization layer is a modular component of a deep learning framework where forward propagation consists of solving a parametrized optimization problem. Consequently, backpropagation entails differentiating the solution to this problem with respect to the parameters. Optimization layers are a promising technology as they are able to incorporate hard constraints into deep neural networks. Moreover, as their output is a (constrained) minimizer of the objective function, it is easier to interpret than the output of a generic layer.

When an optimization layer is fitted to a data-dependent optimization problems of the form equation 1, with the goal of maximizing the quality of the predicted solution, this is usually referred to as decision-focused learning. We narrow our focus to the case where the objective function in equation 1 is linear.Wilder et al. (2019) and Elmachtoub & Grigas (2022) are two of the earliest works applying deep learning techniques to data-driven LPs, and delineate two fundamentally different approaches. Specifically, Wilder et al. (2019) proposes replacing the LP with a continuous and strongly convex relaxation, as described in Section 2. This approach is extended to ILPs in Ferber et al. (2020), and to non-quadratic regularizers in Mandi & Guns (2020). On the other hand, Elmachtoub & Grigas (2022) avoid modifying the underlying optimization problem and instead propose a new loss function; dubbed the Smart Predict-then-Optimize (SPO) loss, to be used instead of the $\ell_{2}$ loss. We emphasize that the SPO loss requires access to the true cost vectors $w(d)$, a setting which we do not consider.
Several works define a continuously differentiable proxy for the solution to the unregularized LP equation 4, which we rewrite here as¹¹1For sake of notational simplicity, the dependence on $d$ is implicit.

In Berthet et al. (2020), a stochastic perturbation is considered:

which is somewhat analogous to Nesterov-Spokoiny smoothing (Nesterov & Spokoiny, 2017) in zeroth-order optimization. For some draws of $Z$ the additively perturbed cost vector $w+\varepsilon z$ may have negative entries, which may cause a solver (e.g. Dijkstra’s algorithm) applied to the associated LP to fail. To avoid this, Dalle et al. (2022) suggest using a multiplicative perturbation of $w$ instead. Pogančić et al. (2019) define a piecewise-affine interpolant to $\ell(x^{\star}(w))$, where $\ell$ is a suitable loss function.
We note a bifurcation in the literature; Wilder et al. (2019); Ferber et al. (2020); Elmachtoub & Grigas (2022) assume access to training data of the form $(d,w(d))$, whereas Pogančić et al. (2019); Berthet et al. (2020); Sahoo et al. (2022) use training data of the form $(d,x^{\star}(d))$. The former setting is frequently referred to as the predict-then-optimize problem. Our focus is on the latter setting. We refer the reader to Kotary et al. (2021); Mandi et al. (2023) for recent surveys of this area.

Bai et al. (2019); El Ghaoui et al. (2021) propose the notion of deep equilibrium model (DEQ), also known as an implicit neural network. A DEQ is a neural network for which forward propagation consists of (approximately) computing a fixed point of a parametrized operator. We note that equation 8 may be reformulated as a fixed point problem,

where $P_{\mathcal{C}}$ is the orthogonal projection²²2For a set $\mathcal{A}\subseteq\mathbb{R}^{n}$, $P_{\mathcal{A}}(x)\triangleq\operatorname*{arg\,min}_{z\in\mathcal{A}}\|z-x\|$. onto $\mathcal{C}$. Thus, DEQ techniques may be applied to constrained optimization layers (Chen et al., 2021; Blondel et al., 2022; McKenzie et al., 2021). However, the cost of computing $P_{\mathcal{C}}$ can be prohibitive, see the discussion in Section 4.

Our work is related to the learning-to-optimize (L2O) framework (Chen et al., 2022; Li & Malik, 2017; Chen et al., 2018; Liu et al., 2023), where an optimization algorithm is learned and its outputs are used to perform inferences. However, traditional L2O methods use a fixed number of layers (i.e. unrolled algorithms). Our approach attempts to differentiate through the solution $x^{\star}_{\Theta}$ and is therefore most closely aligned with works that employ implicit networks within the L2O framework (Amos & Kolter, 2017; Heaton & Fung, 2023; McKenzie et al., 2021; Gilton et al., 2021; Liu et al., 2022). We highlight that, unlike many L2O works, our goal is not to learn a faster optimization method for fully specified problems. Rather, we seek to solve partially specified problems given contextual data by combining learning and optimization techniques.

In all of the aforementioned works, the same fundamental problem is encountered: $\partial{x}_{\Theta}/\partial\Theta$ must be computed where $x_{\Theta}$ is the solution of a (constrained) optimization problem with objective function parametrized by $\Theta$. The most common approach to doing so, proposed in Amos & Kolter (2017) and used in Ruthotto et al. (2018); Wilder et al. (2019); Mandi & Guns (2020); Ferber et al. (2020), starts with the KKT conditions for constrained optimality:

where $\hat{\lambda}$ and $\hat{\nu}\geq 0$ are Lagrange multipliers associated to the optimal solution ${x}_{\Theta}$ (Bertsekas, 1997) and $D(\hat{\nu})$ is a matrix with $\hat{\nu}$ along its diagonal. Differentiating these equations with respect to $\Theta$ and rearranging yields

The matrix and right hand side vector in equation 12 are computable, thus enabling one to solve for $\frac{d{x}_{\Theta}}{d\Theta}$ (as well as $\frac{d\hat{\lambda}}{d\Theta}$ and $\frac{d\hat{\nu}}{d\Theta}$). We emphasize that both primal (i.e., $x_{\Theta}$) and dual (i.e., $\hat{\lambda}$ and $\hat{\nu}$) variables at optimality are required. This constrains the choice of optimization schemes to those that track primal and dual variables, and prohibits the use of fast, primal-only schemes. Following (Amos & Kolter, 2017), most works employing this approach use a primal-dual interior point method, which computes acceptable approximations to $x_{\Theta}$, $\hat{\lambda}$, and $\hat{\nu}$ at a cost of $\mathcal{O}\left(\max\{n,m\}^{3}\right)$, assuming $x\in\mathbb{R}^{n}$ and $A\in\mathbb{R}^{m\times n}$.
Using the stochastic proxy equation 10, Berthet et al. (2020) derive a formula for $dx^{\star}_{\varepsilon}/dw$ which is also an expectation, and hence can be efficiently approximated using Monte Carlo methods. Pogančić et al. (2019) show that the gradients of their interpolant are strikingly easy to compute, requiring just one additional solve of equation 9 with perturbed cost $w^{\prime}$. This approach is extended by Sahoo et al. (2022) which proposes to avoid computing $dx^{\star}_{\varepsilon}/dw$ entirely, replacing it with the negative identity matrix. This is similar in spirit to our use of Jacobian-free Backpropagation (see Section 4).

As we wish to compute ${x}_{\Theta}$ and $\partial{x}_{\Theta}/\partial\Theta$ for high dimensional settings (i.e. large $n$), we eschew second-order methods (e.g. Newton’s method) in favor of first-order methods such as projected gradient descent (PGD). With PGD, a sequence $\{x^{k}\}$ of estimates of $x_{\Theta}$ are computed so that $x_{\Theta}=\lim_{k\rightarrow\infty}x^{k}$ where

This approach works for simple sets $\mathcal{C}$ for which there exists an explicit form of $P_{\mathcal{C}}$, e.g. when $\mathcal{C}$ is the probability simplex (Duchi et al., 2008; Condat, 2016; Li et al., 2023). However, for general polytopes $\mathcal{C}$ no such form exists, thereby requiring a second iterative procedure to compute $P_{\mathcal{C}}(x^{k})$. This projection must be computed at every iteration of every forward pass for every sample of every epoch, and this dominates the computational cost, see McKenzie et al. (2021, Appendix 4).

To avoid this expense, we draw on recent advances in the convex optimization literature, particularly Davis-Yin splitting or DYS (Davis & Yin, 2017). DYS is a three operator splitting technique, and hence is able to decouple the polytope constraints into relatively simple constraints Davis & Yin (2017). This technique has been used elsewhere, e.g. Pedregosa & Gidel (2018); Yurtsever et al. (2021). Specifically, we adapt the architecture incorporating DYS proposed in McKenzie et al. (2021). To this end, we rewrite $\mathcal{C}$ as

While $P_{\mathcal{C}}$ is hard to compute, both $P_{\mathcal{C}_{1}}$ and $P_{\mathcal{C}_{2}}$ can be computed cheaply via explicit formulas, once a singular value decomposition (SVD) is computed for $A$. We verify this via the following lemma (included for completeness, as the two results are already known).

We remark that the SVD of $A$ is computed offline, and is a once-off cost. Alternatively, one could compute $P_{\mathcal{C}_{1}}$ by using an iterative method (e.g. GMRES) to find the (least squares) solution to $Av=Az-b$, whence $P_{\mathcal{C}_{1}}(z)=z-v$, but this would need to be done at each iteration of each forward pass for every sample of every epoch. Thus, it is unclear whether this procides computational savings over the once-off computation of the SVD. However, this could prove useful when $A$ is only provided implicitly, i.e. via access to a matrix-vector product oracle.

The following theorem allows one to approximate $x_{\Theta}$ using only $P_{\mathcal{C}_{1}}$ and $P_{\mathcal{C}_{2}}$, not $P_{\mathcal{C}}$.

A full proof is included in Appendix A; here we provide a proof sketch. Substituting the gradient

for $F_{\Theta}(\cdot,d)$ into the formula for $T_{\Theta}(\cdot)$ given in McKenzie et al. (2021, Theorem 3.2) and rearranging yields equation 16. Verifying the assumptions of McKenzie et al. (2021, Theorem 3.3) are satisfied is straightforward. For example, $\nabla_{z}f_{\Theta}(z;\gamma,d)$ is $1/\gamma$-cocoercive by the Baillon-Haddad theorem (Baillon & Haddad, 1977; Bauschke & Combettes, 2009).
The forward pass of DYS-net consists of iterating equation 15 until a suitable stopping criterion is met. For simplicity, in presenting Algorithm 1 we assume this to be a maximum number of iterations is reached. One may also consider a stopping criterion of $\|z_{k+1}-z_{k}\|\leq\mathrm{tol}$.

From Theorem 2, we deduce the following fixed point condition:

As discussed in McKenzie et al. (2021), instead of backpropagating through every iterate of the forward pass, we may derive a formula for $dx_{\Theta}/d\Theta$ by appealing to the implicit function theorem and differentiating both sides of equation 17:

We notice two immediate problems not addressed by McKenzie et al. (2021): (i) $T_{\Theta}$ is not everywhere differentiable with respect to $z$, as $P_{\mathcal{C}_{2}}$ is not; (ii) if $T_{\Theta}$ were a contraction (i.e. Lipschitz in $z$ with constant less than $1$), then $\mathcal{J}_{\Theta}$ would be invertible. However, this is not necessarily the case. Thus, it is not clear a priori that equation 18 can be solved for $\mbox{d}{{z}_{\Theta}}/\mbox{d}{\Theta}$. Our key result (Theorem 6 below) is to provide reasonable conditions under which $\mathcal{J}_{\Theta}(z_{\Theta})$ is invertible.

Assuming these issues can be resolved, one may compute the gradient of the loss using the chain rule:

This approach requires solving a linear system with $\mathcal{J}_{\Theta}$ which becomes particularly expensive when $n$ is large. Instead, we use Jacobian-free Backpropagation (JFB) (Fung et al., 2022), also known as one-step differentiation (Bolte et al., 2024), in which the Jacobian $\mathcal{J}_{\Theta}$ is replaced with the identity matrix. This leads to an approximation of the true gradient $\mbox{d}{\ell}/\mbox{d}{\Theta}$ using

This update can be seen as a zeroth order Neumann expansion of $\mathcal{J}_{\Theta}^{-1}$ Fung et al. (2022); Bolte et al. (2024). We show equation 19 is a valid descent direction by resolving the two problems highlighted above. We begin by rigorously deriving a formula for $\partial{T_{\Theta}}/\partial{z}$. Recall the following generalization of the Jacobian to non-smooth operators due to Clarke (1983).

The Clarke Jacobian of $P_{\mathcal{C}_{2}}$ is easily computable, see Lemma 9. Define the (set-valued) functions

Then

where $c$ is applied element-wise. We shall now provide a consistent rule for selecting, at any $\bar{z}$, an element of $\partial P_{\mathcal{C}_{2}}(\bar{z})$. If $z_{i}\neq 0$ for all $i$ then $\partial P_{\mathcal{C}_{2}}$ is a singleton. If one or more $z_{i}=0$ then $\partial P_{\mathcal{C}_{2}}$ is multi-valued, so we choose the element of $\partial P_{\mathcal{C}_{2}}$ with $0$ in the $(i,i)$ position for every $z_{i}=0$. We write $dP_{\mathcal{C}_{2}}/dz$ for this chosen element, and note that

This aligns with the default rule for assigning a sub-gradient to $\operatorname{ReLU}$ used in the popular machine learning libraries TensorFlow(Abadi et al., 2016), PyTorch (Paszke et al., 2019) and JAX (Bradbury et al., 2018), and has been observed to yield networks which are more stable to train than other choices (Bertoin et al., 2021).

Given the above convention, we can compute $\partial{T_{\Theta}}/\partial{z}$, interpreted as an element of the Clarke Jacobian of $T_{\Theta}$ chosen according to a consistent rule. Surprisingly, $\partial{T_{\Theta}}/\partial{z}$ may be expressed using only orthogonal projections to hyperplanes. Throughout, we let $e_{i}\in\mathbb{R}^{n}$ be the one-hot vector with $1$ in the $i$-th position and zeros elsewhere, and $a_{i}^{\top}$ be the $i$-th row of $A$. For any subspace $\mathcal{H}$ we denote the orthogonal subspace as $\mathcal{H}^{\perp}$. The following theorem, the proof of which is given in Appendix A, characterizes $\partial{T_{\Theta}}/\partial{z}$.

To show JFB is applicable, it suffices to verify $\|\partial{T_{\Theta}}/\partial{z}\|<1$ when evaluated at ${z}_{\Theta}$. Theorem 4 enables us to show this inequality holds when ${x}_{\Theta}$ satisfies a commonly-used regularity condition, which we formalize as follows.

The significance of Theorem 6, which is proved in Appendix A, is outlined by the following corollary, stating that using $p_{\Theta}$ instead of the true gradient $d\ell/d\Theta$ is still guaranteed to decrease $\ell(\Theta)$, at least for small enough step-size.

Theorem 6 also proves that $\mathcal{J}_{\Theta}(z_{\Theta})$ in equation 18 is invertible, thus also justifying the use of Jacobian-based backprop as well as numerous other gradient approximation techniques (Liao et al., 2018; Geng et al., 2021).

Finally, we note that in practice $p_{\Theta}(x^{K})$ is used instead of $p_{\Theta}(x_{\Theta})$. Fung et al. (2022, Corollary 0.1) guarantees that $p_{\Theta}(x^{K})$ will still be a descent direction for $K$ sufficiently large, see also Bolte et al. (2024, Cor. 1).

We highlight that the pseudogradient computed in the backward pass (i.e., $p_{\Theta}$) does not require any Lagrange multipliers, in contrast with the gradient computed using equation 12. This is why we may use Davis-Yin splitting, as opposed to a primal-dual interior point method, on the forward pass. This allows efficiency gains on the forward pass, as Davis-Yin splitting scales favorably with the problem dimension.

DYS-net is implemented as an abstract PyTorch (Paszke et al., 2019) layer in the provided code. To instantiate it, a user only need specify $A$ and $b$ (see equation 3) and choose an architecture for $w_{\Theta}$. At test time, it can be beneficial to solve the underlying combinatorial problem equation 2 exactly using a specialized algorithm, e.g. commercial integer programming software for the knapsack prediction problem. This can easily be done using DYS-net by instantiating the test-time-forward abstract method.