Learning-Order Autoregressive Models with Application to Molecular Graph Generation

Zhe Wang, Jiaxin Shi, Nicolas Heess, Arthur Gretton, Michalis Titsias

Abstract: Autoregressive models (ARMs) have become the workhorse for sequence generation tasks, since many problems can be modeled as next-token prediction. While there appears to be a natural ordering for text (i.e., left-to-right), for many data types, such as graphs, the canonical ordering is less obvious. To address this problem, we introduce a variant of ARM that generates high-dimensional data using a probabilistic ordering that is sequentially inferred from data. This model incorporates a trainable probability distribution, referred to as an order-policy, that dynamically decides the autoregressive order in a state-dependent manner. To train the model, we introduce a variational lower bound on the exact log-likelihood, which we optimize with stochastic gradient estimation. We demonstrate experimentally that our method can learn meaningful autoregressive orderings in image and graph generation. On the challenging domain of molecular graph generation, we achieve state-of-the-art results on the QM9 and ZINC250k benchmarks, evaluated using the Fréchet ChemNet Distance (FCD), Synthetic Accessibility Score (SAS), Quantitative Estimate of Drug-likeness (QED).

Lay Summary: We have developed an AI model that learns the optimal step-by-step sequence to build complex structures like molecules. While standard models can easily generate text with its predictable order, they struggle with molecules where the best construction path is not obvious. Our model intelligently overcomes this by deciding the most effective piece to add at each stage, such as the next atom or chemical bond. When tested on the challenging task of designing new molecules, this method achieved state-of-the-art results on key industry benchmarks, producing molecules highly rated for their validity, stability, and potential as useful drugs.

Primary Area: Probabilistic Methods->Variational Inference

Keywords: Generative Modeling, Variational Inference, Autoregressive Model, Graph Generation, Molecule Generation

Submission Number: 12612