Bayesian Flow Is All You Need to Sample Out-of-Distribution Chemical Spaces

Generating novel molecules with higher properties than the training space, namely the out-of-distribution generation, is important for de novo drug design. However, it is not easy for distribution learning-based models, for example diffusion models, to solve this challenge as these methods are designed to fit the distribution of training data as close as possible. In this paper, we show that Bayesian flow network, especially ChemBFN model, is capable of intrinsically generating high quality out-of-distribution samples that meet several scenarios. A reinforcement learning strategy is added to the ChemBFN and a controllable ordinary differential equation solver-like generating process is employed that accelerate the sampling processes. Most importantly, we introduce a semi-autoregressive strategy during training and inference that enhances the model performance and surpass the state-of-the-art models. A theoretical analysis of out-of-distribution generation in ChemBFN with semi-autoregressive approach is included as well.

💡 Research Summary

The paper tackles the challenging problem of generating molecules that lie outside the distribution of the training data (out‑of‑distribution, OOD) – a crucial capability for de‑novo drug design and materials discovery. While diffusion models (DMs) have become the dominant generative paradigm for chemistry, they are intrinsically designed to minimize the KL divergence between the model and the training distribution, which makes it difficult for them to deliberately explore regions of chemical space that contain higher‑property compounds.

The authors propose to use a Bayesian Flow Network (BFN), specifically the ChemBFN architecture, as a more natural OOD sampler. BFN differs from diffusion in that it does not require an explicit forward diffusion process nor a learned noise distribution; instead it directly optimizes the parameters of a distribution in a latent space (z) by approximating a reversed stochastic differential equation (SDE). This formulation works for continuous, discretised, and discrete data, and it can be guided more flexibly toward regions that are far from the training set.

Three complementary enhancements are introduced:

1. Online Reinforcement Learning (RL) Reward – A scalar reward term is added to the KL loss. At each continuous time step t, the categorical output distribution e(θ̂; t) is examined; if it corresponds to a chemically valid molecule, a reward of 1 is given, otherwise 0. The reward is weighted by a scaling factor η (empirically set to 0.01). This term encourages the model to allocate probability mass to valid token sequences even when the number of sampling steps is small, thereby raising the validity ratio of generated SMILES/SELFIES.

2. ODE‑like Sampling Procedure – By leveraging the monotonic accuracy schedule β(t) already used in ChemBFN, the authors devise an ODE‑style update for the latent variable:
   z←β(s)(K·e(θ̂; t)−1)+K·β(s)·τ·ε,
   where τ is a temperature coefficient controlling the magnitude of injected Gaussian noise ε∼N(0,I). Setting τ≈0.5 (or 0.05 for heterogeneous datasets) yields a sweet spot where validity exceeds 99 % while preserving a reasonable diversity of structures. The ODE‑like sampler dramatically reduces the number of required steps (from 1 000 to as few as 10) with only modest loss of novelty.

3. Semi‑Autoregressive (SAR) Strategy – Inspection of attention matrices in a trained ChemBFN revealed that entries far from the main diagonal are near zero, suggesting that future tokens contribute little to the current token’s update. By inserting a causal mask (zeroing out attention above the diagonal) the model updates all tokens in a block simultaneously but never uses future tokens to influence the current one. This “semi‑autoregressive” behavior can be applied during training, inference, or both, leading to four possible configurations (normal/normal, normal/SAR, SAR/normal, SAR/SAR). Empirically, the SAR‑SAR configuration combined with RL and ODE sampling delivers the best OOD performance.

Experimental Evaluation
The authors benchmarked their approach on several widely used datasets:

• MOSES and GuacaMol – standard small‑molecule generation suites that assess validity, uniqueness, novelty, and distributional distance (Fréchet ChemNet Distance, FCD).
• ZINC250k – a curated set of ~250 k molecules with associated quantitative estimate of drug‑likeness (QED), synthetic accessibility (SA), and docking scores (DS) against five protein targets. This benchmark evaluates multi‑objective OOD optimisation.
• Protein sequence dataset (≈91 k sequences) – used to test generation of larger biomolecules, with beta‑sheet percentage and solvent‑accessible surface area (SASA) as target properties.

Key findings include:

• Sampling Efficiency – The ODE‑like sampler reduces the required steps from 1 000 to 10 while maintaining >99 % validity.
• RL Impact – Adding the RL reward modestly improves validity and uniqueness, especially when the step count is low.
• Trade‑off Control – Temperature τ governs the balance between validity and diversity; τ < 0.01 yields near‑perfect validity but collapses diversity, whereas τ ≈ 1 maximises novelty at the cost of validity. τ = 0.5 works well for homogeneous molecule sets.
• SAR Benefits – The SAR configuration, particularly when applied both during training and inference (strategy 4), consistently outperforms the baseline ChemBFN across all metrics.
• State‑of‑the‑Art Gains – ChemBFN + RL + ODE (SAR + SAR) achieves validity 0.999, uniqueness 0.992, novelty 0.998, and FCD 0.797 on MOSES, surpassing prior SOTA diffusion‑based models. On GuacaMol it reaches validity 0.863, novelty 0.980. In the multi‑objective ZINC250k experiments, the method attains higher “novel hit ratio” and “novel top‑5 % DS” than competing approaches. For protein sequences, it improves target property attainment by 12‑18 % relative to baselines.

Interpretation and Significance
The work demonstrates that Bayesian flow models are not merely alternative density estimators but can be deliberately steered to explore chemical regions beyond the training manifold. By integrating a lightweight RL reward, a temperature‑scaled ODE sampler, and a semi‑autoregressive architecture, the authors achieve fast, high‑quality OOD generation without extensive architectural overhauls. This positions BFN‑based generators as practical tools for early‑stage drug discovery, where the ability to propose chemically valid, high‑property candidates outside known chemical space is a decisive advantage.

Limitations and Future Directions
The current approach still relies on manually tuned hyper‑parameters (η, τ) that may need dataset‑specific adjustment. SAR introduces causal masks that, while efficient for moderate sequence lengths, could increase memory consumption for very long polymers or proteins. Moreover, the RL reward is binary (valid/invalid) and could be extended to incorporate property‑based shaping (e.g., directly rewarding higher QED or docking scores). Future work may explore automated hyper‑parameter optimisation, memory‑efficient SAR implementations, and richer multi‑objective reward formulations.

In summary, the paper provides a compelling case that “Bayesian Flow Is All You Need” for OOD chemical space sampling, delivering a unified, efficient, and empirically superior framework that advances the frontier of generative chemistry.