[論文レビュー] 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 は非周期性分子と周期結晶の両方を同時に生成する単一の潜在拡散トランスフォーマー フレームワークを導入し、最先端の成果と領域間転移学習の成功を達成します。

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で開発する。
• 潜在空間で拡散モデルを訓練して新しい原子をサンプルし、正当な構造へデコードする。
• 周期性・非周期性系間の転移学習を実証し、拡張性を示す。

提案手法

• 原子を統一属性（カテゴリ的原子種、3D座標、分数座標、格子パラメータ）を持つ集合として表現する。
• VAE を用いて分子と結晶の入力を共有潜在空間へ写像し再構成する。
• 潜在空間でDiffusion Transformer（DiT）を訓練し、新しい潜在埋め込みを生成し、サンプリング時に分類器フリーガイダンスを用いる。
• VAEデコーダで潜在をデコードして、分子（非周期）または結晶（周期）をそれぞれRDKitとPyMatGenで生成する。
• 二段階訓練を採用：第一段階はVAE 再構成、第二段階はクラス条件付きDiTと回転・並進・周期対称性のデータ拡張を用いた潜在拡散。
• 化学に焦点を当てた指標（有効性、安定性、唯一性、新規性）とPosebusters/DFT整合性チェックで評価する；データセット固有のベースラインと比較する。

実験結果

リサーチクエスチョン

• RQ1 単一の統一拡散モデルが分子と結晶の両方を高い有効性とリアリズムで生成できるか。
• RQ2 分子データと結晶データを共同訓練することで転移学習を通じて生成品質が向上するか。
• RQ3 モデルサイズと訓練スケールがエクイバリアントなベースラインに対する性能と速度にどのように影響するか。
• RQ4 MOF や他の周期性/非周期性のハイブリッドなど、より複雑なハイブリッドへ拡張性はあるか。
• RQ5 ドメイン特化型拡散モデルと比較した訓練・推論の効率性はどの程度か。

主な発見

• 共同訓練された ADiT は QM9 分子と MP20 結晶の双方で最先端または競合的な結果を達成。
• 結晶において、ADiT は10,000サンプルで有効性、安定性、Noveltyが複数のベースラインを上回る。
• 分子において、ADiT は有効性と唯一性を満たしつつ物理的現実性を示す Posebusters チェックを満たす。
• 共同訓練は QM9 のみ・MP20 のみのバリアントより有効性と安定性が高く、転移学習の利点を示唆。
• ADiT は等方拡散モデルに比べ推論速度の大幅な改善を提供し、単一の V100 GPU で10,000 サンプルを20分未満で生成可能。
• DiT ノイズ除去器のサイズを DiT-S から DiT-L にスケールすると予測可能な性能向上が得られ、より多くのデータと大型モデルでさらなる恩恵が見込まれる。