Sinkhornをそのまま速くする「Acc‑Sinkhorn」:エントロピー正則化最適輸送で収束が速くなる方法
研究の要点はシンプルです。エントロピー正則化最適輸送(EOT: entropy-regularized optimal transport)を解くために広く使われるSinkhornアルゴリズムの構造を保ちながら、収束速度を上げる新しい変種「Acc‑Sinkhorn」を提案しています。従来のSinkhornが誤差をO(1/k)で減らすのに対し、Acc‑Sinkhornは理論的にO(1/k^2)の速度を達成します。1回の反復あたりの計算コストやスケーリング形式はほぼ変わりません。つまり、実装の手間を大きく増やさずに効率が良くなる点が主な売りです。
やり方の核心は「バイレベル(2段階)最適化の見方」です。Sinkhornでは行のスケーリングと列のスケーリングを交互に行いますが、著者らはこれを内側の変数uを精確に解く手順と、外側の変数vに対する鏡像降下(mirror descent)として再解釈しました。この構造を利用して、Hessian(2次微分に相当する情報)に基づくNesterov型の加速を導入します。重要なのは、加速のために新たな高価な計算を入れるのではなく、既存のSinkhorn反復の外挿(過去の反復の組合せ)だけで済ませている点です。また、不確かな強凸性パラメータを扱うために外側で徐々にパラメータを下げる「ホモトピー」方式を使います。
理論的な成果も示されています。ある検証可能な安定性条件の下でAcc‑SinkhornはO(1/k^2)の収束率を証明します。これを非正則化(正則化を外す)最適輸送の近似問題にあてはめると、目標精度εに対する全体の計算量が(行列がn×nとすると)おおよそ Õ(n^2/ε) になります。従来のSinkhornは同じ設定で Õ(n^2/ε^2) 程度だったため、理論上は精度を上げたときの計算量依存が改善します。論文内では解析を簡単にするために正則化パラメータεを1にスケーリングして扱っていますが、実際の複雑度はこのスケーリングに依存する点に注意が必要です。
実際の例でも効果を示しています。合成データ、色変換(color transfer)、単語アライメントの実験では、正則化が小さい領域でSinkhornに比べて約10倍〜30倍の速度向上が得られたと報告しています。実験ではステップ幅としてα=√(2µ)(µは強凸性に相当するパラメータ)を使い、安定した加速挙動が観察されました。計算はSinkhornと同じO(nm)の回ごとのコストで、GPUでの並列化もそのまま活かせます。
重要な留意点もあります。理論的な完全保証は「可検証な安定性条件」の下で与えられますが、現在の解析はやや保守的で、実験で使われたステップ幅α=√(2µ)を理論的に正当化する部分にはギャップがあります。著者らはこのギャップがメトリック変動を制御する項(文献中の∥yk−x⋆∥の差分に相当)を均一に評価するのが難しいことに由来すると説明しています。安定性条件を取り除くことや、実験で安定だったステップ幅を完全に理論で裏付けることは今後の課題です。そのほか、手法はホモトピースケジュールや正則化の扱いに依存するため、適用時にはこれらの設計が必要になります。