Part VII · coda · the continuous cousin

Where discrete meets differentiable

Every decision in this course is an argmax. Four ways to get a useful gradient through one, and an honest test of whether learning helps inside an exact solver.

needs 30, 07, 15~1½ h slides + 3 h lablab: uv run co lab 31

The last unit. It joins the course's two halves in both directions: continuous tools applied to discrete decisions (relaxations, estimators, argmin layers), and learning applied inside a discrete algorithm (a branching rule for unit 07's branch and bound).

The curriculum's framing of the state of the art is worth keeping in view: learned methods are still behind LKH and CP-SAT on classical problems, and genuinely promising as guidance inside exact solvers rather than as replacements for them. The lab's own result is a modest version of that finding.

How to read this deck

Each of the first three sections has a slide with numbers small enough to follow by hand: Sinkhorn on a 2×2 assignment, Gumbel-max on three logits, and the argmin Jacobian of a one-variable QP. All were computed with the lab's reference solution. The implicit-function derivation is in the notes.

31 · overview

The route

1 · relax the decision

Sinkhorn as a soft assignment; differentiable sorting.

2 · sample the decision

Gumbel-softmax, straight-through, REINFORCE: bias against variance.

3 · differentiate the optimum

The implicit function theorem on KKT: OptNet and cvxpylayers.

4 · learn to optimise

Learned branching and cuts, neural combinatorial optimisation, and where each stands.

getting around

Words with a dotted underline are glossary terms: click one for its definition. G lists every term in the unit, N shows the notes. Optional extra: interactive figures, linked from the slides they support by a small ▸ chip; skip them freely.

section 1

Relax the decision

A permutation becomes a doubly stochastic matrix.

31 · relax

Sinkhorn: the Hungarian algorithm, softened

$$\begin{aligned}\min_{P \in \mathcal{B}}\ & \langle C, P \rangle + \varepsilon \textstyle\sum_{ij} P_{ij}(\log P_{ij} - 1)\end{aligned}$$

$\mathcal{B}$: doubly stochastic matrices (the Birkhoff polytope, whose vertices are the permutations; unit 06). The entropy term makes the minimiser, a transport plan, unique, interior and smooth in $C$:

$$P_{ij} = \exp\big((f_i + g_j - C_{ij}) / \varepsilon\big)$$

Sinkhorn alternates: fix $g$ and choose $f$ so rows sum to 1, then fix $f$ and choose $g$ so columns do. In log domain it's two logsumexp operations per iteration.

the gradient (envelope theorem)

$\nabla_C \big[\min_P \langle C, P\rangle - \varepsilon H(P)\big] = P_\varepsilon^*$

The hard assignment cost has "gradient" equal to the optimal permutation matrix, which jumps. Sinkhorn replaces the jump by a smooth plan. Lab step 2: jax.grad through the iterations equals $P$ to 1e−6.

related relaxations

Gumbel-Sinkhorn (Mena et al. 2018): noise plus Sinkhorn, for learning latent permutations. Differentiable sorting and ranks via optimal transport (Cuturi, Teboul & Vert 2019) or projections onto the permutahedron (Blondel et al. 2020).

31 · relax · example

Sinkhorn, on a 2×2 assignment

Costs $C = \begin{pmatrix}1 & 3\\ 2 & 5\end{pmatrix}$. The hard optimum is the off-diagonal: $3 + 2 = 5$, against $1 + 5 = 6$.

temperature 1, from zero potentials

after	plan $P$	row sums	column sums
fitting rows	$\begin{pmatrix}0.881 & 0.119\\ 0.953 & 0.047\end{pmatrix}$	1, 1	1.83, 0.17
fitting columns	$\begin{pmatrix}0.480 & 0.715\\ 0.520 & 0.285\end{pmatrix}$	1.20, 0.80	1, 1
iteration 2	$\begin{pmatrix}0.383 & 0.628\\ 0.617 & 0.372\end{pmatrix}$	1.01, 0.99	1, 1
converged	$\begin{pmatrix}0.378 & 0.622\\ 0.622 & 0.378\end{pmatrix}$	1, 1	1, 1

what the plan says

