STAR-VAE: Latent Variable Transformers for Scalable and Controllable Molecular Generation

TL;DR

STAR-VAE tackles scalable, conditional molecular generation by marrying a Transformer-based encoder with an autoregressive Transformer decoder in a latent-variable framework trained on SELFIES. A property predictor provides a unified conditioning signal that shapes the latent prior, approximate posterior, and decoder, while LoRA adapters enable efficient, data-limited fine-tuning. Empirically, the model matches or exceeds strong baselines on MOSES and GuacaMol for unconditional generation and demonstrates meaningful, target-specific shifts in docking-score distributions on Tartarus, with latent spaces that are smooth and semantically structured. The approach shows that a modern, large-scale VAE remains competitive for molecular design when combined with principled conditioning and parameter-efficient adaptation, enabling both broad exploration and property-guided optimization in drug discovery.

Abstract

The chemical space of drug-like molecules is vast, motivating the development of generative models that must learn broad chemical distributions, enable conditional generation by capturing structure-property representations, and provide fast molecular generation. Meeting the objectives depends on modeling choices, including the probabilistic modeling approach, the conditional generative formulation, the architecture, and the molecular input representation. To address the challenges, we present STAR-VAE (Selfies-encoded, Transformer-based, AutoRegressive Variational Auto Encoder), a scalable latent-variable framework with a Transformer encoder and an autoregressive Transformer decoder. It is trained on 79 million drug-like molecules from PubChem, using SELFIES to guarantee syntactic validity. The latent-variable formulation enables conditional generation: a property predictor supplies a conditioning signal that is applied consistently to the latent prior, the inference network, and the decoder. Our contributions are: (i) a Transformer-based latent-variable encoder-decoder model trained on SELFIES representations; (ii) a principled conditional latent-variable formulation for property-guided generation; and (iii) efficient finetuning with low-rank adapters (LoRA) in both encoder and decoder, enabling fast adaptation with limited property and activity data. On the GuacaMol and MOSES benchmarks, our approach matches or exceeds baselines, and latent-space analyses reveal smooth, semantically structured representations that support both unconditional exploration and property-aware generation. On the Tartarus benchmarks, the conditional model shifts docking-score distributions toward stronger predicted binding. These results suggest that a modernized, scale-appropriate VAE remains competitive for molecular generation when paired with principled conditioning and parameter-efficient finetuning.

Figures (3)

• Figure 1: The figure shows two scatterplots, where each plot shows the UMAP distribution of the molecules generated with CVAE conditioned with scores ranging from minimum to maximum (colors) for SA Score (Left) and B3DB (Right).
• Figure 2: The figure shows three scatterplots, where each plot shows the UMAP distribution of the molecules generated with CVAE for each target conditioned with docking scores from -10 to 0 (colors).
• Figure 3: Distribution of docking scores for the Tartarus training set (top row), for 1000 VAE‑generated molecules per target (middle row), and for 1,000 CVAE‑generated molecules per target (bottom row). Each histogram displays the frequency of docking‑score values across the respective sample set.