InertialAR: Autoregressive 3D Molecule Generation with Inertial Frames
This work addresses the underexplored area of 3D molecule generation for computational chemistry and drug discovery, representing a novel method rather than an incremental improvement.
The paper tackled the problem of extending transformer-based autoregressive models to 3D molecule generation by addressing challenges in tokenization and architecture design, resulting in state-of-the-art performance on 7 out of 10 evaluation metrics for unconditional generation and all 5 metrics for controllable generation across datasets like QM9, GEOM-Drugs, and B3LYP.
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) tokenizing molecules into a canonical 1D sequence of tokens that is invariant to both SE(3) transformations and atom index permutations, and (2) designing 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. InertialAR devises a canonical tokenization that aligns molecules to their inertial frames and reorders atoms to ensure SE(3) and permutation invariance. Moreover, InertialAR equips the attention mechanism with geometric awareness via geometric rotary positional encoding (GeoRoPE). In addition, it utilizes a hierarchical autoregressive paradigm to predict the next atom-based token, predicting the atom type first and then its 3D coordinates via Diffusion loss. Experimentally, InertialAR achieves state-of-the-art performance on 7 of the 10 evaluation metrics for unconditional molecule generation across QM9, GEOM-Drugs, and B3LYP. Moreover, it significantly outperforms strong baselines in controllable generation for targeted chemical functionality, attaining state-of-the-art results across all 5 metrics.