[論文レビュー] Scalable Autoregressive 3D Molecule Generation

Austin H. Cheng, Chong Sun+1 more|ArXiv.org|May 20, 2025

Nanofabrication and Lithography Techniques被引用数 3

ひとこと要約

Quetzalは3D 分子を原子ごとに構築するスケーラブルな自己回帰モデルで、因果トランスフォーマーとDiffusion MLPを用いて離散的原子タイプと連続的位置を予測し、拡散モデルと競合する品質とより速い生成を実現します。

ABSTRACT

Generative models of 3D molecular structure play a rapidly growing role in the design and simulation of molecules. Diffusion models currently dominate the space of 3D molecule generation, while autoregressive models have trailed behind. In this work, we present Quetzal, a simple but scalable autoregressive model that builds molecules atom-by-atom in 3D. Treating each molecule as an ordered sequence of atoms, Quetzal combines a causal transformer that predicts the next atom's discrete type with a smaller Diffusion MLP that models the continuous next-position distribution. Compared to existing autoregressive baselines, Quetzal achieves substantial improvements in generation quality and is competitive with the performance of state-of-the-art diffusion models. In addition, by reducing the number of expensive forward passes through a dense transformer, Quetzal enables significantly faster generation speed, as well as exact divergence-based likelihood computation. Finally, without any architectural changes, Quetzal natively handles variable-size tasks like hydrogen decoration and scaffold completion. We hope that our work motivates a perspective on scalability and generality for generative modelling of 3D molecules.

研究の動機と目的

• 固定サイズ拡散モデルを超える自己回帰式3D分子生成の動機付け。
• 任意サイズの分子を出力しつつ正確な尤度計算を可能にするスケーラブルなアーキテクチャの開発。
• 離散的な原子タイプ予測子と連続位置モデルを組み合わせて3D構造を生成。
• QM9 および GEOM データセットで拡散モデルと競合するサンプル品質の実証。
• 構造を変更せずに水素装飾やスキャフォールド補完といった柔軟なタスクを示す。

提案手法

• 因果トランスフォーマを用いて次の原子タイプを予測し、Diffusion MLP を用いて次の原子位置をモデル化するアーキテクチャ。
• 次の位置分布 p_pos は DiffLoss 目的で訓練された条件付き拡散モデル。
• 原子タイプ予測は prefix 埋め込みから供給される MLP が出すロジットに対してソフトマックスを用いる。
• 位置条件付けベクトル z_i は次の原子タイプと prefix 表現を消費する第二のトランスフォーマによって生成。
• 訓練は DiffMLP のための複数タイムステップと時間とともにバッチ拡張を伴う全系列監視を使用。
• 生成は自己回帰的に進行し、停止トークンが出るまで原子タイプをサンプリングしてから位置をサンプリングする。トランスフォーマーのパスは O(n)、座標更新は O(n N_diff)。

実験結果

リサーチクエスチョン

• RQ1 従来の標準ベンチマーク（QM9, GEOM）で自己回帰型3D分子モデルは拡散モデルの生成品質に匹敵できるか。
• RQ2 提案する Quetzal アーキテクチャは生成を速くしつつ分子の有効性と一意性を維持できるか。
• RQ3 アーキテクチャを変更せずに水素装飾やスキャフォールド補完といった可変サイズ生成タスクをサポートできるか。
• RQ4 この自己回帰拡散フレームワークで原子ごとの座標の正確な尤度計算は実現可能か。
• RQ5 生成順序は性能と一般化にどのような影響を与え、一般的な順序選択にロバストか。

主な発見

• Quetzal は従来の自己回帰ベースラインを上回り、QM9 および GEOM におけるサンプル品質指標で最先端の拡散モデルと競合。
• 1つのトランスフォーマー・パスと座標用の DiffMLP を用いるため、拡散ベースのアプローチより生成が大幅に高速。
• このモデルは拡散のみの手法では一般的でない、厳密な発散ベースの尤度計算を可能にする。
• Quetzal は水素装飾やスキャフォールド補完といった可変サイズタスクをアーキテクチャ変更なしにネイティブにサポート。
• 水素装飾実験では、訓練データ順序と一致する場合に高い精度を示す。