From Tokens to Blocks: A Block-Diffusion Perspective on Molecular Generation

TL;DR

SoftMol tackles the core limitations of GPT-style molecular language models by introducing a rule-free soft-fragment representation and a block-diffusion generator (SoftBD) that models local chemical context within blocks while autoregressively conditioning on prior blocks. It couples this generative core with adaptive decoding for efficient, high-fidelity sampling and a gated Monte Carlo tree search to assemble fragments toward target proteins, enforcing pharmacological feasibility before docking. Empirically, SoftMol achieves 100% chemical validity, up to 9.7% improvement in binding affinity, 2–3x increases in molecular diversity, and a 6.6x speedup in inference compared to strong baselines, demonstrating state-of-the-art performance in both de novo and target-specific design. The framework leverages a high-quality ZINC-Curated dataset and demonstrates robust robustness to soft-fragment length, offering practical, task-adaptive Granularity for efficient drug-design workflows with potential real-world impact and considerations for responsible deployment.

Abstract

Drug discovery can be viewed as a combinatorial search over an immense chemical space, motivating the development of deep generative models for de novo molecular design. Among these, GPT-based molecular language models (MLM) have shown strong molecular design performance by learning chemical syntax and semantics from large-scale data. However, existing MLMs face two fundamental limitations: they inadequately capture the graph-structured nature of molecules when formulated as next-token prediction problems, and they typically lack explicit mechanisms for target-aware generation. Here, we propose SoftMol, a unified framework that co-designs molecular representation, model architecture, and search strategy for target-aware molecular generation. SoftMol introduces soft fragments, a rule-free block representation of SMILES that enables diffusion-native modeling, and develops SoftBD, the first block-diffusion molecular language model that combines local bidirectional diffusion with autoregressive generation under molecular structural constraints. To favor generated molecules with high drug-likeness and synthetic accessibility, SoftBD is trained on a carefully curated dataset named ZINC-Curated. SoftMol further integrates a gated Monte Carlo tree search to assemble fragments in a target-aware manner. Experimental results show that, compared with current state-of-the-art models, SoftMol achieves 100% chemical validity, improves binding affinity by 9.7%, yields a 2-3x increase in molecular diversity, and delivers a 6.6x speedup in inference efficiency. Code is available at https://github.com/szu-aicourse/softmol

Figures (10)

• Figure 1: Comparison of molecular representation paradigms.Top: Atom-level SMILES provides minimal bias but lacks syntactic robustness. Middle: Rule-based fragments impose rigid vocabulary constraints and heuristic priors. Bottom: Our soft-fragment approach segments SMILES into fixed-length blocks, enabling tunable granularity and high robustness without heuristic rules or auxiliary tokens.
• Figure 2: Overview of the SoftMol framework. (a) Training: The concatenated clean and noised sequences are processed by the Block-Diffusion Transformer with a block-wise attention mask enforcing intra-block bidirectional and inter-block causal dependencies. (b) Sampling: Starting from [BOS], molecules are generated semi-autoregressively via iterative denoising, with previously decoded blocks cached as context. (c) Gated MCTS: Selection traverses the tree via UCT; Expansion uses batched SoftBD generation with duplicate filtering; Simulation applies a tunable feasibility gate---candidates satisfying pharmacological constraints proceed to docking, while failures receive a penalty; Backpropagation updates node statistics.
• Figure 3: Hit Ratio and DS Distributions.Top: Hit Ratio, defined as the proportion of unique generated molecules simultaneously satisfying drug-likeness (QED $> 0.5$, SA $< 5.0$) and binding affinity (DS $<$ median active) criteria. Bottom: Distribution of negative docking scores ($-$DS) for the identified hits satisfying all three criteria. Higher values indicate stronger affinity.
• Figure 4: Effect of Soft-Fragment Length. Heatmaps display Validity, Quality, Docking-Filter, Diversity, Uniqueness, and Sampling Time across the full $K_{\text{train}} \times K_{\text{sample}}$ grid.
• Figure 5: Ablation on Granularity and Search Budget. Hit Ratio and docking scores for varying $K_{\text{sample}}$ across different search budgets $N_{\max} \in \{1000, 5000, 10000\}$. Results are averaged across 5 targets from 50 independent runs.