Differentiable Knapsack and Top-k Operators via Dynamic Programming¶

会议: NeurIPS 2026
arXiv: 2601.21775
代码: https://github.com/google-deepmind/differentiable_knapsack
领域: 优化/理论
关键词: 可微动态规划, 背包问题, Top-k 算子, Fenchel-Young 损失, 组合优化

一句话总结¶

将 0/1 背包与 Top-$k$ 子集选择统一建模为动态规划递推，通过对 Bellman 递归中的最大值算子进行正则化平滑，导出了支持置换等变性与稀疏选择的可微松弛算子、高效向量-雅可比积（VJP）及自回归采样算法。

研究背景与动机¶

在现代机器学习管线中，稀疏隐变量建模、资源受限分配以及决策聚焦学习（Decision-focused Learning）等核心任务均依赖于离散子集的选择。其中，标准的 Top-$k$ 算子在固定基数约束下选择项目，而经典的 0/1 背包问题则进一步推广了这一设定，允许项目具有异构权重并在总容量约束下最大化总体收益。然而，这类离散组合算子本质上是分段常数函数，其雅可比矩阵在几乎所有位置要么处处为零，要么未定义。这直接打破了神经网络的端到端可微计算图，导致反向传播梯度无法有效回传。

为了弥合组合求解器与端到端梯度优化之间的鸿沟，现有方法多采用黑盒扰动技术（如通过高斯噪声扰动并利用 Monte-Carlo 估计梯度）或分段仿射插值平滑。这些方法虽然解耦了底层求解器，但在单次更新中需要反复调用求解器数十次，引入了极其沉重的计算开销。另一类基于最优传输或正则化线性规划的专用可微排序与 Top-$k$ 层，则无法泛化至带有异构物品权重的非齐次背包容量约束。更为关键的是，现有大多数松弛方法通常将正则化形式局限在香农负熵上，使得输出概率处于单纯形内部而无法产生真正稀疏的硬选择，缺乏对离散归纳偏置的精细控制。

针对上述瓶颈，本文的核心切入点是重新审视背包问题与 Top-$k$ 之间内在的动态规划（Dynamic Programming, DP）同构性：两者均可在二维表格上展开为 Bellman 最优子结构递推。核心 idea：将背包与 Top-k 问题统一表示为有向无环图上的 Bellman 动态规划递推，通过对状态转移最大值算子进行广义熵平滑，推导出支持波前并行计算、显式反向 VJP、离散自回归采样及凸 Fenchel-Young 损失的端到端可微松弛算子，并严格证明了香农熵对置换等变性的唯一性与稀疏正则子的充要条件。

方法详解¶

整体框架¶

本文构建了一套将 0/1 背包与 Top-$k$ 算子统一表达为波前可并行计算的平滑动态规划架构。输入为上游神经网络预测的项目分值 $\boldsymbol{\theta} \in \mathbb{R}^n$、离散权重 $\boldsymbol{w} \in \mathbb{N}^n$ 以及整数容量限制 $C$（Top-$k$ 对应 $\boldsymbol{w}=\boldsymbol{1}$ 且 $C=k$）。在前向传播阶段，算法通过平滑 Bellman 递推生成值函数表 $\boldsymbol{V}^\boldsymbol{w}_\Omega$ 与转移梯度缓存 $\boldsymbol{Q}^\boldsymbol{w}_\Omega$；在确定性反向传播中，基于伴随矩阵递归回传梯度得到软选择向量 $\boldsymbol{y}^C_{\boldsymbol{w},\Omega}(\boldsymbol{\theta})$；在随机前向模式下，则利用状态缓存进行 $\mathcal{O}(n)$ 步的离散自回归逆向采样以消除训练-测试偏差。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入: 分值 θ, 权重 w, 容量 C"] --> B["Bellman 递推平滑化<br/>max 替换为 max_Ω 算子"]
    B --> C["波前并行前向递推<br/>计算值表 V 与转移缓存 Q"]
    C --> D["确定性算子推导<br/>伴随表反向回传输出软选择 y"]
    C --> E["自回归逆向采样<br/>单步 Bernoulli 状态转移输出离散硬掩码"]
    D --> F["反向传播 VJP 计算<br/>二阶方向导数回传获取梯度"]
    E --> F
    F --> G["Fenchel-Young 损失优化<br/>凸损失直接计算残差梯度 y - y_true"]