It leans to the cheap off-diagonal (0.622) without committing. At ε = 0.5 it leans harder: 0.731. As ε → 0 it becomes the permutation.

Greedy rounding takes the largest entry, 0.622 at row 1, column 2, then the rest: the optimal assignment, cost 5.

the gradient, read off

jax.grad of the entropic cost with respect to $C$ returns exactly this $P$. Raising $C_{11}$ by a little raises the smooth cost at rate 0.378; the hard cost's rate is 0 or 1 depending on which side of the switch you are.

explore interactively

A closed form for 2×2. A doubly stochastic 2×2 matrix is $\begin{pmatrix}a & 1-a\\ 1-a & a\end{pmatrix}$. Minimising $a(C_{11} + C_{22}) + (1-a)(C_{12} + C_{21})$ plus ε times the negative entropy gives $a / (1 - a) = e^{-\Delta/(2\varepsilon)}$ with $\Delta = (1 + 5) - (3 + 2) = 1$. At ε = 1 that is $e^{-1/2} = 0.607$, so $a = 0.378$; at ε = 0.5 it is $e^{-1}$, so $a = 0.269$. Both match the reference sinkhorn after 50 iterations.

The first half-step. With $g = 0$ each row is a softmax of $-C_{ij}/\varepsilon$: row 1 is $e^{-1}, e^{-3}$ normalised, $(0.881, 0.119)$. The columns then sum to 1.83 and 0.17, and the column step rescales them. The reference entropic_gradient after 200 iterations returns $\begin{pmatrix}0.3775 & 0.6225\\0.6225 & 0.3775\end{pmatrix}$, and round_to_permutation returns $[1, 0]$.

31 · relax

Measured: Sinkhorn is not an assignment solver

Lab then: costs uniform in [0, 1], mean cost per row, greedy rounding of the plan.

n	optimal	Hungarian	ε = 0.05	ε = 0.01	ε = 0.002
50	0.0238	0.002 s	0.0371	0.0320	0.0238 (2.2 s)
100	0.0163	0.013 s	0.0311	0.0261	0.0250 (4.1 s)
200	0.0089	0.062 s	0.0193	0.0217	0.0225 (11.6 s)

A random assignment costs 0.5 per row.

180× slower than your Hungarian at n = 200, and 2.5× the cost after rounding.
Colder didn't help at n = 200, probably because the plan is still spread over many entries and greedy rounding picks badly among them.
The cost differences between good assignments shrink like 1/n, so ε has to shrink with n, and iterations grow like 1/ε.

what it's for

A smooth, batched, GPU-friendly plan with a gradient, as a layer inside a model that must output a matching. There, "approximately a permutation, differentiably" is the requirement, and the Hungarian algorithm can't meet it.

section 2

Sample the decision

Gradients of expectations over discrete choices.

31 · sample

Gumbel, straight-through, REINFORCE

Gumbel-max and Gumbel-softmax

$k = \arg\max_i (\ell_i + G_i)$ with Gumbel noise $G_i$ samples exactly from $\mathrm{softmax}(\ell)$ of the logits. Replacing argmax by $\mathrm{softmax}((\ell + G)/\tau)$ gives a differentiable sample (Jang, Gupta & Poole 2017; Maddison, Mnih & Teh 2017).

straight-through (Bengio, Léonard & Courville 2013)

Forward the hard one-hot, backward the soft sample's gradient: hard + soft − stop_gradient(soft). Biased: the gradient belongs to a different function than the one evaluated.

REINFORCE (Williams 1992) $$\nabla_\ell\, \mathbb{E}[f(k)] = \mathbb{E}\big[f(k)\, (e_k - p)\big]$$

Unbiased, needs no derivative of $f$, and high variance.

Lab then, 10 categories, one sample per estimate, 2 000 estimates (exact gradient norm 0.496):

estimator	bias	spread	noise in the bias
REINFORCE	0.020	1.818	0.041
straight-through, τ = 0.1	0.059	2.529	0.057
straight-through, τ = 0.5	0.044	0.841	0.019
straight-through, τ = 1.0	0.137	0.406	0.009

