A Bayesian Flow Network Framework for Chemistry Tasks

TL;DR

This paper presents ChemBFN, a Bayesian flow network-based language model for chemistry tasks that improves molecule generation quality, supports conditional generation, and can be fine-tuned for various predictive tasks.

Contribution

Introduces ChemBFN, a novel Bayesian flow network model for chemistry that enhances sampling accuracy and enables multi-task fine-tuning within a single framework.

Findings

Improved molecule diversity with fewer sampling steps

Effective conditional generation via classifier-free guidance

State-of-the-art performance in regression and classification tasks

Abstract

In this work, we introduce ChemBFN, a language model that handles chemistry tasks based on Bayesian flow networks working on discrete data. A new accuracy schedule is proposed to improve the sampling quality by significantly reducing the reconstruction loss. We show evidence that our method is appropriate for generating molecules with satisfied diversity even when a smaller number of sampling steps is used. A classifier-free guidance method is adapted for conditional generation. It is also worthwhile to point out that after generative training, our model can be fine-tuned on regression and classification tasks with the state-of-the-art performance, which opens the gate of building all-in-one models in a single module style. Our model has been open sourced at https://github.com/Augus1999/bayesian-flow-network-for-chemistry.