关键设计¶

1. 统一动态规划递推与 Bellman 算子平滑：消除求解不可微性

传统 0/1 背包的子问题最优值满足 Bellman 递归：跳过第 $i$ 项收益为 $V^{\boldsymbol{w}}[i-1,c]$，选取第 $i$ 项（当 $w_i \leq c$）收益为 $\theta_i + V^{\boldsymbol{w}}[i-1,c-w_i]$。对于 Top-$k$ 问题，可完全视作 $w_i=1, C=k$ 的特例，唯一区别在于将边界行 $V^{\boldsymbol{1}}[0, 1:]$ 初始化为 $-\infty$ 以强制满足严格等式约束。针对标准 $\max$ 算子的分段常数瓶颈，本文利用强凸正则化函数 $\Omega: \triangle^2 \to \mathbb{R}$ 将两分支比较平滑化： $$V^\boldsymbol{w}_\Omega[i,c] = \begin{cases} V^\boldsymbol{w}_\Omega[i-1,c], & \text{若 } w_i > c \\ \mathsf{max}_\Omega\left(\theta_i + V^\boldsymbol{w}_\Omega[i-1,c-w_i],\; V^\boldsymbol{w}_\Omega[i-1,c]\right), & \text{否则} \end{cases}$$ 其中平滑最大值算子定义为 $\mathsf{max}_\Omega(a,b) = \max_{\boldsymbol{q} \in \triangle^2} \{ q_1 a + q_2 b - \Omega(\boldsymbol{q}) \}$。由此定义的可微背包与 Top-$k$ 算子为最终状态值相对于输入分值的梯度 $\boldsymbol{y}^C_{\boldsymbol{w},\Omega}(\boldsymbol{\theta}) = \nabla_{\boldsymbol{\theta}} V^\boldsymbol{w}_\Omega[n,C]$，天然落入可行解集的凸包 $\mathsf{conv}(\mathcal{Y}^C_{\boldsymbol{w}})$ 内。

2. 正则化选择的理论表征：置换等变性与稀疏选择的充要条件

正则化函数 $\Omega$ 直接决定了松弛算子的代数对称性与几何边界特性。本文在理论上建立了两个关键结论：其一，对于任何可分凸正则子 $\Omega(\boldsymbol{q}) = \omega(q_1) + \omega(q_2)$，松弛算子具有置换等变性（即对输入分值与权重的置换满足 $\sigma \circ \boldsymbol{y}^C_{\boldsymbol{w},\Omega} = \boldsymbol{y}^C_{\sigma(\boldsymbol{w}),\Omega} \circ \sigma$）当且仅当 $\Omega$ 为缩放香农负熵 $-\gamma H^s(\boldsymbol{q}) = \gamma (q_1 \log q_1 + q_2 \log q_2)$。这是因为香农熵是唯一能够保证多元平滑最大值算子满足结合律（associativity）的正则子。其二，松弛算子能够触及单纯形边界输出真正稀疏选择（即可行顶点集 $\mathcal{Y}^C_{\boldsymbol{w}} \subseteq \mathsf{range}(\boldsymbol{y}^C_{\boldsymbol{w},\Omega})$）的充要条件是 $\omega'(t)$ 在端点 $(0, 1)$ 处的导数有界。此时定义跳过阈值 $\tau_\Omega = \omega'(1^-) - \omega'(0^+)$，当跳过局部优势 $\Delta_{i,c} = V^\boldsymbol{w}_\Omega[i-1,c] - (\theta_i + V^\boldsymbol{w}_\Omega[i-1,c-w_i]) \geq \tau_\Omega$ 时，输出分量精确置零 $(\boldsymbol{y}^C_{\boldsymbol{w},\Omega}(\boldsymbol{\theta}))_i = 0$。Gini 熵（$\tau_\Omega=\gamma$）和 1.5-Tsallis 熵（$\tau_\Omega=2\gamma$）均能产生稀疏解，而香农熵因梯度在边界发散只能产生全密集解。

3. 波前并行化与向量-雅可比积（VJP）高效回传：摆脱定制 CUDA 依赖

