Chaitanya K. Joshi Xiang Fu Yi-Lun Liao Vahe Gharakhanyan Benjamin Kurt Miller Anuroop Sriram Zachary Ward Ulissi

摘要

扩散模型是生成3D原子系统的标准工具。然而，对于不同类型的原子系统（如分子和材料），尽管其基本物理原理相同，生成过程通常高度特定于目标系统。我们引入了全原子扩散变压器（All-atom Diffusion Transformer, ADiT），这是一种统一的潜在扩散框架，可以使用同一模型同时生成周期性材料和非周期性分子系统：(1) 自编码器将分子和材料的统一全原子表示映射到共享的潜在嵌入空间；(2) 扩散模型被训练以生成新的潜在嵌入，自编码器可以解码这些嵌入以采样新的分子或材料。在MP20、QM9和GEOM-DRUGS数据集上的实验表明，联合训练的ADiT能够生成真实且有效的分子及材料，其性能与专门针对分子和晶体的模型相当，达到了当前最佳水平。ADiT在自编码器和扩散模型中使用标准变压器并具有最小的归纳偏差，这使得其在训练和推理过程中相比等变扩散模型显著加速。将ADiT扩展至5亿参数可预见地提升了性能，标志着向广泛通用的基础模型迈进了一步，为生成化学领域的发展提供了支持。开源代码：https://github.com/facebookresearch/all-atom-diffusion-transformer