MolSnap: Snap-Fast Molecular Generation with Latent Variational Mean Flow

TL;DR

MolSnap introduces a causality-aware transformer and a variational mean flow framework for fast, diverse, and high-quality molecular generation conditioned on text, outperforming existing methods in efficiency and novelty.

Contribution

The paper presents a novel causality-aware transformer and a variational mean flow model that enhance molecular generation quality, diversity, and efficiency compared to prior approaches.

Findings

Achieves up to 74.5% novelty in generated molecules.

Maintains 100% validity across all benchmarks.

Requires only one or five NFEs for generation, improving computational efficiency.

Abstract

Molecular generation conditioned on textual descriptions is a fundamental task in computational chemistry and drug discovery. Existing methods often struggle to simultaneously ensure high-quality, diverse generation and fast inference. In this work, we propose a novel causality-aware framework that addresses these challenges through two key innovations. First, we introduce a Causality-Aware Transformer (CAT) that jointly encodes molecular graph tokens and text instructions while enforcing causal dependencies during generation. Second, we develop a Variational Mean Flow (VMF) framework that generalizes existing flow-based methods by modeling the latent space as a mixture of Gaussians, enhancing expressiveness beyond unimodal priors. VMF enables efficient one-step inference while maintaining strong generation quality and diversity. Extensive experiments on four standard molecular benchmarks demonstrate that our model outperforms state-of-the-art baselines, achieving higher novelty (up to 74.5\%), diversity (up to 70.3\%), and 100\% validity across all datasets. Moreover, VMF requires only one number of function evaluation (NFE) during conditional generation and up to five NFEs for unconditional generation, offering substantial computational efficiency over diffusion-based methods.