[论文解读] All-atom Diffusion Transformers: Unified generative modelling of molecules and materials

Chaitanya K. Joshi, Xiang Fu+5 more|ArXiv.org|Mar 5, 2025

Machine Learning in Materials Science被引用 3

一句话总结

ADiT 引入一个单一潜在扩散 Transformer 框架，联合生成非周期性分子和周期性晶体，达到最先进的结果并实现跨领域的迁移学习。

ABSTRACT

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

研究动机与目标

• 在分子与晶体系统之间推动统一的扩散方法，分享其潜在的物理基础。
• 通过 VAE 为分子和材料建立共享潜在表示。
• 在潜在空间中训练扩散模型以生成新的原子，然后解码回有效结构。
• 展示周期性与非周期性系统之间的迁移学习并展示可扩展性。

提出的方法

• 将原子表示为具有统一属性的集合（类别原子类型、三维坐标、分数坐标、晶格参数）。
• 使用变分自编码器（VAE）将分子与晶体的输入映射到共享潜在空间并进行重构。
• 在潜在空间中训练扩散 Transformer (DiT) 以生成新的潜在嵌入，在采样过程中使用无分类器引导。
• 通过 VAE 解码器解码潜在向量，用 RDKit 产生有效的分子（非周期性）或晶体（周期性），分别由 PyMatGen 进行。
• 采用两阶段训练：第一阶段进行 VAE 重构，第二阶段进行带类别条件 DiT 的潜在扩散，并对旋转平移与周期对称进行数据增强。
• 用化学相关指标（有效性、稳定性、唯一性、新颖性）以及 Posebusters/DFT 常识性检查进行评估；与数据集特定基线进行比较。

实验结果

研究问题

• RQ1 单一统一的扩散模型是否能够在高有效性与真实感下同时生成分子和晶体？
• RQ2 在分子与晶体数据上进行联合训练是否通过迁移学习提升生成质量？
• RQ3 模型大小和训练规模如何相对于等变基线影响性能与速度？
• RQ4 该方法是否可扩展到更复杂的混合物，如 MOFs 以及其他周期性/非周期性混合系统？
• RQ5 与领域特定扩散模型相比，在训练与推理方面的效率如何？

主要发现

• 联合训练的 ADiT 在 QM9 分子与 MP20 晶体上实现了最先进或具有竞争力的结果。
• 在晶体方面，ADiT 在 10,000 个样本中相对于多种基线显示出更高的有效性、稳定性和新颖性。
• 在分子方面，ADiT 在有效性和唯一性方面达到或超过基线，同时满足 Posebusters 对物理真实感的检查。
• 联合训练比仅使用 QM9 或 MP20 的变体具有更高的有效性和稳定性，表明迁移学习的收益。
• ADiT 相对于等变扩散模型提供显著的推理加速，在单个 V100 GPU 上生成 10,000 个样本不足 20 分钟。
• 将 DiT 去噪器规模从 DiT-S 提升到 DiT-L 可带来可预测的性能提升，暗示在更多数据和更大模型下还有进一步收益。