InertialAR: Autoregressive 3D Molecule Generation with Inertial Frames

2025-10-31

AI Summary

This work addresses the dual challenges of SE(3) equivariance and atomic permutation invariance in autoregressive 3D molecular generation. Methodologically: (1) we propose an inertia-frame-aligned molecular serialization strategy ensuring both SE(3) and permutation invariance; (2) we introduce Geometric Rotation Positional Encoding (GeoRoPE), the first explicit positional encoding designed to model 3D rotational symmetry; (3) we design a discrete-continuous hybrid attention mechanism jointly modeling atom types and 3D coordinates; and (4) we adopt a diffusion-inspired coordinate prediction loss to enhance geometric fidelity. Our geometry-aware Transformer achieves state-of-the-art performance on 7 out of 10 metrics across QM9, GEOM-Drugs, and B3LYP benchmarks. Moreover, it attains comprehensive superiority—outperforming all baselines across all five evaluation metrics—in controllable functional-group generation tasks.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning
• Computer Vision: 3D Computer Vision

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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) 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.

Problem

• Developing SE(3) and permutation invariant tokenization for 3D molecules
• Modeling hybrid discrete-continuous atom tokens with geometric awareness
• Achieving state-of-the-art performance in unconditional and controllable molecule generation

Innovation

• Canonical tokenization ensures SE(3) and permutation invariance
• Geometric rotary positional encoding enhances attention mechanism
• Hierarchical autoregressive paradigm predicts atom type and coordinates