Norms of vectors. REINFORCE's bias is within its noise. Straight-through trades variance for bias as τ grows: at τ = 1, 4.5× less spread and a bias 15 noise levels wide.

31 · sample · example

Gumbel-max, on three logits

Logits $\ell = (1, 0, -1)$, so $\mathrm{softmax}(\ell) = (0.665,\ 0.245,\ 0.090)$.

one sample, by hand

	$i = 1$	$i = 2$	$i = 3$
uniform $U_i$	0.086	0.237	0.801
$G_i = -\log(-\log U_i)$	−0.899	−0.365	1.507
$\ell_i + G_i$	0.101	−0.365	0.507

The argmax is $i = 3$: this time the noise lifted the least likely option to the top. It should, 9% of the time.

20 000 samples

	$i = 1$	$i = 2$	$i = 3$
frequency	0.671	0.242	0.087
softmax	0.665	0.245	0.090

Gumbel-softmax replaces the argmax by $\mathrm{softmax}((\ell + G)/\tau)$: with the noise above and $\tau = 1$, a soft vector leaning to $i = 3$ instead of a hard pick, and it has a gradient in $\ell$.

explore interactively

section 3

Differentiate the optimum

The KKT conditions of unit 29, differentiated.

31 · argmin

The implicit function theorem on KKT

An argmin layer: $x^*(q) = \arg\min_x \tfrac12 x^\top Q x + q^\top x$ s.t. $Ax \le b$. With the active set $S$ fixed (strict complementarity), the KKT conditions are linear:

$$\begin{bmatrix} Q & A_S^\top \\ A_S & 0 \end{bmatrix} \begin{bmatrix} x \\ \lambda_S \end{bmatrix} = \begin{bmatrix} -q \\ b_S \end{bmatrix}$$

The implicit function theorem: the same matrix times $[\partial x/\partial q;\ \partial \lambda/\partial q]$ equals $[-I;\ 0]$. One linear solve gives the argmin Jacobian, no matter how the solver found $x^*$.

derivation in the notes (N)

OptNet (Amos & Kolter 2017) for QPs; cvxpylayers (Agrawal et al. 2019) for disciplined parametrized convex programs.

Lab then, box-constrained QP, n = 30, 19 variables at a bound:

method	max error vs finite differences	time
implicit (yours)	6.8e−10	1.1 ms + solve
JAX through 500 unrolled projected-gradient steps	6.8e−10	4.9 ms
finite differences	—	60 solves, 0.34 s

where it breaks

At a degenerate point (a constraint active with zero multiplier) the active set can change under an infinitesimal change of $q$, and $x^*(q)$ isn't differentiable there. For LPs, $x^*$ is piecewise constant: the Jacobian is zero almost everywhere, the same problem as argmax.

Derivation

On a neighbourhood where the active set stays $S$, the solution satisfies $Qx + q + A_S^\top \lambda_S = 0$ and $A_S x = b_S$. Both are linear in $(x, \lambda_S, q)$. Differentiating with respect to $q$ gives $Q\,\partial x + I + A_S^\top \partial\lambda = 0$ and $A_S\,\partial x = 0$, which is the block system on the slide. Strict complementarity is what keeps $S$ fixed under small changes of $q$.

The LP remark is why differentiable optimisation for combinatorial problems adds smoothing: perturbed optimisers (Berthet et al. 2020) average the argmax over random perturbations of the costs, and regularised LPs add a quadratic or entropic term. Sinkhorn in section 1 is exactly the entropic version for assignment.

