InertialAR: Autoregressive 3D Molecule Generation with Inertial Frames

AI summary

Overview

Research area: Machine learning / generative modeling for 3D molecular structures, specifically autoregressive (AR) sequence modeling combined with geometric symmetry handling (SE(3) invariance and atom-index permutation invariance).

Technical level: Advanced. The paper assumes familiarity with Transformer attention, autoregressive likelihood factorization, diffusion loss, and molecular symmetry groups.

Scope: The paper introduces InertialAR, a hierarchical autoregressive Transformer that tokenizes 3D molecules into canonical sequences using inertial frames, adds geometry-aware positional encodings, and decodes atom types and 3D coordinates separately.

What This Paper Is About

Autoregressive models work well for text and images, but extending them to 3D molecules is hard for two reasons: there is no obvious way to turn a 3D structure into a single canonical 1D token sequence that is unchanged by rotations, translations, and atom renumbering, and each "token" is a hybrid of a discrete atom type plus three continuous coordinates.

The paper's goal is to solve both problems so that a standard Transformer can generate 3D molecular conformations one atom at a time, and to test whether this approach can match or beat diffusion-based models that dominate the field.

Key Contributions

1. Generation-oriented canonical tokenization. Each molecule is aligned to its canonical inertial frame (computed from the inertia tensor) to remove rotational and translational ambiguity, then its atoms are deterministically reordered using the canonical atom ranking from RDKit's canonical SMILES serialization. This maps any 3D structure to a unique, SE(3)-invariant and permutation-invariant token sequence. The authors state this is the first work to turn inertial-frame canonicalization into a generation-ready tokenization for autoregressive 3D generation.

2. Geometric Positional Encoding (GeoPE). A positional encoding scheme that injects geometry into attention through two branches: ManhattanPE, which rotates query and key components by an axis-wise sinusoidal encoding so that attention scores depend only on relative positions, and EuclideanPE, which encodes pairwise Euclidean distances via an RBF kernel approximated with the Nyström method so the architecture stays compatible with standard Transformer attention.

3. A hierarchical autoregressive decoding paradigm. At each step the model predicts the atom type first with a cross-entropy loss, then predicts the 3D coordinates conditioned on that type using Diffusion Loss (denoising a Gaussian-perturbed coordinate). The projection weight matrices are block-diagonal so the type and Euclidean branches are projected independently.

4. Comprehensive empirical evaluation. State-of-the-art results on 8 of 10 evaluation metrics for unconditional generation across QM9, GEOM-Drugs, and B3LYP, plus state-of-the-art results on all 5 metrics for class-conditional generation on QM9 using classifier-free guidance.

Main Findings

• Unconditional generation on QM9: InertialAR achieves the highest scores among compared methods on Validity (97.3%), Atom Stability (99.2%), and Molecule Stability (94.5%), with Valid & Unique at 92.5%. The strongest autoregressive baseline reported, Geo2Seq, reaches 97.1% validity and 93.2% molecule stability but only 81.7% Valid & Unique.

• Unconditional generation on GEOM-Drugs: InertialAR attains the best Atom Stability at 87.0% and reports 98.0% validity, ahead of the compared baselines on atom stability.

• Large-scale B3LYP results: On the 1M-molecule B3LYP subset, InertialAR scores 99.3% valid, 98.8% valid & unique, 84.8% atom stability, and 24.7% molecule stability. The two baselines evaluated there are EDM (92.9 / 92.8 / 80.6 / 0.8) and Geo2Seq (73.3 / 2.7 / 10.0 / 0.0). The paper describes the autoregressive baseline Geo2Seq as performing poorly on this dataset.

• Class-conditional generation: Conditioning on the 5 most frequent Molecule Class IDs (functional-group configurations), InertialAR reaches an average hit rate of 83.3%, versus the reported baseline values for the same classes (EDM ranging from 8.9% to 37.5%, Geo2Seq from 33.8% to 49.4%). For example, on Class 7 (with Ether) InertialAR records a 90.9% hit rate and 97.5% molecule stability; on Class 3 (with Hydroxyl), 85.7% hit rate and 99.4% molecule stability.

• Molecule editing via classifier-free guidance: Increasing the CFG guidance scale by 0.3 (Δs = 0.3) transforms molecules that lacked the required functional groups into ones that satisfy the target class while producing more plausible 3D geometries.

• Ablations: Replacing Diffusion Loss with direct L2 regression causes a sharp degradation in generation quality, and removing deterministic atom reordering consistently harms validity and uniqueness.

• Robustness of the canonicalization: The paper reports that principal-moment degeneracy affects only 0.007% of QM9 and is effectively zero on GEOM-Drugs, and that perturbation-induced instabilities are negligible. The authors explicitly state their contribution is not a theoretically perfect canonicalization method, but a demonstration that this simple strategy is robust enough for state-of-the-art results.

• Not reported: The paper does not report model parameter counts, training compute, sampling wall-clock time, or memory usage.

Methodology in Plain English

