🤖 AI

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

The paper presents NEAT, a novel neighborhood-guided autoregressive set transformer that achieves state-of-the-art quality in 3D molecule generation with substantial speed advantages on the QM9 and GEOM-Drugs datasets by ensuring permutation invariance at the atomic level through an order-independent training strategy.

Original authors: Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

Published 2026-05-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to build a complex 3D structure, such as a molecular model made of multi-colored spheres (atoms) connected by rods (bonds). The goal is for the robot to create new, valid structures from scratch rather than merely copying existing ones.

This work introduces a new robot architect named NEAT (Neighborhood-Guided, Efficient, Autoregressive Set Transformer). Here is how it works, simply explained:

The Problem: The "Order" Confusion

Most previous robots attempted to build molecules by following a strict sequence, like students entering a classroom one after another. They had to decide: "Who comes first? Who comes second?"

The Problem: In a molecule, atoms actually have neither a "first" nor a "second" place. They all exist simultaneously in a 3D cluster. Forcing them into an order is like trying to describe a snowflake by listing its crystals in a specific sequence. It is arbitrary and confusing. If the robot chooses the wrong order, the molecule falls apart.
The Old Way: Previous methods tried to solve this by inventing a "canonical" (standard) order, which, however, distorted the robot. It learned the order of the list better than the shape of the molecule.

The Solution: The "Neighborhood" Strategy

NEAT changes the game. Instead of a sequence, it views the molecule as a set.

How it works: Imagine you are building a house brick by brick. Instead of asking, "Which brick comes next in the list?", NEAT looks at the current wall you have built and asks, "Who are the neighbors standing directly next to this wall?"
The Magic: NEAT examines the "boundary" of the current structure. It knows that every atom touching the current structure is a valid candidate to be added next. The order does not matter to it; it cares only about the neighborhood.
The Result: This makes the robot "permutation invariant." It does not matter whether you view the molecule from the left, right, or upside down; the robot sees the same neighbors and builds the same valid structure.

The Construction Process: Two Steps Simultaneously

When NEAT decides to add a new atom, it does two things at once:

What is it? It guesses the atom type (e.g., "This should be a carbon sphere").
Where is it? It uses a special mathematical technique called Flow Matching to determine exactly where in 3D space this atom should float. Think of it as a river current gently pushing a leaf (the atom) from a random starting point to its perfect, stable position within the molecule.

Why is NEAT Special?

The authors claim NEAT is a major advancement for three reasons:

It is fast: Since it does not need to simulate thousands of noisy steps like other robots (which is like trying to carve a statue by working on a block of stone for hours), NEAT builds molecules much faster. The work states that it is 5 to 16 times faster than its competitors.
It is flexible: Since it is not based on a fixed order, you can give it a "starter kit" (a partial molecule) and ask it to finish the job. It can complete the structure regardless of how you hand over the starter kit. Other robots often fail if the starter kit does not match their specific "order" rules.
It is accurate: In standard tests (datasets named QM9 and GEOM-Drugs), NEAT builds molecules that are chemically valid and unique, matching or surpassing the best existing methods.

The Catch (Limitations)

The authors are honest about what NEAT cannot yet do:

Size Limits: It can only build molecules as large as those it saw during training. It cannot extrapolate to build a huge molecule if it only learned on small ones.
Bond Ambiguity: It relies on knowing which atoms are connected (bonds) to determine the neighborhood. If you give it a strange crystal structure where the bonds are unclear, it might get confused.
No "Magical" 3D Rules: It does not have built-in physical rules for 3D rotation (SE(3) equivariance), but the authors found that randomly rotating the training data was sufficient for it to work well anyway.

Summary

In short, NEAT is a new way to teach computers to design 3D molecules. Instead of forcing them to follow a rigid, artificial list, it teaches them to look at their immediate neighbors and build naturally. It is faster, more flexible, and just as accurate as current state-of-the-art methods.

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

Problem Statement

Current computational workflows for discovering bioactive molecules rely heavily on virtual screening or de novo generation. While virtual screening is limited by the finite size of chemical libraries and increasing computational costs, generative modeling offers a promising alternative. However, the generation of 3D molecular structures presents unique challenges. Molecules are inherently three-dimensional, and their properties are often determined by conformation.

Existing state-of-the-art models for 3D generation predominantly utilize diffusion and flow-matching techniques. These methods frequently suffer from quadratic scaling in message-passing updates, lack inherent adaptability to variable molecular sizes (often requiring a fixed number of atoms determined at initialization), and struggle to handle dynamic atom insertions during denoising.

Autoregressive models offer a flexible alternative that naturally supports variable atom counts and requires only one encoder call per token. However, standard autoregressive transformers require a sequential arrangement of tokens. In molecular graphs, atoms have no inherent sequential order. Previous work has addressed this by using canonical atom orderings, yet these approaches are not permutation-invariant with respect to atoms. They distort the prediction of the next token in favor of specific ordering conventions, thereby limiting the model's ability to generalize or perform tasks such as prefix completion, where the starting atom sequence may not match the canonical ordering of the final molecule. Conversely, Graph Neural Network (GNN) approaches that ensure permutation invariance often generate molecules with low chemical validity.

Methodology

The authors propose NEAT (Neighborhood-Guided, Efficient, Autoregressive Set Transformer), the first autoregressive transformer model for 3D molecular generation that respects the permutation invariance of atoms without relying on canonical atom orderings.

Core Architecture

NEAT treats molecular graphs as sets of atoms rather than sequences. The architecture consists of three main components:

