arXiv · 2512.05844
NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular 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.
Explore related subjects
Keep this discovery
Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer. 2025-12-05. NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation. https://arxiv.org/abs/2512.05844
Cite the original work for its findings. Save a collection to share your selection of sources.