统一生成分子与材料的原子级扩散模型,性能媲美专用模型。
All-atom Diffusion Transformers: Unified generative modelling of molecules and materials
- 用统一编码器将分子和材料映射到共享隐空间,再通过Transformer扩散生成
- 在MP20、QM9等数据集上生成结果真实有效,性能达顶尖水平
- 模型轻量高效,适合构建通用化学生成基础模型
扩散模型是3D原子系统生成建模的标准工具。然而,针对不同原子系统(如分子与材料),生成过程通常高度依赖目标体系,尽管其底层物理一致。本文提出全原子扩散Transformer(ADiT),一种统一的潜在扩散框架,可使用同一模型联合生成周期性材料与非周期性分子系统:(1) 自编码器将分子与材料的统一全原子表示映射至共享隐空间;(2) 扩散模型训练生成新隐向量,自编码器可解码为新的分子或材料样本。在MP20、QM9和GEOM-DRUGS数据集上的实验表明,联合训练的ADiT能生成真实且有效的分子与材料,性能达到与分子及晶体专用模型相当的顶尖水平。ADiT采用标准Transformer结构,诱导偏置极小,训练与推理速度显著快于等变扩散模型。将模型扩展至五亿参数后性能稳步提升,标志着迈向通用化学生成基础模型的重要一步。开源代码:https://github.com/facebookresearch/all-atom-diffusion-transformer
原文摘要 · Abstract (English)
Diffusion models are the standard toolkit for generative modelling of 3D atomic systems. However, for different types of atomic systems -- such as molecules and materials -- the generative processes are usually highly specific to the target system despite the underlying physics being the same. We introduce the All-atom Diffusion Transformer (ADiT), a unified latent diffusion framework for jointly generating both periodic materials and non-periodic molecular systems using the same model: (1) An autoencoder maps a unified, all-atom representations of molecules and materials to a shared latent embedding space; and (2) A diffusion model is trained to generate new latent embeddings that the autoencoder can decode to sample new molecules or materials. Experiments on MP20, QM9 and GEOM-DRUGS datasets demonstrate that jointly trained ADiT generates realistic and valid molecules as well as materials, obtaining state-of-the-art results on par with molecule and crystal-specific models. ADiT uses standard Transformers with minimal inductive biases for both the autoencoder and diffusion model, resulting in significant speedups during training and inference compared to equivariant diffusion models. Scaling ADiT up to half a billion parameters predictably improves performance, representing a step towards broadly generalizable foundation models for generative chemistry. Open source code: https://github.com/facebookresearch/all-atom-diffusion-transformer
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。