由于递推式中第 $i$ 行状态仅依赖于第 $i-1$ 行，算法在容量维度 $c \in \{1, \dots, C\}$ 上具备天然的波前并行（Wavefront Parallelism）结构，将传统 $\mathcal{O}(nC)$ 串行步骤降为 $\mathcal{O}(n)$ 次并行步。在反向传播中，算子的雅可比矩阵对应平滑值函数的 Hessian 矩阵 $\nabla^2 \mathsf{max}^C_{\boldsymbol{w},\Omega}(\boldsymbol{\theta})$，具有对称性。本文设计了 reverse-over-forward 算法：前向递推维护方向导数 $\dot{V}^\boldsymbol{w}_\Omega$，反向递推通过伴随变量矩阵 $\dot{\boldsymbol{E}}^\boldsymbol{w}_\Omega$ 回传敏感度，能够在无自定义底核的前提下，通过 Numba JIT 在 CPU/GPU 批处理与容量维度实现极高执行吞吐。

4. 自回归底层分布与 Fenchel-Young 损失：消除训练-测试偏差

为了解决连续松弛层在离散推理任务中的“训练时软向量、测试时硬掩码”不一致问题，本文证明了在平滑 DP 转移矩阵 $\boldsymbol{Q}^\boldsymbol{w}_\Omega[i,c] = \partial V^\boldsymbol{w}_\Omega[i,c] / \partial \theta_i$ 下，存在一个良定义的离散概率分布 $\pi^{\boldsymbol{w},C}_{\boldsymbol{\theta},\Omega}$，满足边缘期望精确等于松弛算子 $\mathbb{E}[Y] = \boldsymbol{y}^C_{\boldsymbol{w},\Omega}(\boldsymbol{\theta})$。前向传播可从 $i=n$ 倒序执行 Bernoulli 自回归单遍采样；反向传播则将该采样步视作确定性期望节点，利用 VJP 回传代理梯度。此外，作为监督输出层时，基于 Fenchel 共轭定义的 Fenchel-Young 损失 $L_{\Omega^C_{\boldsymbol{w}}}(\boldsymbol{\theta}; \boldsymbol{y}) = \mathsf{max}^C_{\boldsymbol{w},\Omega}(\boldsymbol{\theta}) + \Omega^C_{\boldsymbol{w}}(\boldsymbol{y}) - \langle \boldsymbol{\theta}, \boldsymbol{y} \rangle$ 保证全局凸性，其梯度形式极其优雅： $$\nabla_{\boldsymbol{\theta}} L_{\Omega^C_{\boldsymbol{w}}}(\boldsymbol{\theta}; \boldsymbol{y}) = \boldsymbol{y}^C_{\boldsymbol{w},\Omega}(\boldsymbol{\theta}) - \boldsymbol{y}$$ 直接由松弛预测与真实标签之差驱动，无需计算昂贵的二阶图导数。

实验关键数据¶

主实验¶

论文在 PyEPO 决策聚焦学习基准上评估背包预测优化性能，同时在容量与库存约束下的动态商品陈列（Dynamic Assortment, $n=20, T=80$）强化学习任务中对比基准算法。

任务 / 场景	评估指标	本文 (DP FY-Loss / SRL)	经典基准 (PFY / DBB / PPO)	相对优势 / 增益
PyEPO 决策学习 ($n=50$)	测试相对遗憾度 (Regret $\downarrow$)	0.024	0.052 (PFY) / 0.048 (DBB)	相对遗憾度降低 50.0%
PyEPO 决策学习 ($n=100$)	测试相对遗憾度 (Regret $\downarrow$)	0.031	0.071 (PFY) / 0.065 (DBB)	表现更加优异且计算时间随容量线性可控
动态陈列强化学习 ($T=80$)	期望累积收益 (Revenue $\uparrow$)	162.8 ± 3.1 (Tsallis)	134.5 (PPO) / 148.2 (标准 SRL)	超越经典强化学习，逼近作弊专家上限 (171.4)
动态陈列梯度方差	梯度协方差迹 ($\text{Tr}(\text{Cov}) \downarrow$)	0.18 (Gini) / 0.21 (Tsallis)	1.45 (PFY-SRL)	梯度估计方差降低一个数量级

消融实验¶

在离散变分自编码器（DVAE, Stacked MNIST 数字混合重建, $k=3, n=10$）中，针对不同正则化函数（香农负熵、Gini 负熵、1.5-Tsallis 负熵）以及前向模式（确定性 vs. 随机自回归采样）进行消融对比：

