Stable Global Weighting of Flow Mixtures using Simplex Exponential Moving Average

技術要約：Simplex Exponential Moving Averageを用いたフロー混合物の安定的なグローバル重み付け

問題提起
正規化フロー（Normalising flows）は、単純な基底分布を表現力豊かな確率密度へと変換することにより、変分推論（VI）の柔軟性を大幅に拡張してきました。しかし、個々のフローアーキテクチャ（例：RealNVP, MAF, RBIG）は、不均一な事後分布の幾何学的構造に対して汎化できないことがよくあります。異なるフローは、それぞれ異なる帰納バイアス、最適化ダイナミクス、および数値的特性を持っており、単一のアーキテクチャが多様なターゲット分布に対して一貫して高い性能を発揮することは困難です。混合ベースのアプローチは、相補的なフローを組み合わせるための自然な戦略を提供しますが、既存のフレームワークには以下のような重大な制限があります：

1. 共同最適化（Joint Optimisation）： 混合重みとエキスパートのパラメータを同時に学習することは、不均一な最適化ダイナミクスの結合により、不安定な学習や「エキスパートの支配（expert domination）」を引き起こす可能性があります。
2. 固定されたグローバル重み付け（Fixed Global Weighting）： 学習中に固定係数を推定し、推論時にそれらを一定に保つ手法は、相補的なエキスパートの振る舞いを活用する柔軟性に欠け、結果として効率的な容量割り当てが損なわれることがよくあります。

したがって、最も適切なアーキテクチャに関する事前知識を必要とせず、また共同学習による不安定さを招くこともない、安定した原理的な最適化フレームワークが必要です。

手法：AMF-VI-sEMA
本論文では、エキスパートの学習と混合の最適化を分離するために設計された2段階のフレームワークである AMF-VI-sEMA（Adaptive Mixture of Flows - Variational Inference with Simplex Exponential Moving Average）を提案します。

• ステージ1：独立したエキスパートの特化（Independent Expert Specialisation）
  不均一なエキスパートのセット（具体的にはRealNVP、MAF、RBIG）を、それぞれ異なる構造的レジームに特化するように独立して学習させます。各エキスパート $q_k(z|\phi_k)$ は、訓練データに対する自身の対数尤度を最大化するように最適化されます。一度学習が完了すると、エキスパートのパラメータ $\phi_k$ は凍結されます。このステージにより、各コンポーネントが独自の最適化ダイナミクスの下で収束することが保証され、共同学習における調整の不整合を回避できます。

• ステージ2：Simplex EMA (sEMA) グローバル重み付け
  エキスパートのパラメータが凍結された状態で、フレームワークは確率単体（probability simplex）上のグローバルな混合重みベクトル $\pi$ を学習します。これは、尤度駆動型の温度制御メカニズムを通じて達成されます：

  1. 責任度の計算（Responsibility Calculation）： ホールドアウトされた検証用データのバッチに対し、エキスパートの平均対数尤度を温度パラメータ $\tau$ でスケーリングしたソフトマックスを用いて、責任度 $r_k(z)$ を算出します。
  2. 平滑化と平均化（Smoothing and Averaging）： 責任度は検証バッチ全体で平均化され、過渡的な変動を抑えるために、オプションとして一様分布へと平滑化されます。
  3. Simplex EMA 更新（Simplex EMA Update）： グローバルな重み $\pi$ は、単体上で直接指数移動平均（EMA）を介して更新されます：$\pi \leftarrow \alpha \pi + (1-\alpha)\bar{r}$。ゼロ化を防ぐために小さな下限値 $\epsilon$ が適用され、その後に再正規化が行われます。

この設計は、サンプルごとのゲーティングや重みへのバックプロパゲーションを必要とせずに、適応的に容量を再配分する、扱いやすくデータに依存しないゲーティングメカニズムを生み出します。フレームワークは、コンポーネントの崩壊を診断するために、有効なエキスパート数（$N_{eff} = \exp(H(\pi))$）を監視します。

主な貢献

