Differentiable optimization layers (POUNCE-backed)

Differentiable optimization layers (POUNCE-backed)#

discopt exposes optimization solves as differentiable JAX layers: you can jax.grad through an LP/QP/NLP solution with respect to its parameters, and use the result inside a larger jit/vmap/grad pipeline.

The gradient does not differentiate through the solver’s iterations. It uses the implicit function theorem at the KKT point — sensitivity analysis in the style of sIPOPT [Pirnay et al., 2012] and OptNet [Amos and Kolter, 2017] — so the forward solve can be any backend that returns a primal-dual point. Here the forward solve is POUNCE, a pure-Rust port of the filter interior-point method of Wächter and Biegler [2006], while JAX supplies the model derivatives the adjoint needs.

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np
import discopt.modeling as dm

A differentiable NLP layer#

make_nlp_layer(model, parameters) returns solve(p) -> (obj, x, lam), differentiable with respect to the parameter vector p. We solve \(\min_x (x_0-p)^2 + x_1^2\) s.t. \(x_0+x_1=1\), \(0\le x\le 2\), whose optimum is \(x_0=(1+p)/2,\ x_1=(1-p)/2\).

from discopt._relax.pounce_layer import make_nlp_layer

m = dm.Model("parametric_nlp")
p = m.parameter("p", value=0.3)
x0 = m.continuous("x0", lb=0.0, ub=2.0)
x1 = m.continuous("x1", lb=0.0, ub=2.0)
m.minimize((x0 - p) ** 2 + x1 ** 2)
m.subject_to(x0 + x1 == 1.0)

layer = make_nlp_layer(m, [p])
obj, x, lam = layer(jnp.array([0.3]))
print(f"obj = {float(obj):.4f},  x = {np.asarray(x)}")

obj = 0.2450,  x = [0.65 0.35]

Differentiate the solution with jax.grad / jax.jacobian. The sensitivities come from the KKT adjoint and match the analytic values (\(dx_0/dp=\tfrac12,\ dx_1/dp=-\tfrac12\)).

dobj_dp = jax.grad(lambda pp: layer(pp)[0])(jnp.array([0.3]))
dx_dp = jax.jacobian(lambda pp: layer(pp)[1])(jnp.array([0.3]))
print(f"d(obj)/dp = {float(dobj_dp[0]):.4f}")
print(f"dx/dp     = {np.asarray(dx_dp).ravel()}")

d(obj)/dp = -0.7000
dx/dp     = [ 0.5 -0.5]

The layer composes under jit and vmap:

batched = jax.vmap(lambda pp: layer(pp)[0])(jnp.array([[0.2], [0.4], [0.6]]))
print("objectives over p in {0.2, 0.4, 0.6}:", np.asarray(batched))

objectives over p in {0.2, 0.4, 0.6}: [0.32 0.18 0.08]

Differentiable LP/QP layers#

The same implicit-KKT idea gives differentiable LP and QP layers [Amos and Kolter, 2017]. They keep their closed-form KKT sensitivity rule but solve the forward on POUNCE’s interior-point (analytic-center) point, which keeps the sensitivity system nonsingular.

from discopt._relax.differentiable_lp import lp_solve_grad

# min c.x  s.t.  A x = b,  0 <= x <= 1
c = jnp.array([1.0, -2.0, 0.5])
A = jnp.array([[1.0, 1.0, 1.0]])
b = jnp.array([1.5])
x_l = jnp.zeros(3)
x_u = jnp.ones(3)
obj = lp_solve_grad(c, A, b, x_l, x_u)
dc = jax.grad(lp_solve_grad)(c, A, b, x_l, x_u)
print(f"LP optimal objective = {float(obj):.4f}")
print(f"d(obj)/dc = {np.asarray(dc)}")

LP optimal objective = -1.7500
d(obj)/dc = [-4.98917109e-09  1.00000001e+00  4.99999996e-01]

Summary#