Unrolling (differentiating through the solver's iterations) and implicit differentiation agree here because 500 projected-gradient steps converge. With fewer iterations, unrolling gives the gradient of the truncated algorithm, which can be a different and sometimes more useful object.

31 · argmin · example

An argmin Jacobian, in one dimension

$x^*(q) = \mathrm{clip}(-q, 0, 1)$. Dots: the three cases in the table. Rings: the two kinks, where the active set changes and no derivative exists.

$\min_x \tfrac12 x^2 + qx$ subject to $0 \le x \le 1$. Unconstrained the minimiser is $-q$; the box clips it.

three cases

$q$	$x^*$	active set, multiplier	$dx^*/dq$
−0.5	0.5	none	−1
−2	1	$x \le 1$, $\lambda = 1$	0
0.3	0	$x \ge 0$, $\lambda = 0.3$	0

With nothing active, the KKT system is $1 \cdot x = -q$: slope $-1$. With a bound active, its equation pins $x$, and the solve returns 0. At $q = 0$ the bound $x \ge 0$ is active with multiplier 0: degenerate, and the two one-sided slopes disagree.

section 4

Learn to optimise

Guidance inside exact solvers, and the honest state of the field.

31 · learning

Learned branching, cuts, and neural combinatorial optimisation

inside exact solvers

Learning to branch: imitate strong branching's choices cheaply (Khalil et al. 2016; Alvarez, Louveaux & Wehenkel 2017). Graph neural networks on the variable–constraint graph (Gasse et al. 2019) beat SCIP's default reliability branching on some benchmark families.
Cut selection (Tang, Agrawal & Faenza 2020; Paulus et al. 2022), node selection, primal heuristics, and solver configuration.

The solver keeps its guarantees: a bad learned choice costs time, never correctness.

instead of solvers: neural CO

Pointer networks (Vinyals, Fortunato & Jaitly 2015), attention models trained with REINFORCE (Kool, van Hoof & Welling 2019), POMO (Kwon et al. 2020) for routing.

the honest current state

On classical problems at benchmark sizes, neural constructive methods remain behind LKH for TSP and behind CP-SAT, HGS and MIP solvers on their home ground, especially in generalisation to larger or differently distributed instances. As learned guidance inside exact or classical solvers, the results are genuinely promising.

Bengio, Lodi & Prouvost, EJOR 290 (2021), is the methodological survey.

31 · learning

Measured: the lab's learned rule, judged by unit 28's method

A linear ranker over 5 features, trained on 954 strong-branching decisions (10 knapsacks), tested on 20 held-out knapsacks with larger n:

rule	SGM nodes	SGM s	fewest nodes on
most-fractional	541	0.35	0
pseudocost	510	0.35	0
strong	386	0.84	20
learned	536	0.37	0

paired on nodes	ratio [95% CI]	Holm p
learned / most-fractional	0.99 [0.94, 1.03]	0.65
learned / pseudocost	1.05 [1.00, 1.10]	0.10
learned / strong	1.38 [1.29, 1.46]	0.0003

verdict: it does not beat most-fractional, and here's why

No evidence of a difference: paired interval [0.94, 1.03] over 20 held-out instances.
It learned fractionality: weight 2.27 against at most 0.62 on the other features, so it's nearly most-fractional.
The features can't carry the signal: held-out imitation accuracy 0.58, against 0.55 for most-fractional's own agreement with strong branching. Strong branching decides from child LP bounds, which none of the five features contains.

Next experiment, same protocol: add node-LP features (reduced costs, dual values, a one-pivot bound estimate), as Gasse et al. did.

31 · lab

Lab 31 — what you build

about 3 hours · uv run co lab 31

Sinkhorn in log domain, and greedy rounding.
Gradients through Sinkhorn in JAX, checked against the envelope theorem.
Gumbel-max, Gumbel-softmax, straight-through, and exact against REINFORCE gradients.
The argmin Jacobian of a QP by the implicit function theorem.
A learned branching rule: features, strong-branching data, a ranker.

uv run co test 31
uv run co then 31     # Sinkhorn vs Hungarian · estimators · argmin layer · learned branching, judged by unit 28

31 · checkpoint

Done when

checkpoint

All five lab steps are ✓.
Your imitation-learned branching rule beats most-fractional branching on node count on held-out instances, or you can say precisely why it did not, in terms of unit 28's methodology rather than anecdote.
For each of the four techniques, you can say what it buys and what it costs, with the numbers.

Reading

Cuturi, NeurIPS 2013.
Jang, Gupta & Poole, ICLR 2017.
Amos & Kolter, ICML 2017.
Gasse et al., NeurIPS 2019.
Bengio, Lodi & Prouvost, EJOR 290 (2021).

The capstone: one problem, four traditions, one harness.