1. 分離された最適化フレームワーク： 本論文は、エキスパートの学習と混合の適応を分離する2段階の戦略を導入しています。エキスパートを凍結してから重みを更新することで、共同の責任重み付き最適化による不安定性を回避しつつ、アーキテクチャに依存しない手法を実現しています。
2. 実用的な適応型重み付けメカニズム： 著者らは、計算効率の高い責任ベースのSimplex EMA更新を提案しています。この手法は、共同最適化やサンプルごとのゲーティングを必要とせずにグローバルな混合重みを推定し、最小限の計算オーバーヘッドで安定したエキスパートの協調を提供します。
3. 包括的な実証評価： フレームワークは、6つの標準的な2D合成ファミリ（Banana, X-Shaped, Bimodal, Multimodal, Two-moons, Rings）と4つの実在または低次元のベイズ・ターゲット（BLR, BPR, Weibull, Real-GMM2）を含む10個のベンチマーク問題で評価されています。評価は標準的な対数尤度を超え、Wasserstein-2距離、最大平均偏差（MMD）、キャリブレーション分析、およびランダムシードに対する堅牢性まで多岐にわたります。
4. 最適化ダイナミクスの分析： 本論文は、エキスパートの利用状況、混合重みの進化、およびハイパーパラメータの感度に関する分析を提供し、なぜ不均一なフロー混合物が成功するのか、そしてsEMAメカニズムがいかに学習を安定させるのかについての洞察を提供しています。

結果
実証評価により、AMF-VI-sEMAが前身であるAMF-VIよりも一貫して優れた性能を示し、単一フローのベースライン（ResFlowやEM-Mixなど）で見られる破滅的な輸送失敗を回避していることが示されました。

• 性能： 本フレームワークは、すべての合成および実データセットにおいて、最高またはそれに準ずる負の対数尤度（NLL）を達成しています。また、頑健な輸送挙動（Wasserstein-2）と不一致指標（MMD）を維持しており、単一のアーキテクチャで見られる幾何学的な失敗（例：曲がった幾何学構造において高い輸送誤差を示すNICEやResFlow）を回避しています。
• 安定性： 本手法は、すべてのデータセットにおいて有効なエキスパート数 $N_{eff} > 1.4$ を維持しており、混合物が単一のエキスパートに崩壊していないことを示しています。
• 堅牢性： クロスシード評価は、NLL、KLダイバージェンス、およびWasserstein距離において低い分散を示しており、フレームワークが再現可能であり、ランダムな初期化に対して堅牢であることを裏付けています。
• 適応性： 学習されたグローバルな重みは、ターゲットの幾何学構造に適応します。例えば、回帰型の事後分布（BLR, BPR）では、メカニズムはRealNVPを抑制してMAFとRBIGを優先させる一方、リング状の構造（Rings）では、RealNVPに質量を強く集中させます。

意義と主張
本論文は、AMF-VI-sEMAが、どのフローアーキテクチャが最も適切であるかという事前知識を必要とせずに、不均一なランドスケープにおける事後分布近似のための堅牢で信頼できる最適化戦略を提供すると主張しています。本研究の意義は以下の通りです：

• 安定性： sEMAメカニズムは、時間的な正則化効果を導入することで、ゲーティングのダイナミクスを安定させ、激しいエキスパートの切り替えを緩和します。
• 簡潔さと効率性： このアプローチはアーキテクチャに依存せず、基礎となるフローモデルの変更を必要とせず、共同学習と比較して計算オーバーヘッドが極めて小さいものです。
• 信頼性： 共同の責任重み付き学習の脆弱性を回避することで、本フレームワークは複雑で非凸な、あるいはマルチモーダルな事後分布の幾何学構造をモデリングするための実用的なソリューションを提供します。

著者らは、単純な単峰性のターゲットに対しては特化した単一フローアーキテクチャの方が好ましい場合があることを認めていますが、AMF-VI-sEMAは多様な幾何学構造において、それらに対する堅牢性において一貫して優れていると述べています。今後の課題として、条件付きまたはアモータイズされたゲーティング戦略の探索や、高次元設定への拡張が示唆されています。