Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching

要約

非正常密度またはエネルギー関数からサンプリングする拡散プロセスを学習するための非常にスケーラブルで効率的なアルゴリズムであるAdjointサンプリングを紹介します。
これは、エネルギー評価やモデルサンプルの数よりもはるかに勾配的な更新を可能にする最初のポリシーアプローチであり、同様の方法で以前に調査されたよりもはるかに大きな問題設定にスケーリングできるようにします。
私たちのフレームワークは、理論的には確率的最適制御に基づいており、サンプルがターゲット分布に向かってプッシュする是正措置を必要とせずに訓練することができる補間マッチングと同じ理論的保証を共有しています。
デカルト座標とねじれ座標の両方で分子をモデル化するために、主要な対称性と周期的境界条件を組み込む方法を示します。
古典的なエネルギー関数に関する広範な実験を通じて、私たちのアプローチの有効性を実証し、多くの分子システムで償却されたコンフォーマー生成を実行するニューラルネットワークベースのエネルギーモデルにさらに拡大します。
高度にスケーラブルなサンプリング方法の開発に関するさらなる研究を奨励するために、これらの挑戦的なベンチマークをオープンソースにすることを計画しています。

要約(オリジナル)

We introduce Adjoint Sampling, a highly scalable and efficient algorithm for learning diffusion processes that sample from unnormalized densities, or energy functions. It is the first on-policy approach that allows significantly more gradient updates than the number of energy evaluations and model samples, allowing us to scale to much larger problem settings than previously explored by similar methods. Our framework is theoretically grounded in stochastic optimal control and shares the same theoretical guarantees as Adjoint Matching, being able to train without the need for corrective measures that push samples towards the target distribution. We show how to incorporate key symmetries, as well as periodic boundary conditions, for modeling molecules in both cartesian and torsional coordinates. We demonstrate the effectiveness of our approach through extensive experiments on classical energy functions, and further scale up to neural network-based energy models where we perform amortized conformer generation across many molecular systems. To encourage further research in developing highly scalable sampling methods, we plan to open source these challenging benchmarks, where successful methods can directly impact progress in computational chemistry.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google