一行装饰器把 Dijkstra 等组合求解器变成可微神经网络层

PMinervini · x · 2026-10-13

Pasin Manurangsi(EdinburghNLP)展示了一个极简 Python 装饰器,可以把任意组合求解器(如 Dijkstra 最短路径、IR 中的排序函数)一行代码变成神经网络的可微层。原理基于其 AAAI 2023 论文 Adaptive Perturbation-Based Gradient Estimation for Discrete Latent Variable Models,是 NeurIPS 2021 Implicit MLE 方法的扩展,用于对黑盒组合求解器做梯度反传,代码已开源为 torch-adaptive-imle(MIT 协议)。配图给出示例:先用 @aimle(targetdistribution=...) 装饰一个普通 batchedsolver 函数,即可获得可微版本,支持最短路径等离散优化嵌入端到端训练。