SinkhornアルゴリズムとJAXによる実装: 最適輸送の高速化
最適輸送問題に以前から興味があり、最近のDeepSeekのMHC (Manifold Constrained Hyper Connections) 論文でも言及されていたため、その中心となるSinkhornアルゴリズムについて詳細を調べ、JAXで実装してみた際のメモ。
Warning
この記事は Google Antigravity を使用して作成されました。 あくまで私自身の勉強した結果の備忘録としてのメモと思っていただければと思います。 （正確性はかならずしも担保されません。） 作成過程で知らないことが多くあり、非常に勉強になりました。