Sparsity-Constrained Optimal Transport

要約

タイトル: スパース性制約最適輸送
要約:
– 最適輸送 (OT) は、近年神経ネットワークの損失またはマッチングレイヤーとして、正則化された最適輸送がますます使用されるようになっています
– エントロピー正則化 OT は Sinkhorn アルゴリズムを使用して計算できますが、すべてのソースがすべてのターゲットと (分数の) マッチングされる、完全に密な輸送計画につながります
– この問題に対処するために、いくつかの作品は二次正則化を調査しています。この正則化はスパース性を保持し、制約のないスムーズな (半)デュアル目的につながります。
– しかしながら、二次正則化は輸送計画の非ゼロ数 (カージナリティ) を直接制御することができません。
– この論文では、輸送計画に明示的なカージナリティ制約を持つ OT の新しいアプローチを提案しています。
– カージナリティ制約は、画像パッチのような入力トークンを神経ネットワークのような専門家モデルとマッチングするために使用される sparse mixture of experts のアプリケーションにおいて、計算パフォーマンス上非常に重要です。
– カージナリティ制約の非凸性にもかかわらず、対応する (半)デュアル問題は扱いやすく、最初の勾配法で解決することができます。
– この手法は、非正則化OT ($k = 1$の場合に復元) と二次正則化OT ($k$が十分に大きい場合に復元) の中間地点として考えることができます。
– 目的関数の滑らかさは$k$が増加するにつれて増加し、収束速度と最適計画のスパースさのトレードオフが生じます。

要約(オリジナル)

Regularized optimal transport (OT) is now increasingly used as a loss or as a matching layer in neural networks. Entropy-regularized OT can be computed using the Sinkhorn algorithm but it leads to fully-dense transportation plans, meaning that all sources are (fractionally) matched with all targets. To address this issue, several works have investigated quadratic regularization instead. This regularization preserves sparsity and leads to unconstrained and smooth (semi) dual objectives, that can be solved with off-the-shelf gradient methods. Unfortunately, quadratic regularization does not give direct control over the cardinality (number of nonzeros) of the transportation plan. We propose in this paper a new approach for OT with explicit cardinality constraints on the transportation plan. Our work is motivated by an application to sparse mixture of experts, where OT can be used to match input tokens such as image patches with expert models such as neural networks. Cardinality constraints ensure that at most $k$ tokens are matched with an expert, which is crucial for computational performance reasons. Despite the nonconvexity of cardinality constraints, we show that the corresponding (semi) dual problems are tractable and can be solved with first-order gradient methods. Our method can be thought as a middle ground between unregularized OT (recovered in the limit case $k=1$) and quadratically-regularized OT (recovered when $k$ is large enough). The smoothness of the objectives increases as $k$ increases, giving rise to a trade-off between convergence speed and sparsity of the optimal plan.

arxiv情報

提供元, 利用サービス

arxiv.jp, OpenAI