Energy Rank Alignment: Using Preference Optimization to Search Chemical Space at Scale

要約

考えられる分子の数は原子の数と組み合わせて増加するため、化学空間の探索は非常に困難な問題です。
化合物のデータベースで訓練された大規模な自己回帰モデルは強力なジェネレーターを生み出しましたが、望ましい特性を持つ分子を生成するための堅牢な戦略がまだ不足しています。
この分子探索問題は、大規模な言語モデルの「アライメント」問題によく似ていますが、多くの化学タスクについては、特定の簡単に評価できる報酬関数があります。
ここでは、エネルギー ランク アラインメント (ERA) と呼ばれるアルゴリズムを導入します。このアルゴリズムは、明示的な報酬関数を利用して、自己回帰ポリシーの最適化に使用する勾配ベースの目標を生成します。
我々は、このアルゴリズムが近接ポリシー最適化 (PPO) および直接優先最適化 (DPO) と密接に関連しているが、エネルギー関数の役割を果たす報酬を備えた理想的なギブス ボルツマン分布に収束するミニマイザーを備えていることを理論的に示します。
さらに、このアルゴリズムは拡張性が高く、強化学習を必要とせず、ペアごとの優先観測の数が少ない場合に DPO と比較して良好なパフォーマンスを発揮します。
私たちは、このアプローチを分子トランスフォーマーを整列させて外部で指定された特性を持つ分子を生成するために導入し、化学空間のさまざまな部分を探索しながら堅牢に実行することを発見しました。
ここでは化学的検索に焦点を当てていますが、LLM アライメントのための AI 監視タスクでも優れた結果が得られ、この方法が拡張可能で汎用性があることが示されています。

要約(オリジナル)

Searching through chemical space is an exceptionally challenging problem because the number of possible molecules grows combinatorially with the number of atoms. Large, autoregressive models trained on databases of chemical compounds have yielded powerful generators, but we still lack robust strategies for generating molecules with desired properties. This molecular search problem closely resembles the ‘alignment’ problem for large language models, though for many chemical tasks we have a specific and easily evaluable reward function. Here, we introduce an algorithm called energy rank alignment (ERA) that leverages an explicit reward function to produce a gradient-based objective that we use to optimize autoregressive policies. We show theoretically that this algorithm is closely related to proximal policy optimization (PPO) and direct preference optimization (DPO), but has a minimizer that converges to an ideal Gibbs-Boltzmann distribution with the reward playing the role of an energy function. Furthermore, this algorithm is highly scalable, does not require reinforcement learning, and performs well relative to DPO when the number of preference observations per pairing is small. We deploy this approach to align molecular transformers to generate molecules with externally specified properties and find that it does so robustly, searching through diverse parts of chemical space. While our focus here is on chemical search, we also obtain excellent results on an AI supervised task for LLM alignment, showing that the method is scalable and general.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google