Differentiable Knapsack and Top-k Operators via Dynamic Programming

TL;DR

This work introduces a unified differentiable dynamic programming framework to embed Knapsack and Top-$k$ operators into neural networks. By replacing the non-differentiable max in Bellman recursions with differentiable max$_\Omega$, it enables deterministic or stochastic forward passes, exact VJPs, and Fenchel-Young losses, while preserving the ability to handle non-uniform weights and cardinality constraints. A key theoretical contribution is showing Shannon entropy is the unique separable regularizer that yields permutation-equivariant relaxed operators, and that certain regularizers can induce sparsity in selections. Empirically, the DP-based operators improve decision-focused learning, constrained dynamic assortment RL, and Fenchel-Young discrete VAE tasks, offering efficient gradients and scalable performance with sparse or dense relaxations as needed.

Abstract

Knapsack and Top-k operators are useful for selecting discrete subsets of variables. However, their integration into neural networks is challenging as they are piecewise constant, yielding gradients that are zero almost everywhere. In this paper, we propose a unified framework casting these operators as dynamic programs, and derive differentiable relaxations by smoothing the underlying recursions. On the algorithmic side, we develop efficient parallel algorithms supporting both deterministic and stochastic forward passes, and vector-Jacobian products for the backward pass. On the theoretical side, we prove that Shannon entropy is the unique regularization choice yielding permutation-equivariant operators, and characterize regularizers inducing sparse selections. Finally, on the experimental side, we demonstrate our framework on a decision-focused learning benchmark, a constrained dynamic assortment RL problem, and an extension of discrete VAEs.

Figures (9)

• Figure 1: Illustration of our relaxed Top-$k$ and Knapsack operators. We plot the sum of the first two coordinates of the relaxed operator for ${\bm{\theta}}=(\theta_1,\theta_2,\frac{1}{2},1)^\top$. In the first row, we use $\bm{y}^k_{\bm{1},\Omega}$ with $k=2$ and ${\bm{w}}=\bm{1}$. In the second row, we use $\bm{y}^C_{\bm{w},\Omega}$ with $C=2$ and ${\bm{w}}=(2,0,1,1)^\top$, breaking the symmetry in $\theta_1$ and $\theta_2$. Using $\Omega\equiv0$ yields piecewise-constant item selections. Gini and $1.5$-Tsallis regularization yield a sparse a differentiable (a.e. for Gini) operator, while Shannon regularization yields a dense and differentiable one.
• Figure 2: Scaling and performance of $\bm{y}^C_{\bm{w},\Omega}$ as an output layer. Lower computational time and test relative regret are better.
• Figure 3: Constrained dynamic assortment results. We report mean values and $95\%$ CIs. Left: Expected revenue, estimated on $10^4$ test episodes. Our method outperforms realizable baselines and approaches the expert oracle, which exploits hidden information. Middle: Expected trace of gradient covariance, estimated via repeated exploration sampling, target aggregation, and gradient estimation on replay buffer batches. Right: Average wall-clock time for gradient estimation. Our approach achieves speedups over perturbation-based SRL by avoiding repeated solver calls and reusing intermediate outputs from \ref{['algo:value']} for sampling and differentiation via \ref{['algo:layer', 'algo:sample']}.
• Figure 4: DVAE training dynamics. Dashed lines indicate a stochastic forward pass. Top: Test reconstruction MSE. Left: Sparsity of the latent representation ${\bm{y}}$ (average number of non-zero entries). Right: Train KL divergence between continuous latent style variables ${\bm{z}}$ and unit Gaussian prior.
• Figure 5: Example DP tables for the smoothed Knapsack (left) and Top-$k$ (right) recursions.

Theorems & Definitions (15)

• definition 1: Smoothed DP values and operators.
• proposition 1: Characterization of equivariance
• proposition 2: Characterization of sparsity
• proposition 3: Underlying distribution
• proof
• proof
• lemma 1
• lemma 2
• proof
• lemma 3