Differentiable Knapsack and Top-k Operators via Dynamic Programming
This work addresses the problem of enabling gradient-based optimization for discrete selection operators in machine learning, which is incremental as it builds on existing dynamic programming and smoothing techniques.
The paper tackled the challenge of integrating discrete knapsack and top-k operators into neural networks by proposing a differentiable relaxation framework using dynamic programming, achieving efficient parallel algorithms and demonstrating applications in decision-focused learning, constrained assortment RL, and discrete VAEs.
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.