ウォープレベルGPUリダクションを用いた高速な対数領域シンコーン型最適輸送

arXiv cs.LG / 2026/5/5

📰 ニュースDeveloper Stack & InfrastructureModels & Research

共有:

要点

この論文では、エントロピー正則化付き最適輸送（OT）に対する対数領域シンコーン法をネイティブCUDAで実装したFastSinkhornを提案しています。
ウォープレベルのシャッフルリダクションと共有メモリのタイル処理に加え、計算をすべて対数領域で行うことで、小さい正則化パラメータ（ε = 1e−4）でも数値的に安定した計算を実現します。
密なOT問題（n = m = 8192）でのベンチマークでは、POTライブラリに対して12×、GPU対応PyTorch基線に対して5.9×の高速化を示し、GPUメモリ使用量は256MBに抑えています。
画像の色変換や3D点群マッチングといった用途で検証し、実用性を裏付けるための収束解析も行っています。
総じて、フレームワーク由来のオーバーヘッドを減らしつつ、対数領域OTの安定性を保つには、ネイティブなGPUカーネルの丁寧な設計が有効であると主張しています。

要旨: Sinkhornアルゴリズムによるエントロピー正則化最適輸送（OT）は機械学習における基本的な手法となっていますが、既存の実装は小さな正則化パラメータに対して数値的に不安定になるか、あるいは深層学習フレームワークによる大きなオーバーヘッドを伴います。私たちは、ワープレベルのシャッフルによる縮約と共有メモリのタイル化を組み合わせて、数値安定性を損なうことなく高いGPU利用率を実現する、ログ領域Sinkhornアルゴリズムの軽量なネイティブCUDA実装であるFastSinkhornを提示します。私たちのソルバは完全にログ領域で動作するため、標準領域の手法が失敗するような正則化パラメータであるepsilon = 10^{-4} まででも頑健な計算が可能です。n = m = 8192 の密なOT問題において、私たちの実装は、広く用いられているPOTライブラリに対して12倍の高速化を達成し、GPUで加速したPyTorchのベースラインに対しても5.9倍の高速化を実現しつつ、使用するGPUメモリは256 MBのみです。画像の色転送、3D点群の対応付け、収束解析においてソルバを検証し、注意深い数値的取り扱いを備えたネイティブCUDAカーネルが、大規模な最適輸送計算のための実用的で効率的な基盤を提供することを示します。