The paper treats molecule generation like language modeling, one atom at a time. The first problem is making the "sentence" unique. If you list atoms in an arbitrary order or leave a molecule in an arbitrary rotation, the same molecule produces many different sequences, which makes training hard. The authors fix this in two steps: they rotate and translate the molecule so it sits in its own natural coordinate system derived from its mass distribution (its inertia tensor), using a designated fourth atom to resolve which way the axes should point; then they renumber the atoms using the ordering RDKit already produces when it writes a canonical SMILES string. After this, each molecule corresponds to exactly one token sequence.

The second problem is that attention has no idea where atoms are in space. The authors give each atom a positional encoding. One part, ManhattanPE, works like rotary position encoding but extends it to three axes, so the dot product between two atoms' queries and keys depends only on how far apart they are, not on where the whole molecule happens to be. The other part, EuclideanPE, explicitly encodes the straight-line distance between pairs of atoms, using a low-rank approximation so that the distance information shows up naturally in the attention dot product without modifying the attention formula itself.

The third problem is that an atom token is half discrete (an element) and half continuous (three coordinates). Rather than predicting both at once, the model first picks the atom type with a standard classification loss, then, given that choice, predicts the coordinates by training a small denoiser to remove added Gaussian noise from the true position and sampling by iterative denoising. For conditional generation, the model is trained with labels sometimes dropped, and at inference the conditional and unconditional predictions are mixed with a guidance scale, which lets the user trade off class adherence against diversity.

Why This Matters

Impact on research. The work argues that autoregressive models can be a competitive alternative to diffusion for 3D molecule generation, which has been dominated by diffusion and equivariant graph networks. It also shows an input-canonicalization route to symmetry handling that avoids building equivariant architectures or relying on data augmentation, and it demonstrates that the approach scales to a much larger and more chemically diverse dataset than the commonly used QM9.

Real-world applications:

• Drug discovery: targeted generation of candidate molecules carrying specific functional groups such as hydroxyl, ether, or heteroaromatic rings, which is exactly the class-conditional task evaluated here.
• Lead optimization and molecule editing: the CFG results show that raising the guidance scale can add a missing functional group while keeping the geometry plausible, resembling a structure-editing tool.
• Materials science: the authors propose extending the method to periodic material discovery, where 3D structure generation is similarly central.
• Foundation model components for scientific domains: the paper positions the autoregressive formulation as a possible building block for broader multimodal scientific models.

Industry relevance. Pharmaceutical and materials companies rely heavily on generating and screening 3D conformations; an autoregressive model that supports variable-length generation natively and avoids the iterative sampling cost associated with diffusion (the paper notes diffusion's computationally intensive sampling as a motivation) is directly relevant to high-throughput virtual screening pipelines. The released code at github.com/HaoruiLi46/InertialAR lowers the barrier for adoption.

Future Directions

• Extending the framework to protein structure modeling, where the same SE(3)-invariance and sequence-ordering questions arise at a much larger scale.
• Applying it to periodic material discovery, which involves additional symmetries beyond molecular SE(3) alone.
• Integrating InertialAR into broader multimodal frameworks, as the authors suggest.
• Open questions the paper leaves unresolved include whether the inertial-frame canonicalization remains stable for structures with near-degenerate principal moments at scale, and how far the approach extends to datasets even larger than the 1M-molecule B3LYP subset used here.

Target Audience

Machine learning researchers working on generative models for scientific data, especially those interested in autoregressive sequence models, SE(3)-invariant representation, or positional encoding design. It is also relevant to computational chemists and cheminformatics practitioners who generate or screen 3D molecular conformations and want to understand the current state of controllable 3D generation. Readers without background in Transformer attention, diffusion losses, or molecular symmetry will find the method sections demanding.

Authors’ abstract

Transformer-based autoregressive models have emerged as a unifying paradigm across modalities such as text and images, but their extension to 3D molecule generation remains underexplored. The gap stems from two fundamental challenges: (1) how to tokenize molecules into a canonical 1D sequence of tokens that is invariant to both SE(3) transformations and atom index permutations, and (2) how to design an architecture capable of modeling hybrid atom-based tokens that couple discrete atom types with continuous 3D coordinates. To address these challenges, we introduce InertialAR. It first performs generation-oriented canonical tokenization by aligning each molecule to a canonical inertial frame and reordering atoms, thereby converting arbitrary 3D structures into a unique, SE(3)- and permutation-invariant sequence of tokens for autoregressive generation. Built upon this canonical tokenization, we propose geometric positional encoding (GeoPE), which endows Transformer attention with 3D geometric awareness. Finally, InertialAR utilizes a hierarchical autoregressive paradigm to decode the next atom, consecutively predicting the atom type and 3D coordinates via Diffusion Loss. Experimentally, InertialAR achieves state-of-the-art performance on 8 of the 10 evaluation metrics for unconditional generation across QM9, GEOM-Drugs, and B3LYP. Moreover, it significantly outperforms baselines in controllable generation for targeted chemical functionality, attaining state-of-the-art results across all 5 metrics. Code is available at github.com/HaoruiLi46/InertialAR.