NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation

TL;DR

NEAT is a novel set transformer model for 3D molecular generation that is permutation-invariant and achieves state-of-the-art quality with faster generation times.

Contribution

Introduces NEAT, a neighborhood-guided, order-agnostic transformer that improves molecular generation by ensuring permutation invariance and efficiency.

Findings

Achieves state-of-the-art results on QM9 and GEOM-Drugs datasets.

Offers significant speed improvements over existing models.

Ensures permutation invariance in atom-level generation.

Abstract

Transformer-based autoregressive models offer an efficient alternative to diffusion- and flow-matching-based approaches for generating 3D molecules. One challenge remains: standard transformer architectures require a sequential ordering of tokens, which is not inherently defined for the atoms in a molecule. Prior works have addressed this by using canonical atom orderings. However, these approaches are not permutation invariant w.r.t. atoms and bias next-token prediction towards ordering conventions. We overcome this limitation by introducing a novel neighborhood-guided training strategy. Our model, NEAT (Neighborhood-Guided, Efficient, Autoregressive Set Transformer) treats molecular graphs as sets of atoms and learns an order-agnostic distribution over admissible tokens at the graph boundary, thereby ensuring atom-level permutation invariance. NEAT achieves state-of-the-art generation quality on the QM9 and GEOM-Drugs datasets while offering a significant speed advantage over existing baselines.