[論文レビュー] Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks

David Alvarez-Melis, Yair Schiff+1 more|Digital Access to Scholarship at Harvard (DASH) (Harvard University)|Jun 1, 2021

Topological and Geometric Data Analysis参考文献 49被引用数 9

ひとこと要約

本稿では、入力凸ニューラルネットワーク（ICNN）を用いて凸関数をパrameter化することで、確率測度空間上の汎関数を最適化するスケーラブルな手法JKO-ICNNを提案する。この手法により、収束保証のもとで安定した高次元最適化が可能となり、低次元PDEおよび高次元分子生成の両方で検証された。特に、薬物様性（QED）が向上した。

ABSTRACT

Gradient flows are a powerful tool for optimizing functionals in general metric spaces, including the space of probabilities endowed with the Wasserstein metric. A typical approach to solving this optimization problem relies on its connection to the dynamic formulation of optimal transport and the celebrated Jordan-Kinderlehrer-Otto (JKO) scheme. However, this formulation involves optimization over convex functions, which is challenging, especially in high dimensions. In this work, we propose an approach that relies on the recently introduced input-convex neural networks (ICNN) to parametrize the space of convex functions in order to approximate the JKO scheme, as well as in designing functionals over measures that enjoy convergence guarantees. We derive a computationally efficient implementation of this JKO-ICNN framework and experimentally demonstrate its feasibility and validity in approximating solutions of low-dimensional partial differential equations with known solutions. We also demonstrate its viability in high-dimensional applications through an experiment in controlled generation for molecular discovery.

研究の動機と目的

• 高次元における確率測度上の汎関数最適化の計算的課題に対処すること。
• 特に高次元空間において、JKOスキームにおける凸関数最適化の困難さを克服すること。
• 深層学習を用いたWasserstein空間上の勾配フローのためのスケーラブルで微分可能なフレームワークを開発すること。
• 理論的収束保証のもとで、分子生成のような複雑な確率汎関数の実用的最適化を可能にすること。
• 既知の解が存在する低次元PDEおよび高次元の現実世界の応用の両方において、本手法の実現可能性を示すこと。

提案手法

• JKOスキームにおける凸ポテンシャル関数を、入力に対して凸性を保証する入力凸ニューラルネットワーク（ICNN）でパrameter化する。
• Brenierの定理を用いて、JKO変分問題を凸関数に関する最適化問題に再定式化し、微分可能な最適化を可能にする。
• 有限サンプルと自動微分を用いた適応的勾配降下法を実装し、エンドツーエンドの学習を可能にする。
• 分子埋め込みからのQEDスコアを予測する凸な代替モデル（リーマン・ICNN）を用い、ポテンシャル汎関数における凸性を保証する。
• 学習可能な重みを備えた分布間距離測度（Sinkhorn距離またはMMD）を統合し、ターゲット分布へ向かう流れをガイドする。
• 未観測のサンプルに一般化可能なトランスポート写像を学習することで計算コストを amortized し、再訓練なしに効率的な推論を可能にする。

実験結果

リサーチクエスチョン

• RQ1 ICNNは、高次元確率最適化に必要なJKOスキームにおける凸関数空間を効果的にパrameter化できるか？
• RQ2 JKO-ICNNフレームワークは、高次元におけるスケーラブルで微分可能な最適化を実現しつつ、収束保証を維持できるか？
• RQ3 JKO-ICNNは、解析解が既知の低次元PDEの解をどれほど正確に近似できるか？
• RQ4 JKO-ICNNは、高次元分子生成における複雑な汎関数（例：薬物様性（QED））を効果的に最適化できるか？
• RQ5 異なる分布間距離測度（Sinkhorn対MMD）および重み付け戦略が、モード崩壊および最適化の安定性に与える影響は何か？

主な発見

• JKO-ICNNは、既知の解析解が存在する低次元PDEの解を正確に近似し、収束性と精度を示した。
• QM9データセットにおいて、λ₂ = 1,000でMMDを用いた場合、中央値QEDは0.315から0.452に上昇し、相対的に44.8%の向上を達成した。
• Sinkhorn距離とλ₂ = 10,000を用いた場合、中央値QEDは0.419に達し、初期分布から相対的に33.0%の向上を示した。
• 分子生成において、有効性が92.7%、一意性が81.9%を達成し、モードカバレッジと構造的品質の両面で効果的であることが示された。
• 実験から、高値のλ₂（10,000）またはMMDベースの距離測度がモード崩壊を引き起こすことが判明し、ハイパーパramータ選択への感受性が顕在された。
• 学習されたトランスポート写像は未観測のサンプルに一般化可能であり、再訓練なしに効率的な推論が可能であり、計算のamortizationが実現された。