结构配置	正则化类型	重建 MSE ($\downarrow$)	隐变量稀疏度 (非零项数, 目标=3)	连续隐变量 KL 发散状态
确定性前向	香农熵 (Shannon)	0.0382	9.85 (完全密集, 无法归零)	稳定收敛
确定性前向	Gini 熵 (Gini)	0.0241	3.08 (快速收敛至真实基数)	稳定收敛
确定性前向	1.5-Tsallis 熵	0.0246	3.12 (高度稀疏)	稳定收敛
随机采样前向	Gini 熵 (Stochastic)	0.0255	3.00 (严格等于离散顶点)	稳定收敛
直通估计器基线 (ST)	Gumbel Top-$k$	0.0495	3.00 (硬采样)	灾难性发散 ($D_{\text{KL}} \to \infty$)

关键发现¶

正则化先验对隐变量选择至关重要：带有边界有界梯度的 Gini 和 Tsallis 正则化能够驱动非活跃项精确收敛到 0，使得隐变量非零项均值从香农熵的 9.85 骤降至 3.08，彻底消除了连续分布拖尾带来的重构模糊。
自回归采样与代理梯度的组合表现出极强的数值稳定性：相较于 Straight-Through (ST) 估计器导致的 KL 散度灾难性发散，本文平滑 DP 引导的期望梯度使连续属性与离散身份表征得以稳定解耦。

亮点与洞察¶

将组合算法结构重塑为图神经网络层：并非将离散组合问题粗暴包裹为黑盒采样器，而是深入动态规划的 Bellman 状态转移内部，通过对标量转移算子的局部平滑直接换取整图全局可微性，兼具数学优雅性与计算效率。
置换等变性与稀疏性的权衡定理：严谨指出了香农熵对于置换等变性的唯一性，阐明了在离散神经网络层中追求“完全排列无偏”与“严格稀疏边界”之间不可调和的内在数学张力。
一处计算、双向复用的波前流水线：单次调用前向平滑递归即可同时完成值函数求解、状态转移矩阵缓存，并直接支撑后续的反向 VJP 计算与离散逆向采样，极大降低了强化学习与端到端训练的系统开销。

局限与展望¶

算法的计算与空间复杂度依赖于背包容量离散网格 $\mathcal{O}(nC)$，当容量 $C$ 扩展为极大规模实数或连续高维约束时，需要借助状态离散化或缩放近似，可能产生离散化误差。
自回归逆向采样虽然只需单遍 $\mathcal{O}(n)$ 步，但由于时间维度的串行依赖，在项目数 $n$ 极大（如数万量级）时相比一键并行抽样存在延迟瓶颈，未来可探索自回归状态转移的并行推演与剪枝策略。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从广义动态规划与 Bellman 算子平滑视角统一了背包与 Top-$k$ 算子，理论表征深刻。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖运筹决策学习、组合强化学习与生成模型 DVAE 三大多元场景，主消融数据完备。
写作质量: ⭐⭐⭐⭐⭐ 算法推导逻辑严密，定理证明详实，伪代码与系统架构高度自洽。
价值: ⭐⭐⭐⭐⭐ 为离散组合约束集成进现代端到端可微深度学习提供了极具通用性与工程实用价值的基础算子。

任务 / 场景	评估指标	本文 (DP FY-Loss / SRL)	经典基准 (PFY / DBB / PPO)	相对优势 / 增益
PyEPO 决策学习 (\(n=50\))	测试相对遗憾度 (Regret \(\downarrow\))	0.024	0.052 (PFY) / 0.048 (DBB)	相对遗憾度降低 50.0%
PyEPO 决策学习 (\(n=100\))	测试相对遗憾度 (Regret \(\downarrow\))	0.031	0.071 (PFY) / 0.065 (DBB)	表现更加优异且计算时间随容量线性可控
动态陈列强化学习 (\(T=80\))	期望累积收益 (Revenue \(\uparrow\))	162.8 ± 3.1 (Tsallis)	134.5 (PPO) / 148.2 (标准 SRL)	超越经典强化学习，逼近作弊专家上限 (171.4)
动态陈列梯度方差	梯度协方差迹 (\(\text{Tr}(\text{Cov}) \downarrow\))	0.18 (Gini) / 0.21 (Tsallis)	1.45 (PFY-SRL)	梯度估计方差降低一个数量级