[论文解读] Differentiable Top-k Operator with Optimal Transport

Yujia Xie, Hanjun Dai+6 more|arXiv (Cornell University)|Feb 16, 2020

Domain Adaptation and Few-Shot Learning参考文献 37被引用 15

一句话总结

本文提出 SOFT（可扩展的基于最优传输的可微分），一种利用熵正则最优传输（EOT）对 top-k 操作进行可微分近似的新型方法，从而在机器学习模型中实现端到端训练。通过将 top-k 选择建模为正则化最优传输问题，并结合 Sinkhorn 算法与 KKT 条件，SOFT 实现了对 top-k 操作的梯度传播，显著提升了 k-NN 分类、束搜索和稀疏注意力机制的性能。

ABSTRACT

The top-k operation, i.e., finding the k largest or smallest elements from a collection of scores, is an important model component, which is widely used in information retrieval, machine learning, and data mining. However, if the top-k operation is implemented in an algorithmic way, e.g., using bubble algorithm, the resulting model cannot be trained in an end-to-end way using prevalent gradient descent algorithms. This is because these implementations typically involve swapping indices, whose gradient cannot be computed. Moreover, the corresponding mapping from the input scores to the indicator vector of whether this element belongs to the top-k set is essentially discontinuous. To address the issue, we propose a smoothed approximation, namely the SOFT (Scalable Optimal transport-based diFferenTiable) top-k operator. Specifically, our SOFT top-k operator approximates the output of the top-k operation as the solution of an Entropic Optimal Transport (EOT) problem. The gradient of the SOFT operator can then be efficiently approximated based on the optimality conditions of EOT problem. We apply the proposed operator to the k-nearest neighbors and beam search algorithms, and demonstrate improved performance.

研究动机与目标

• 解决标准 top-k 操作的不可微性问题，该问题导致神经网络无法实现端到端训练。
• 开发一种平滑、连续的 top-k 操作近似方法，支持基于梯度的优化。
• 通过统一的可微分 top-k 操作，实现可微分的 k-NN 分类、束搜索和稀疏注意力机制。
• 通过熵正则化与 Sinkhorn 算法，在控制近似误差的同时保持计算效率。

提出的方法

• 将 top-k 操作重新表述为最优传输（OT）问题的解，将 top-k 元素的选择参数化为一个传输计划。
• 对 OT 问题施加熵正则化，以生成可微分的近似，从而得到熵正则最优传输（EOT）公式。
• 利用 Sinkhorn 算法高效计算最优传输计划，实现可扩展且可微分的 top-k 选择。
• 基于 EOT 问题的 Karush-Kuhn-Tucker（KKT）条件，推导 SOFT top-k 操作的梯度，从而支持反向传播通过 top-k 层。
• 将 SOFT 操作集成到下游模型中，如 k-NN 分类器、束搜索和稀疏注意力层。
• 以 O(n) 的计算复杂度实现该方法，其中 n 为输入元素的数量，确保可扩展性。

实验结果

研究问题

• RQ1 top-k 操作能否在保持其选择语义的前提下实现可微分？
• RQ2 熵正则最优传输能否提供一种平滑、连续的近似，对离散的 top-k 操作进行误差可控的逼近？
• RQ3 所提出的可微分 top-k 操作是否能提升 k-NN 分类、束搜索和注意力机制中的端到端训练性能？
• RQ4 SOFT 操作的梯度在反向传播过程中表现如何？是否与期望的优化目标一致？

主要发现

• SOFT top-k 操作使得原本因 top-k 不可微而必须采用两阶段训练的模型，实现了端到端训练。
• 在 k-NN 分类中，SOFT top-k 操作在准确率上优于基线方法，包括基于 softmax 和微调 CNN 的基线模型。
• 在神经机器翻译的束搜索中，SOFT top-k 操作通过支持可微解码与梯度流动，提升了生成质量。
• 在稀疏注意力机制中，该方法实现了显著的性能提升，证明了可微分 top-k 在注意力模型中的有效性。
• 在小规模实验中验证，SOFT 操作的梯度能够正确地将得分推向或移出 top-k 集合，具体取决于目标标签。