Atomic Design Transformer: Scaffold-Conditioned 3D Molecule Generation via xTB-Reward Reinforcement Learning

Abstract

We present an SE(3)-invariant transformer for 3D-molecule generation, the Atomic Design Transformer (ADT). ADT places atoms one at a time, autoregressively. SE(3) invariance is achieved by tokenization: each new atom's position is encoded in the local coordinate frame of a previously placed atom. The backbone is a plain causal transformer. The token stream fully specifies a 3D structure together with its chemical-bond graph G, without any bond-order assignment. To score generated molecules we introduce the xTB topology-preservation rate (XTP): the fraction of molecules for which an xTB GFN2 relaxation preserves G specified by the token stream. For XTP-accepted molecules we also report the relaxation energy and the root mean square of the atomic displacement (RMSD). We evaluate two ADT models. The first is ADT pretrained on the GEOM-Drugs $\le\!30$-heavy-atom dataset; we benchmark scaffold-conditioned 3D generation across seven drug-like scaffolds from the model. It reaches an XTP of ${\sim}54\%$ and a valid-molecule yield $N^{\mathrm{gen}}/N$ of ${\sim}50\%$, where $N^{\mathrm{gen}}/N$ is the fraction of samples that are distinct, topology-preserving, and chemically valid. The second model continues from the first by reinforcement learning against the verifiable xTB reward (RLVR), using no external molecules. RLVR raises XTP to ${\sim}98\%$ and $N^{\mathrm{gen}}/N$ to ${\sim}95\%$, while approximately preserving the GEOM-Drugs size and composition distributions. Finally, we present an Inverse-Kinematics Transformer that recovers XTP for large molecules, where discretization error accumulates. ADT thus enables direct 3D generation.