[論文レビュー] 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

ひとこと要約

本稿では、エントロピー最適輸送（EOT）を用いた、微分可能かつスケーラブルなトップ-k演算子SOFT（Scalable Optimal transport-based diFferenTiable）を提案する。EOTを用いてトップ-k演算子の微分可能近似を構築し、機械学習モデルにおけるエンドツーエンド学習を可能にする。トップ-k選択を正則化された最適輸送問題として定式化し、SinkhornアルゴリズムとKKT条件を活用することで、トップ-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.

研究の動機と目的

• 標準的なトップ-k演算子の非微分性がニューラルネットワークにおけるエンドツーエンド学習を妨える問題に対処すること。
• 勾配ベース最適化を可能にする、滑らかで連続的なトップ-k演算子の近似を構築すること。
• 統一的な微分可能トップ-k演算子を用いて、微分可能なk-NN分類、ビームサーチ、スパースアテンション機構を実現すること。
• エントロピー正則化とSinkhornアルゴリズムを用いて、近似誤差を制御しつつ計算効率を維持すること。

提案手法

• トップ-k演算子を、選択された上位k個の要素を輸送計画としてパrameter化した最適輸送（OT）問題の解として再定式化する。
• OT問題にエントロピー正則化を適用し、微分可能な近似を生成することで、エントロピー最適輸送（EOT）の定式化を導出する。
• Sinkhornアルゴリズムを用いて最適輸送計画を効率的に計算し、スケーラブルかつ微分可能なトップ-k選択を実現する。
• EOT問題のカールシュ＝クーン＝タッカー（KKT）条件を用いて、SOFTトップ-k演算子の勾配を導出することで、バックプロパゲーションによるトップ-k層の勾配伝播を可能にする。
• k-NN分類器、ビームサーチ、スパースアテンション層などの下流モデルにSOFT演算子を統合する。
• 計算量が入力要素数nに対してO(n)となるように実装し、スケーラビリティを確保する。

実験結果

リサーチクエスチョン

• RQ1 トップ-k演算子を、選択の意味論を保持したまま微分可能にできるか？
• RQ2 エントロピー最適輸送が、誤差を制御可能な滑らかで連続的な離散的トップ-k演算子への近似を提供できるか？
• RQ3 提案された微分可能トップ-k演算子が、k-NN分類、ビームサーチ、アテンション機構におけるエンドツーエンド学習性能を向上できるか？
• RQ4 バックプロパゲーション中にSOFT演算子の勾配はどのように振る舞い、望ましい最適化目的と整合しているか？

主な発見

• SOFTトップ-k演算子により、従来はトップ-kの非微分性のため2段階学習が必要だったモデルのエンドツーエンド学習が可能になった。
• k-NN分類では、ソフトマックスベースおよび事前学習済みCNNベースのベースライン手法よりも高い精度を達成した。
• ニューラル機械翻訳におけるビームサーチでは、勾配伝播を可能にする微分可能なデコーディングが可能になり、生成品質が向上した。
• スパースアテンション機構においても顕著な性能向上が得られ、アテンションベースのモデルにおける微分可能トップ-kの有効性が実証された。
• トイラベル実験により、SOFT演算子の勾配が、目的ラベルに応じてスコアをトップ-k集合へ向かってまたはそれから遠ざけつつ適切にプッシュしていることが検証された。