Differentiable MILP / MIQP

Differentiable MILP / MIQP#

Optimization layers let a solver sit inside a differentiable program: you can backpropagate a downstream loss through argmin. For convex problems this is the OptNet / cvxpylayers construction [Agrawal et al., 2019, Amos and Kolter, 2017], which discopt already provides for LP, QP, and NLP via implicit differentiation of the KKT system (the sIPOPT sensitivity [Pirnay et al., 2012]).

Mixed-integer programs are harder: the optimal integer assignment is a piecewise-constant function of the parameters, so the objective is not globally smooth and the naive gradient is zero almost everywhere [Ferber et al., 2020]. This matters for decision-focused learning [Elmachtoub and Grigas, 2022], where we want to train an upstream predictor through a downstream discrete decision.

discopt handles integer models with fix-and-differentiate: solve the MILP/MIQP, fix the integers at their optimal values, and differentiate the resulting continuous restriction at the incumbent.

The fix-and-differentiate gradient#

For a parametric mixed-integer program

\[ \min_{x,\,y\in\mathbb{Z}^k}\; f(x,y;p)\quad\text{s.t.}\quad g(x,y;p)\le 0, \]

let \(y^\*(p)\) be the optimal integer assignment. Wherever \(y^\*\) is locally stable in \(p\), fixing \(y=y^\*\) leaves a continuous problem whose optimum \(x^\*(p)\) varies smoothly, and the envelope theorem gives the exact objective sensitivity

\[ \frac{d\,\text{obj}^\*}{dp} = \frac{\partial \mathcal{L}}{\partial p}\Big|_{x^\*,\,\lambda^\*} = \frac{\partial f}{\partial p} + \lambda^{\*\top}\frac{\partial g}{\partial p}, \]

evaluated at the fixed-integer optimum with its continuous multipliers \(\lambda^\*\). At parameter values where \(y^\*\) switches (breakpoints) the objective has a kink and the gradient is one-sided — the value returned is that of the incumbent assignment’s restriction, which is the right object when the discrete decision is being held fixed for learning.

import os
os.environ.setdefault('JAX_PLATFORMS', 'cpu')
os.environ.setdefault('JAX_ENABLE_X64', '1')

import numpy as np
import discopt.modeling as dm
from discopt._relax.differentiable import differentiable_solve

A fixed-charge MILP#

Minimize \(p\,y + x\) subject to \(x \ge 3 - y\), with \(y\in\{0,1\}\) and \(x\in[0,10]\). For a small fixed charge \(p<1\) it pays to ‘switch on’ \(y=1\), giving \(x=2\) and objective \(p+2\) — so \(d\,\text{obj}^\*/dp = 1\).

def fixed_charge(p_value):
    m = dm.Model('fixed_charge')
    p = m.parameter('p', value=p_value)
    y = m.binary('y')
    x = m.continuous('x', lb=0, ub=10)
    m.subject_to(x >= 3 - y)
    m.minimize(p * y + x)
    return m, p

m, p = fixed_charge(0.5)
res = differentiable_solve(m)
print('status   :', res.status)
print('objective:', round(float(res.objective), 6))
print('y*, x*   :', float(np.asarray(res.x['y'])), float(np.asarray(res.x['x'])))
print('d(obj)/dp:', round(float(res.gradient(p)), 6))

status   : optimal
objective: 2.5
y*, x*   : 1.0 2.0
d(obj)/dp: 1.0

The analytic gradient matches a central finite difference of the re-solved objective — confirming it is the true sensitivity, not a relaxation surrogate.

def fd_gradient(make_model, p0, eps=1e-5):
    mp, _ = make_model(p0 + eps)
    mm, _ = make_model(p0 - eps)
    return (differentiable_solve(mp).objective - differentiable_solve(mm).objective) / (2 * eps)

print('analytic:', round(float(res.gradient(p)), 6))
print('finite-diff:', round(float(fd_gradient(fixed_charge, 0.5)), 6))

analytic: 1.0
finite-diff: 1.0

A MIQP with a binding constraint#

Minimize \((x-p)^2 + \tfrac12 y\) subject to \(x \le 0.5\), \(y\in\{0,1\}\). For \(p>0.5\) the continuous optimum wants \(x=p\) but the constraint binds at \(x=0.5\) with \(y=0\), so the gradient flows through the active constraint: \(d\,\text{obj}^\*/dp = 2(p-0.5)\).

def miqp_binding(p_value):
    m = dm.Model('miqp_binding')
    p = m.parameter('p', value=p_value)
    y = m.binary('y')
    x = m.continuous('x', lb=-5, ub=5)
    m.subject_to(x <= 0.5)
    m.minimize((x - p) ** 2 + 0.5 * y)
    return m, p

m2, p2 = miqp_binding(2.0)
r2 = differentiable_solve(m2)
print('objective:', round(float(r2.objective), 6), ' expected', (0.5 - 2.0) ** 2)
print('analytic :', round(float(r2.gradient(p2)), 6), ' expected 2*(2-0.5) =', 2 * (2.0 - 0.5))
print('finite-diff:', round(float(fd_gradient(miqp_binding, 2.0)), 6))

OA: generating OA cuts for 0 of 1 rows (1 convex; the rest are either nonconvex or already exact in the master)

OA: generating OA cuts for 0 of 1 rows (1 convex; the rest are either nonconvex or already exact in the master)

OA: generating OA cuts for 0 of 1 rows (1 convex; the rest are either nonconvex or already exact in the master)

objective: 2.25  expected 2.25
analytic : 3.0  expected 2*(2-0.5) = 3.0

finite-diff: 3.0

A JAX-composable layer for decision-focused learning#

make_milp_objective_layer wraps the MILP solve + fix-and-differentiate gradient as a jax.custom_vjp, so it slots into jax.grad / jax.jit pipelines — the composition decision-focused learning needs [Elmachtoub and Grigas, 2022]. Here we backpropagate a downstream squared loss on the optimal objective through the integer solve.

import jax
import jax.numpy as jnp
from discopt._relax.pounce_layer import make_milp_objective_layer

m, p = fixed_charge(0.5)
layer = make_milp_objective_layer(m, [p])

# Downstream loss: drive the optimal objective toward a target of 2.0.
def loss(p_vec, target=2.0):
    return (layer(p_vec) - target) ** 2

p_vec = jnp.array([0.5])
print('obj       :', round(float(layer(p_vec)), 6))
print('d loss/dp :', round(float(np.asarray(jax.grad(loss)(p_vec))[0]), 6))

obj       : 2.5

d loss/dp : 1.0

The gradient 2 * (obj - target) * d(obj)/dp = 2 * (2.5 - 2.0) * 1 = 1.0 is computed end to end through the branch-and-bound solve.

Summary#

  • discopt differentiates MILP / MIQP / integer MINLP by fixing the integers at the optimum and applying the continuous envelope-theorem sensitivity [Pirnay et al., 2012] to the restriction.

  • The objective gradient is exact wherever the optimal integer assignment is locally stable; at breakpoints it is the incumbent assignment’s one-sided value — the right convention for decision-focused learning [Elmachtoub and Grigas, 2022, Ferber et al., 2020].

  • make_milp_objective_layer exposes this as a jax.custom_vjp for use inside jax.grad / jax.jit / jax.vmap, mirroring the LP/QP/NLP layers [Agrawal et al., 2019, Amos and Kolter, 2017].