REVIEW 4 major objections 4 minor 12 references
Sampling Out-of-Distribution Chemical Spaces via Bayesian Flow
T0 review · 4 major / 4 minor · reviewed 2026-08-11 · deepseek-v4-flash
Pith's one-line read Bayesian flow networks can sample molecules beyond their training distribution, with better docking scores than prior models.
desk verdict Useful empirical trick (causal mask for ChemBFN) and strong benchmark numbers, but the OOD claim is not actually demonstrated because the training-set docking distribution is never reported. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central mechanism is the semi-autoregressive causal mask applied to ChemBFN, a Bayesian flow network that generates molecules token-by-token in fixed-length strings such as SMILES or SELFIES. In a Bayesian flow network the generator does not learn a noise-removal process; it updates the parameters of a categorical distribution over each token through a sequence of Bayesian updates, and a neural network estimates the clean-token distribution at each time step. The causal mask forces each token to be updated using only itself and tokens to its left, although all tokens are still updated in parallel; the paper justifies this by observing that trained attention matrices concentrate on the diagonal. Two auxiliary pieces carry the practical gains: a reinforcement-learning term added to the training loss that rewards intermediate token distributions corresponding to valid molecules, and an ODE-like sampling update in latent space, $z \leftarrow \beta(s)(K e(\hat\theta; t)-1)+\sqrt{K\beta(s)}\tau\epsilon$, where $\tau$ is a temperature that trades validity against diversity.
What would settle it
Measure the attention matrices of a trained ChemBFN at increasing classifier-free guidance strengths while conditioning on a far-from-distribution property vector; if the total weight far from the main diagonal rises above a small threshold, the hard causal mask is removing information the model relies on. A second check is to train on a dataset rich in long-range dependencies, such as macrocycles, and compare the fully semi-autoregressive strategy with the unmasked baseline; if the valid-SMILES ratio or docking scores drop, the sparsity assumption fails for that distribution.
Extended reading notes
Core claim
On the paper's own terms, the central discovery is that out-of-distribution generation is not a defect to be corrected in Bayesian flow networks but an intrinsic property that can be amplified. The load-bearing observation is that in a trained ChemBFN, self-attention entries far from the diagonal are essentially zero; therefore a hard causal mask can be applied to the bidirectional token updates, producing a semi-autoregressive model that still learns chemistry. Depending on whether the mask is active during training, during sampling, or both, the model moves into different regions of chemical space, and when a classifier-free guidance vector points beyond the training distribution, the masked variants achieve lower (better) docking scores on the top five percent of novel hits than all compared state-of-the-art baselines on all five proteins. In the SELFIES version the same masks raise the fraction of generated molecules that pass all novelty and property filters from under six percent to more than 25 percent.
Load-bearing premise
The load-bearing premise is the empirical observation that in a trained ChemBFN the attention weights far from the main diagonal are essentially zero, so a hard causal mask can be applied without discarding information the model needs; the abstract promises a theoretical analysis, but the paper supplies only this empirical observation, and the argument collapses if that sparsity is not stable under guidance strength, dataset, or training run.
Editorial extensions
If this is right
- If the claims are correct, molecule generators based on Bayesian flow networks can be steered to regions of chemical space with better predicted binding affinity than anything in the training set, which is exactly the regime de novo drug design needs.
- The semi-autoregressive mask can be applied at sampling time to an already trained ChemBFN and still improve out-of-distribution metrics, meaning the benefit does not require retraining from scratch.
- The combination of the reinforcement-learning term and the ODE-like solver maintains high validity at only 10 sampling steps, so large-scale virtual screening no longer requires thousands of generative steps.
- Switching the molecular representation from SMILES to SELFIES, which is designed to be robust, multiplies the novel hit ratio several-fold while keeping the top docking scores better than all compared baselines.
- The same semi-autoregressive recipe transfers from small molecules to protein sequences, where generated proteins reach objective values above the training range while retaining naturalness comparable to natural proteins.
Reading between the lines
- An editorial caution: the abstract promises a theoretical analysis of out-of-distribution generation, but the body contains no theorem or formal argument; the case for the causal mask rests on the empirical observation of attention sparsity, so the theory is really a conjecture about the stability of that sparsity.
- A natural testable extension is to measure attention-weight mass outside the diagonal while classifier-free guidance strength is increased; if far-from-diagonal mass grows, the hard mask would begin deleting information the model needs, and the OOD gain should shrink or invert.
- Another extension is to apply the same causal-mask strategy to other discrete generative models, such as masked language models or discrete diffusion models, and compare OOD metrics; the mechanism is not obviously specific to Bayesian flow networks, though the paper does not claim this.
- Because the protein-sequence experiment uses only 64 generated sequences, the transfer claim is suggestive rather than established; a larger protein benchmark would tell whether the strategy generalizes beyond small molecules.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes ChemBFN, a Bayesian-flow-network model for molecule generation, and claims it is a natural out-of-distribution sampler that can generate molecules with better docking scores and higher hit ratios than state-of-the-art baselines. The authors introduce an auxiliary reinforcement-learning term, an ODE-like sampling scheme, and a semi-autoregressive strategy based on applying causal masks to the ChemBFN transformer. Experiments cover MOSES and GuacaMol validity/diversity benchmarks, ZINC250k docking-score optimization against five proteins, and two protein-sequence generation tasks. The abstract promises a theoretical analysis of out-of-distribution generation with the semi-autoregressive approach, and the conclusion states that ChemBFN is 'naturally a controllable out-of-distribution sampler.'
Significance. If the central claims are correct, the work would be significant: it would show that a Bayesian flow network can extrapolate beyond its training distribution for de novo drug design, with notably lower docking scores and higher hit ratios than prior diffusion or RL methods, and with sampling accelerated to roughly ten steps. The paper also ships publicly available code and reports standard deviations over five runs for the small-molecule tasks, which strengthens reproducibility. However, the main OOD claim is currently not backed by a direct comparison to the training-set property distribution, and the promised theoretical analysis is absent; these gaps are load-bearing and need to be addressed before the significance can be assessed.
major comments (4)
- [Datasets and Benchmarks, Eq. (2) and Tables 3-4] The central claim that ChemBFN is an out-of-distribution sampler is not established because the comparison to the training distribution is missing. Eq. (2) defines fDS as the median docking score of the training set, so any molecule passing the filter merely beats the median. Tables 4, 6, and 8 report 'novel top 5% docking score' on the molecules passing this filter, but the paper never reports the training-set distribution of docking scores (minimum, percentiles, or top-5% cutoff) for any of the five proteins. A top-5% score of -12.9 kcal/mol for PARP1 could lie inside the training distribution if the training set contains molecules with more negative scores. The UMAP and FCD analyses in Figures 3-5 measure distances in a ChemNet embedding, which is not equivalent to property-based out-of-distribution behavior. I ask the authors to report, per protein, the training distribution of docking scores and to compare generated scores against a training-set top-5% or tail-quantile cutoff; the OOD label should be reserved for molecules that exceed that threshold.
- [Abstract and Semi-autoregressive Training and Sampling] The abstract promises 'a theoretical analysis of out-of-distribution generation in ChemBFN with semi-autoregressive approach,' but the body contains no theorem, proposition, or formal argument. The section 'Semi-autoregressive Training and Sampling' supports the use of causal masks only with the empirical observation that attention weights far from the diagonal are 'extremely close to zero'; no threshold, no stability analysis across datasets or under guidance, and no argument that a hard mask preserves the learned conditional distribution are given. As the semi-autoregressive strategy is presented as a main contribution that lets the model 'surpass the state-of-the-art models,' this load-bearing justification needs either a rigorous statement of the claimed property or removal of the theoretical-analysis claim.
- [Methods, Algorithm 1] The ODE-like update as written is inconsistent with the formula quoted in the text. The text states z_i = z_{i-1} + alpha_i (K e(theta_hat_{i-1}; t_{i-1}) - 1) + sqrt(K alpha_i) epsilon, but Algorithm 1 assigns z <- beta(s) (K e(theta_hat; t) - 1) + sqrt(K beta(s)) tau epsilon, overwriting z at each step instead of accumulating. With the overwriting update, the sampler's z depends only on the last network evaluation, so it cannot implement the sequential refinement that the paper claims; the reported 10-step sampling results therefore do not correspond to the algorithm described. Please correct the pseudocode to match the text (or vice versa) and confirm that the experimental implementation accumulates the updates.
- [Conditional Generation of Protein Sequences] The protein OOD claim is supported by only 32x2 generated sequences per objective with no standard deviations across repeated runs, and Figure 8 shows point clouds without a quantitative comparison to the training distribution beyond 'higher objective values.' In addition, the text states the model determines the relationship 'unsupervised,' yet the model is trained with scalar objective labels and guided sampling; that description is contradictory. Please either provide repeated-run statistics and a formal comparison to training-set extrema or quantiles, or soften the claim.
minor comments (4)
- [Introduction, first paragraph] The phrase 'the models were pool at generating' should read 'poor at generating,' and the discussion near Tables 3-4 refers to 'SOAT models' instead of 'SOTA models.'
- [Methods, Eq. (1)] The notation e^{(k)}(theta_hat; t) is not defined; as written, the RL term appears to maximize the probability of the argmax token for valid molecules rather than the expected reward of a policy gradient, so the connection to REINFORCE should be stated more precisely.
- [Fast Sampling, Table 2] The temperature tau is a free parameter that strongly controls the validity-diversity trade-off (e.g., tau < 0.01 yields validity >= 99.5% but uniqueness < 60%); the sensitivity of the downstream docking metrics to tau should be reported or justified, since only tau = 0.5 and tau = 0.05 are mentioned without a systematic scan.
- [Computational Details] The unconditional rate and guidance strength are given, but the classifier-free guidance implementation for conditioning on continuous scalar values (QED, SA, DS, protein objectives) is not described; a few sentences on how the conditioning vector is injected would improve reproducibility.
Circularity Check
No central circularity: OOD docking results are externally evaluated; only the auxiliary RL validity gain is partly by construction.
-
self definitional
[Methods, Eq. (1); Results, Table 2 and surrounding text]
"To increase valid ratio of generated molecules, a naive strategy is to raise the possibilities that at any time step t the output distributions correspond to valid molecules. Therefore, we set the criterion c as c := {e(hat theta; t) corresponds to a valid molecule }. ... We observed that (1) RL strategy marginally improved the validity and diversity (uniqueness and novelty) of generated samples."
The auxiliary RL loss is defined with c equal to molecular validity, so the model is explicitly rewarded whenever the predicted output distribution corresponds to a valid molecule. Reporting that adding RL improves validity is then a direct consequence of the training objective rather than an independent empirical finding. This is a minor by-construction result, confined to the validity/diversity side of the paper; the central OOD docking claims are measured by external QuickVina 2 scores and are not fitted by this loss.
full rationale
The paper's central claim is that ChemBFN can generate out-of-distribution molecules with better docking scores than SOTA baselines. Those results are evaluated with QuickVina 2 on molecules that pass external filters, and no model parameter is fitted to the reported top-5% docking scores, so the headline comparison is not circular. The semi-autoregressive strategy is justified by an empirical observation of near-zero off-diagonal attention entries and then tested across four train/sample mask combinations; this is an experimental design choice, not a derivation that reduces to its own conclusion. The main circularity-adjacent issue is the auxiliary RL term in Eq. (1), whose reward criterion is explicitly molecular validity; the resulting validity improvement in Table 2 is therefore partially built into the loss rather than predicted independently. That issue is minor and does not affect the docking-score comparisons. The abstract promises a theoretical analysis of OOD generation, but no theorem or formal argument appears in the body; this is missing support, not circularity, and it does not raise the circularity score. Likewise, the absence of training-set docking-score percentiles weakens the OOD interpretation, but that is an evidential limitation rather than a self-referential derivation. Overall, the paper is largely self-contained against external benchmarks, with only a minor by-construction component in the RL validity result.
Assumptions & free parameters
free parameters (5)
- eta (RL scaling constant) =
0.01
- tau (temperature in ODE-like sampling) =
0.5 for MOSES, 0.05 for GuacaMol
- guidance strength =
0.5 for small molecules, 1.0 for proteins
- unconditional rate =
0.2
- LoRA rank and alpha =
rank=4, alpha=1
assumptions (5)
- standard math BFN training and sampling equations of Graves et al. are correct and applicable to SMILES/SELFIES tokens.
- standard math The ODE-like sampling update of Xue et al., Atkinson et al., and Qu et al. is a valid approximation of the BFN generative process.
- domain assumption The ZINC250k OOD filters (QED>0.5, SA<5, DS below training median, SNN<0.4) and QuickVina 2 docking scores are meaningful proxies for drug-likeness and out-of-distribution quality.
- ad hoc to paper Attention weights far from the diagonal in a trained ChemBFN are negligible, so a hard causal mask can be applied without retraining.
- ad hoc to paper The amino-acid-wise tokenizer for proteins is a faithful extension of ChemBFN.
Cite this review
Pith. "Pith review of Sampling Out-of-Distribution Chemical Spaces via Bayesian Flow." pith.science (2026). https://pith.science/paper/Z6Y34652
@misc{pith2026241211439,
author = {Pith},
title = {Pith review of: Sampling Out-of-Distribution Chemical Spaces via Bayesian Flow},
year = {2026},
howpublished = {\url{https://pith.science/paper/Z6Y34652}},
note = {Machine review of arXiv:2412.11439}
}
read the original abstract
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.
Figures
Figures from the paper (6 more)
Reference graph
Works this paper leans on
-
[1]
(1) Bohacek, R. S.; McMartin, C.; Guida, W. C. The art and practice of structure-based drug design: a molecular modeling perspective. Medicinal research reviews 1996, 16, 3–50. (2) Ertl, P. Cheminformatics analysis of organic substituents: identification of the most common substituents, calculation of substituent properties, and automatic identifica- tion...
arXiv 1996
-
[2]
(50) Gruver, N.; Stanton, S.; Frey, N.; Rudner, T
Bioinformatics 2015, 31, 2214–2216. (50) Gruver, N.; Stanton, S.; Frey, N.; Rudner, T. G. J.; Hotzel, I.; Lafrance-Vanasse, J.; Ra- jpal, A.; Cho, K.; Wilson, A. G. Protein design with guided discrete diffusion. Proceed- ings of the 37th International Conference on Neural Information Processing Systems. Red Hook, NY, USA,
work page 2015
-
[5]
Journal of chemical information and computer sciences 1988, 28, 31–36
Introduction to methodology and encoding rules. Journal of chemical information and computer sciences 1988, 28, 31–36. (37) Krenn, M.; H¨ ase, F.; Nigam, A.; Friederich, P.; Aspuru-Guzik, A. Self-referencing em- bedded strings (SELFIES): A 100% robust molecular string representation. Machine Learning: Science and Technology 2020, 1, 045024. (38) Williams,...
arXiv 1988
-
[7]
UMAP: Uniform Manifold Approximation and Pro- jection for Dimension Reduction
24 (51) McInnes, L.; Healy, J.; Melville, J. UMAP: Uniform Manifold Approximation and Pro- jection for Dimension Reduction. 2020; https://arxiv.org/abs/1802.03426. (52) Ho, J.; Salimans, T. Classifier-Free Diffusion Guidance. 2022; https://arxiv.org/ abs/2207.12598. (53) Olivecrona, M.; Blaschke, T.; Engkvist, O.; Chen, H. Molecular de-novo design through...
arXiv 2020
-
[12]
(63) Cock, P. J. A.; Antao, T.; Chang, J. T.; Chapman, B. A.; Cox, C. J.; Dalke, A.; Friedberg, I.; Hamelryck, T.; Kauff, F.; Wilczynski, B.; de Hoon, M. J. L. Biopython: freely available Python tools for computational molecular biology and bioinformatics. Bioinformatics 2009, 25, 1422–1423. (64) Lin, Z.; Akin, H.; Rao, R.; Hie, B.; Zhu, Z.; Lu, W.; Smeta...
work page 2009
-
[13]
(65) Shrake, A.; Rupley, J. A. Environment and exposure to solvent of protein atoms. Lysozyme and insulin. Journal of molecular biology 1973, 79, 351–371. 26 TOC Graphic Property 1 Property 2 Property 3 BNF GeneratingTraining 27
work page 1973
-
[31]
(12) Zhung, W.; Kim, H.; Kim, W
(11) Schwalbe-Koda, D.; G´ omez-Bombarelli, R.Machine Learning Meets Quantum Physics ; Springer International Publishing, 2020; p 445–467. (12) Zhung, W.; Kim, H.; Kim, W. Y. 3D molecular generative framework for interaction- guided drug design. Nature Communications 2024, 15,
work page 2020
-
[48]
(54) Jeon, W.; Kim, D. Autonomous molecule generation using reinforcement learning and docking to develop potential novel inhibitors. Scientific reports 2020, 10, 22104. (55) Jin, W.; Barzilay, D.; Jaakkola, T. Hierarchical Generation of Molecular Graphs us- ing Structural Motifs. Proceedings of the 37th International Conference on Machine Learning. 2020;...
arXiv 2020
Show all 12 references
-
[2022]
Fixing Weight Decay Regularization in Adam
25 (61) Loshchilov, I.; Hutter, F. Fixing Weight Decay Regularization in Adam. 2019; https: //arxiv.org/abs/1711.05101. (62) RDKit: Open-source cheminformatics. https://www.rdkit.org, Accessed: 2024-11-
2019 arXiv
-
[2024]
MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space
22 (34) Qu, Y.; Qiu, K.; Song, Y.; Gong, J.; Han, J.; Zheng, M.; Zhou, H.; Ma, W.- Y. MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space. 2024; https://arxiv.org/abs/2404.12141. (35) Tao, N.; Abe, M. Bayesian Flow Network Framework for Chemistry Tasks. Journal...
2024 arXiv
-
[2688]
P.; Murcko, M
(13) Walters, W. P.; Murcko, M. Assessing the impact of generative AI on medicinal chem- istry. Nature biotechnology 2020, 38, 143–145. (14) Klarner, L.; Rudner, T. G. J.; Morris, G. M.; Deane, C. M.; Teh, Y. W. Context- Guided Diffusion for Out-of-Distribution Molecular and P...
2020 arXiv
-
[4348]
(59) Sterling, T.; Irwin, J. J. ZINC 15–ligand discovery for everyone. Journal of chemical information and modeling 2015, 55, 2324–2337. (60) Hu, E. J.; Shen, Y.; Wallis, P.; Allen-Zhu, Z.; Li, Y.; Wang, S.; Wang, L.; Chen, W. LoRA: Low-Rank Adaptation of Large Language Models...
2015
Reviewed August 11, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.