Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization

技術要約：ワッサースタイン空間における差の凸計画法とそのMMD最適化への応用

問題設定
確率測度の空間上での汎関数最適化は、変分推論から生成モデリングに至るまで、機械学習のタスクにおいて中心的な役割を果たしている。標準的なアプローチは、ワッサースタイン勾配流を離散化することにより、ワッサースタイン空間 $(P_2(\mathbb{R}^d), W_2)$ 上で直接最適化を行うことである。凸汎関数（KLダイバージェンス、f-ダイバージェンスなど）に対しては効果的であるが、実用的に重要な目的関数の多くは、ワッサースタイン測地線に沿って非凸である。顕著な例として、二乗ワッサースタイン距離、スライス・ワッサースタイン距離、および最大平均偏差（MMD）が挙げられる。ワッサースタイン勾配降下法（WGD）のような標準的な一次手法は、これらの非凸な設定では苦戦することが多く、繊細なノイズ注入や、分布の重みを変化させたり暗黙的なステップを混合したりする代替空間（例：ワッサースタイン-MMDまたはワッサースタイン-フィッシャー・ラオ空間）での最適化を必要とする。

手法
著者らは、非凸汎関数に特化した新しい粒子ベースのアルゴリズムを提案している。このアルゴリズムは、差の凸（Difference-of-Convex; DC）分解をワッサースタイン空間へと持ち込むものである。コアとなる手法は、古典的な凸凹（Convex-Concave）手順（CCCP）をワッサースタイン空間へと拡張した、**ワッサースタイン凸凹手順（WCCCP）**である。

1. ワッサースタイン空間におけるDC分解： 目的汎関数 $F$ を、$F(\mu) = F_+(\mu) - F_-(\mu)$ と分解する。ここで、$F_+$ と $F_-$ は共に $P_2(\mathbb{R}^d)$ 上の全凸（totally convex）な汎関数である。凹部分をポテンシャルエネルギー（線形汎関数）に限定していた先行研究とは異なり、本手法では $F_-$ を任意の微分可能な汎関数とすることができる。
2. アルゴリズム構成： 各反復 $k$ において、与えられた測度 $\mu_k$ に対し、アルゴリズムは主要化最小化（majorization-minimization）問題を解く。これは、$\mu_k$ の周りで凹部分 $F_-$ を線形化することで上界を構築し、輸送写像 $T \in L^2(\mu_k)$ の空間上で、以下の凸サーロゲートを最小化するものである：
   $$T_{k+1} = \arg\min_{T \in L^2(\mu_k)} \left( F_+(T\#\mu_k) - \langle \nabla_{W_2}F_-(\mu_k), T - \text{Id} \rangle_{L^2(\mu_k)} \right)$$
   次のイテレートは $\mu_{k+1} = (T_{k+1})\#\mu_k$ として定義される。
3. 実装： 内部の最小化は $L^2(\mu_k)$ 上の勾配降下法によって行われる。経験的分布（粒子）の場合、これは粒子の位置の更新へと翻訳される。著者らは、WCCCPと、ワッサースタイン空間におけるミラー降下法およびブレグマン近接降下法との関連性も確立している。
4. MMDの詳細： 最大平均偏差（MMD）に対しては、カーネル $k(x,y) = \psi(x-y)$ 自体を、凸成分 $\psi = \psi_+ - \psi_-$ へと分割することによる明示的なDC分解を提案している。これには、ジョルダン分解や代数的な分割（例：$\cosh/\sinh$）を用いて、ガウスカーネルやリーゼスカーネルを含む移動不変カーネルに対する具体的な分解を提供している。

主要な貢献

• WCCCPアルゴリズム： 凹部分を線形ポテンシャルに制限することなく、ワッサースタイン空間へのDC最適化を拡張する、ワッサースタイン凸凹手順の導入。
• 理論的収束性：
  • イテレートに沿って目的関数の値 $F(\mu_k)$ が非増加することの証明。
  • **漸近的定常性（almost stationarity）**の結果の確立。各成分の滑らかさと強凸性（全凸性）の仮定の下で、イテレート間の距離の二乗とワッサースタイン勾配のノルムに関する劣線形収束率を証明している。具体的には、$\min_{k} \|\nabla_{W_2}F(\mu_k)\|^2 = O(1/K)$ である。
  • 本スキームが、特定の条件下でブレグマン近接降下法およびミラー降下法と等価であることの提示。
• MMDのための明示的なDC分解： 本論文は、移動不変カーネルを持つMMD汎関数を、収束保証のための必要な凸性と滑らかさの条件を満たすDC形式へと分解する体系的な方法を提供している。
• 実証的検証： エネルギー距離およびガウスカーネルを用いたMMDに関する実験において、適切に選択されたDC分解は、標準的なWGDおよびワッサースタイン近接勾配法（前方後方法）よりも高速かつ安定した収束を示すことが示された。

結果

• エネルギー距離： 合成データ（スパイラルおよびキャット形状）およびCIFAR10画像において、WCCCPはWGDおよび前方後方（FB）スキームよりも大幅に速く収束する。このアルゴリズムは、複雑なターゲット分布に一致させるためにエネルギー距離を最小化することに成功している。
• ガウスカーネルを用いたMMD： ガウスカーネルを用いたMMDの最小化は、バンド幅と初期値に対して極めて敏感であることで知られている。著者らは、WGDおよび（標準的な分解を用いた）FBスキームが、粒子がターゲットから離れていく局所解に陥ることが多いことを示している。対照的に、特定のカーネル分解（ジョルダンまたは $\cosh/\sinh$）を用いたWCCCPは、より優れた収束を実現する。ジョルダン分解は、より滑らかで数値的不安定性が少ないことが指摘されている一方、$\cosh/\sinh$ 分解は、速度は遅くなるものの、すべての粒子が最終的に収束することを保証する。
• 比較： 結果は、WCCCPの性能がDC分解の選択に強く依存していることを浮き彫りにしている。「適切に選ばれた」分解こそが、標準的な勾配降下法を凌駕するための鍵となる。

意義と主張
本論文は、標準的な勾配手法が非凸性のために失敗する領域に対し、ワッサースタイン空間における非凸汎関数の最適化のための堅牢なフレームワークを提供すると主張している。CCCPをこの設定へと持ち上げることで、著者らは、従来のMMD特有のアプローチに見られる重み混合の問題を回避する、理論的に裏付けられた代替案を提示している。本研究は、アルゴリズムが一般的である一方で、その実用的な成功はカーネルの具体的なDC分解に依存することを強調している。著者らは謙虚に、異なる分解の影響をより深く理解し、自動的かつ適応的な選択戦略を設計するためには、今後の研究が必要であると述べている。さらに、現在の理論的解析は内部の最適化問題の厳密な解を仮定しており、不正確なソルバーを考慮に入れることが、将来の研究における必要な方向性として特定されている。