[论文解读] Scalable Autoregressive 3D Molecule Generation

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

Nanofabrication and Lithography Techniques被引用 3

一句话总结

Quetzal 是一个可扩展的自回归模型，逐原子构建三维分子，预测离散原子类型与连续位置，使用因果变换器和扩散多层感知机，质量可与扩散模型相媲美，生成更快。

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.

研究动机与目标

• 推动自回归三维分子生成超越固定尺寸扩散模型的局限。
• 开发一个可扩展架构，输出任意大小的分子，同时实现精确似然计算。
• 结合离散原子类型预测器与连续位置模型来生成三维结构。
• 在 QM9 和 GEOM 数据集上演示与扩散模型具有竞争的样本质量。
• 展示如氢装饰和骨架完成等灵活任务，而不需要架构改变。

提出的方法

• 模型架构结合因果变换器用于预测下一个原子类型，和扩散多层感知机用于建模下一个原子的位置。
• 下一个位置分布 p_pos 是一个条件扩散模型，使用 DiffLoss 目标进行训练。
• 原子类型预测通过对来自前缀嵌入的 MLP 产生的对数概率进行 softmax。
• 位置条件向量 z_i 由第二个变换器产生，该变换器接收下一个原子类型和前缀表示。
• 训练采用带有多时间步的全序列监督，并在 DiffMLP 上随时间进行批量扩展。
• 生成自回归地进行，先采样原子类型再采样位置，直到产生停止标记为止，需要 O(n) 次变换器前向和 O(n N_diff) 次坐标更新。

实验结果

研究问题

• RQ1 一个自回归的三维分子模型能否在标准基准上（QM9、GEOM）达到与扩散模型相当的生成质量？
• RQ2 提出的 Quetzal 架构是否在保持高效有效性的同时实现更快的生成速度？
• RQ3 模型是否能够在不改变架构的前提下支持如氢装饰和骨架完成等变尺寸生成任务？
• RQ4 在这种自回归扩散框架中，是否可以实现对每个原子坐标的精确似然计算？
• RQ5 生成顺序如何影响性能和泛化能力，模型是否对常见的排序选择具有鲁棒性？

主要发现

• Quetzal 在 QM9 和 GEOM 的样本质量指标上超越了先前的自回归基线，并与最先进扩散模型竞争。
• 由于对每个原子仅需一次变换器前向和 DiffMLP 更新坐标，生成速度显著快于基于扩散的方法。
• 该模型实现了精确的基于散度的似然计算能力，这在纯扩散方法中并不常见。
• Quetzal 自然支持氢装饰和骨架完成等变尺寸任务，无需架构变更。
• 氢装饰实验在原子排序与训练数据排序一致时表现出高准确性。