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

2025-11-04

AI Summary

This study addresses three key challenges in molecular generation: broad chemical space coverage, weak property controllability, and low generation efficiency. Methodologically, we propose a Transformer-based variational autoencoder (VAE) framework that employs SELFIES for syntax-valid molecular representation; adopts an encoder–decoder latent variable architecture to jointly model unconditional and property-conditioned generation; and integrates LoRA (Low-Rank Adaptation) for efficient fine-tuning. Evaluated on the GuacaMol and MOSES benchmarks, our approach achieves state-of-the-art (SOTA) performance or surpasses leading baselines across multiple metrics. In the Tartarus docking task, it significantly improves the distribution of predicted binding affinities, demonstrating superior property precision control, structural diversity, and chemical validity. Collectively, these results validate the framework’s effectiveness in balancing broad chemical space exploration with targeted property optimization and high-fidelity generation.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation
• Computer Vision: Generative Adversarial Networks (GANs) for Vision

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• Developing scalable generative models for vast chemical space exploration
• Enabling conditional molecular generation based on specific properties
• Providing fast molecular generation with guaranteed syntactic validity

Innovation

• Transformer-based VAE with SELFIES encoding
• Conditional latent variable for property-guided generation
• Low-rank adapters for efficient fine-tuning