テクニカル・サマリー：最小MMD推定における勾配フローの観点

問題提起

最小最大平均埋め込み距離（Maximum Mean Discrepery; MMD）推定は、パラメータ推論のための堅牢な尤度フリーの代替手法として台頭している。これは、パラメトリックなモデル族 $\{P_\theta : \theta \in \Theta\}$ と未知のデータ生成分布 $Q$ との間の不一致を最小化するものである。最小MMD推定量の統計的性質（誤設定に対する堅牢性、尤度フリーな計算可能性など）は十分に確立されているが、関連する最適化問題 $\arg \min_{\theta \in \Theta} \text{MMD}^2(P_\theta, Q)$ の性質については、依然として理解が進んでいない。

既存の最適化手法は通常、標準的な（確率的）勾配降下法（GD）に依存している。しかし、MMD目的関数は一般に非凸であるため、標準的なGDは定常点への収束は保証するものの、大域的最小値への収束は保証しない。最小MMD最適化に関する既存の理論的収束保証（例：Chérief-Abdellatif and Alquier, 2022）は、実際には成立しないことが多い凸性の仮定に依存している。これは、目的関数のランドスケープが高度に非凸となる単純なガウス混合モデルによって示される通りである。本研究が取り組む中心的な未解決問題は、凸ではない目的関数に対しても収束保証を提供できる、最小MMD推定のための勾配ベースの最適化スキームを設計できるか否かである。

手法

著者らは、確率測度の空間における最近のMMD勾配フローの進展に着想を得た、**前処理付き勾配降下法（Preconditioned Gradient Descent; PGD）**スキームを提案する。この手法は、ノンパラメトリックな勾配フローとパラメトリックな最適化の間の溝を、以下の3つの主要なコンポーネントを通じて埋めるものである。

1. 適応的ノンパラメトリックMMD降下

著者らはまず、連続の方程式を通じて、分布 $P_t$ が確率測度の空間 $\mathcal{P}_2(\mathbb{R}^d)$ 内で直接進化するノンパラメトリックな降下スキームを分析する：
$$P_{t+1} = (\text{Id} - \gamma_t \nabla f_{\ell_t, P_t, Q}) \# P_t$$
ここで、$f_{\ell, P, Q}$ はウィットネス関数であり、$\nabla f$ は勾配場である。決定的なことに、著者らは適応的カーネル長スケールのスケジュール $(\ell_t)_{t \in \mathbb{N}}$ を導入している。この長スケールは、最初は大きく、ターゲットとなる $\ell > 0$ に向かって単調に減少していく。このカーネルの「アニーリング（焼きなまし）」は、目的関数を平滑化して局所的なトラップを回避し、最適化のランドスケープを改善することが示されている。軌跡的な勾配優位条件の下で、この適応的ノンパラメトリック・スキームがターゲット分布 $Q$ に大域的に収束することを証明している（定理3.1）。

2. パラメトリック部分空間への投影

これをパラメトリックモデル $P_\theta = (G_\theta)\#\rho$ に適用するために、著者らはノンパラメトリックな勾配フローを近似する降下方向を設計した。パラメーター $\theta$ に関する標準的なMMDの勾配を計算する代わりに、ノンパラメトリックな勾配場 $\nabla f_{\ell_t, P_{\theta_t}, Q} \circ G_{\theta_t}$ を、パラメーター $\theta$ の微小な更新によって実現可能な方向の線形部分空間へと投影する。
更新方向 $\Delta \theta_t$ は、以下の最小二乗問題を解くことで得られる：
$$\Delta \theta_t = \arg \min_{u \in \mathbb{R}^p} \mathbb{E}_{z \sim \rho} \left[ \| J_\theta G_{\theta_t}(z) u - \nabla f_{\ell_t, P_{\theta_t}, Q}(G_{\theta_t}(z)) \|^2 \right]$$
これは、前処理行列（ジェネレータのヤコビアンの経験的なフィッシャー類似行列）とリッジ正則化項 $\lambda_t$ を含む閉形式の解を与える：
$$\Delta \theta_t^{PGD} = \left( \mathbb{E}[J_\theta G_{\theta_t}^\top J_\theta G_{\theta_t}] + \lambda_t I \right)^{-1} \mathbb{E}[J_\theta G_{\theta_t}^\top \nabla f_{\ell_t, P_{\theta_t}, Q}(G_{\theta_t})]$$
これは、MMD自体の幾何学から導出された前処理器を用いる自然勾配降下法（NGD）とは異なり、PGDの前処理器は、パラメトリックな更新方向とノンパラメトリックな更新方向の間の不一致を最小化することから導出されている。