Data Augmentation Module (Neighborhood Guidance): Instead of sequential positional encodings, the model employs a sampling strategy to mimic molecular growth. Given a molecular graph, a connected subgraph (source set) is sampled. The target set is defined as the "boundary nodes" (1-hop neighbors) of the source set. This provides a set of valid token continuations for the next step, independent of the order in which atoms were added.
Set Transformer Backbone: The source set is encoded using a Set Transformer (based on NanoGPT/GPT-2 with bidirectional attention and without positional embeddings). Atom types are linearly embedded, and positions are encoded using Fourier features. Additive pooling aggregates these node representations into a set-based representation ( $z_n$ ). This design ensures permutation invariance with respect to the order of atoms in the source set.
Flow-Matching Head: The set-based representation $z_n$ $z_{n}$ is used to predict the next token, factored into atom type and position: $p(a_{n+1}, x_{n+1} | z_n) = p(a_{n+1} | z_n) \cdot p(x_{n+1} | a_{n+1}, z_n)$ $p (a_{n + 1}, x_{n + 1} ∣ z_{n}) = p (a_{n + 1} ∣ z_{n}) \cdot p (x_{n + 1} ∣ a_{n + 1}, z_{n})$ .
- Atom Type: Modeled as a categorical distribution over a linear layer and Softmax.
- Atom Position: Modeled using Conditional Flow Matching (CFM). The model learns a vector field to transport a random initial position distribution to the true atom positions.

Training Strategy

Neighborhood Guidance: During training, the model is not supervised by a single next token but by the set of all valid boundary nodes. This eliminates the need for canonicalization.
Teacher Forcing: The prediction of the atom type and the reconstruction of the position are decoupled. The model predicts the type, and the Flow-Matching head learns the position conditioned on this type.
Augmentation: To improve robustness against error accumulation, Gaussian noise is injected into a fraction of the atoms in the source set during training, and positions are re-centered relative to the noisy source set.
Loss Function: The total loss is the sum of the Cross-Entropy loss for atom types and the Flow-Matching loss for positions. A linear assignment problem is solved to couple source and target distributions, minimizing transport costs and avoiding crossing paths.

Inference

Generation begins with a random initial atom. At each step, the current set of atoms is encoded, the next atom type is predicted, and a random initial position is sampled. The Flow-Matching vector field is then integrated (using Euler or Euler-Maruyama methods) to determine the final 3D coordinates. This process repeats until a stop token is predicted.

Key Contributions

Permutation-Invariant Autoregressive Generation: NEAT introduces a novel training strategy using neighborhood guidance and Set Transformers to achieve 3D molecular generation that is invariant to atom ordering, overcoming the limitations of canonical orderings in autoregressive models.
Efficiency: By treating the input as a set and using a single encoder call per token, NEAT avoids the quadratic scaling of diffusion-based message passing as well as the need for multiple integration steps per atom generation step found in some flow-matching baselines.
Prefix Completion Capability: The order independence of NEAT enables high-quality prefix completion (generating molecules around a given scaffold) without requiring the prefix to match a specific canonical ordering, a task where sequence-based models like QUETZAL struggle.
State-of-the-Art Performance: The model achieves competitive or superior results on standard benchmarks (QM9 and GEOM-Drugs) compared to both autoregressive and diffusion/flow-matching baselines.

Experimental Results

The model was evaluated on the QM9 (small organic molecules) and GEOM-Drugs (drug-like molecules) datasets.

Generation Quality: On QM9, NEAT outperformed previous autoregressive models (G-SchNet, G-SphereNet, Symphony, Geo2Seq) in terms of xyz2mol validity and uniqueness. It achieved an xyz2mol validity of 95.2% and a uniqueness of 91.6%, securing second place regarding lookup validity and uniqueness. On GEOM, NEAT achieved the highest lookup validity (99.7%) and the second-best xyz2mol validity (93.9%) as well as uniqueness (91.6%) among autoregressive models, matching the performance of state-of-the-art flow-matching models like FlowMol3.
Prefix Completion: In a task where 1,000 molecules were generated from 100 diverse prefixes, NEAT significantly outperformed QUETZAL. NEAT achieved 97.8% lookup validity and 92.7% xyz2mol validity, whereas QUETZAL achieved only 27.8% and 28.1%, respectively. This demonstrates NEAT's ability to generalize structural features independent of atom ordering conventions.
Runtime Efficiency: NEAT shows significant speed advantages. It runs approximately 5 times faster than SemlaFlow and 16 times faster than FlowMol3. In worst-case scenarios, NEAT's complexity is $O(BN^3)$ (where $N$ is the number of atoms), while diffusion/flow models scale as $O(BTN^2)$ (where $T$ is the number of integration steps). Since $N$ is typically smaller than $T$ for drug-like molecules, NEAT offers a practical speed advantage.

Significance and Claims

The work claims that NEAT provides a strong foundation for employing transformer architectures in complex drug design tasks, specifically for conditional 3D molecular generation. By eliminating the dependency on canonical atom orderings, the model enables flexible tasks such as scaffold-based generation (prefix completion), which were previously difficult for autoregressive approaches.

The authors note that while the current architecture is not explicitly SE(3)-equivariant, the use of random rotation augmentation during training effectively compensates for this, resulting in only minor variances in predictions. The work suggests that integrating NEAT with an SE(3)-equivariant transformer is a promising direction for future research.

The authors emphasize that although the model generates high-quality 3D structures, molecule generation is only one step in a complex pipeline requiring extensive validation, synthesis, and testing. They also acknowledge limitations, such as the dependency on edge information for neighborhood guidance (which may limit applicability to inorganic crystals) and the lack of extrapolation beyond the maximum number of atoms seen during training.