3. 理論的保証

著者らは、以下の2つの主要な条件の下で、PGDスキームの全域的収束を確立している（定理4.6）：

1. 勾配優位性（Gradient Dominance）: 勾配場のノルムとMMDの値（または誤設定の場合の最適MMDとのギャップ）を関連付ける条件。
2. 投影残差（Projection Residual）: ノンパラメトリックな勾配場がパラメトリックモデルの接空間に対して直交しておらず、かつ投影誤差が有界であることを保証する条件（仮定3）。

この定理は、適応的な長スケールがスケジュールされている限り、PGDスキームが非凸の設定においてもMMD目的関数を大域的最小値（または誤設定の場合の最良の近似値）へと導くことを証明している。

主な結果

本論文は、PGDスキームを支持する理論的および実証的な証拠を提供している。

• 理論的収束: 定理4.6は、MMD目的関数の凸性を仮定することなく、PGDの大域的漸近収束を確立している。これは、凸性を必要とした従来の研究からの重要な転換である。
• ノンパラメトリック収束: 定理3.1は、適応的な長スケールを用いたノンパラメトリックMMD降下がターゲット分布に収束することを証明し、「アニーリング」の直感を検証している。
• 実証的パフォーマンス:
  • ガウス混合モデル: 2モードのガウス混合モデルの例では、標準的なGDは局所解に捕まるが、PGDは大域的最小値への収束に成功する。
  • G-and-k 分布: この尤度フリーのベンチマークにおいて、PGDとNGDは標準的なGDを上回る。PGDはより少ない反復回数で最適値に到達し、困難なパラメータに対してもGDより優れたハンドリングを示す。
  • ロトカ・ヴォルテラ・モデル: データ汚染（外れ値）が存在する場合、PGDは様々な汚染レベルにおいて一貫して大域的最小値を回復するが、GDおよびNGDは初期値の悪さや高い汚染レベルの下で失敗する。
  • 複合適合度検定: トグルスイッチ遺伝子発現モデルにおいて、PGDは、GDが非凸性を緩和するためにマルチスタート・スキーム（400回の初期化）を使用した場合でも、較正されたType-Iエラーとより高い統計的検出力を達成する。PGDは、単一の初期化でこれを実現する。
• アブレーション研究: 実験により、適応的な長スケールと前処理器の両方が性能に寄与しており、その相対的な重要性は問題（例：ロトカ・ヴォルテラでは長スケールが、G-and-kでは前処理器がより重要）によって異なることが示されている。

意義と主張

本論文は、最小MMD推定の統計理論と、その実用的な計算との間の溝を埋めることを主張している。その主な意義は、目的関数の凸性を仮定せずに、最小MMD推定量の大域的収束保証を提供した最初の事例であることにある。

MMD勾配フローの幾何学的洞察を活用し、適応的な長スケール・スケジュールを導入することで、著者らは、非凸なランドスケープにおいて局所解を脱出する最適化スキームを設計することが可能であることを示している。提案されたPGDスキームは、明示的な（ただし非自明な）条件の下で理論的な収束保証を提供するだけでなく、モデルの誤設定やデータの汚染を含む多様な統計的推論タスクにおいて、標準的な勾配降下法や自然勾配降下法よりも堅牢かつ効率的に優れていることを実証している。

著者らは、自身の分析が、実用において検証が困難な、適応的長スケールと投影残差に関する事前条件に依存していることを指摘し、限界についても謙虚に述べている。また、これらの収束結果は漸近的であり、サンプルサイズと反復回数の両方に関する非漸近的なレートをまだ提供していないことも認めている。これらの手法を他の最小距離推定法に拡張することが今後の課題として示唆されている。