SlowGuess/NeurIPS_2023_OCR huggingface dataset:
While deep probabilistic models attained competitive results in terms of RMSD, we found their uncertainty estimates to be consistently less reliable than those of GP-based models with discrete string kernels or shallow continuous kernels on hand-crafted features (e.g. fragprints), limiting their suitability for Bayesian optimisation and active learning. Our results suggest that for small to mid-sized molecular datasets the Tanimoto kernel combined with fragment representations in particular is a very compelling option, with good accuracy, calibration, and runtime across all tasks.
5.2 Bayesian Optimisation
Building on these results, we employed the two best-performing kernels, namely the Tanimoto-fragment kernel and the SMILES string kernel, to undertake Bayesian optimization BO over the photoswitch and ESOL datasets. BO is run for 20 iterations of sequential candidate selection (EI acquisition) where candidates are drawn from  $95\%$  of the dataset. The models are initialised with  $5\%$  of the dataset. In the case of the photoswitch dataset, this corresponds to just 19 molecules. The results are provided in Figure 4. In this ultra-low data setting—common to many areas of synthetic chemistry [3]—both models significantly outperform the random search baseline, highlighting the real-world use-case for such models in supporting human chemists to prioritise candidates for synthesis.
Furthermore, one may observe that BO performance is tightly coupled to regression and UQ performance. In the case of the photoswitch dataset, the better-performing Tanimoto model on regression and UQ also achieves better BO performance, while on the ESOL dataset, the string kernel performs best. Additionally, we run BO on the Buchwald-Hartwig dataset using the Tanimoto kernel for the bit-vector representations DRFP and OHE, and the RBF kernel for RXNFP. All three representations perform similarly and outperform the random search.
Finally, we also investigate the performance of GP-based models for preferential BO—a setting in which the acquisition strategy only requires rank-based preferences of candidates, as opposed to their absolute objective function values [65, 66, 67, 68]. We use a Tanimoto-fragment kernel GP model to perform molecular Bayesian optimization on the photoswitch dataset using binary preference data alone and present the full results in Appendix E.4.
Table 2: An overview of existing open-source GP, BO and molecular machine learning libraries. With GAUCHE, we provide a modular, robust and easy-to-use framework that combines state-of-the-art probabilistic modelling and black-box optimization techniques with bespoke molecular representations and kernels to make them more easily accessible to the broader scientific community.
Library	Gaussian Processes	Bayesian Optimisation	Molecular Representations	Chemistry Tutorials	Graph Kernels	Bit Vector Kernels	String Kernels
GPyTorch [26]	✓	X	X	X	X	X	X
GPflow [69, 70]	✓	X	X	X	X	X	X
BoTorch [27]	✓	✓	X	X	X	X	X
DeepChem [71]	X	X	✓	✓	X	X	X
GraKel [59]	X	X	X	X	✓	X	X
FlowMO [72]	✓	X	✓	✓	X	✓	✓
GAUCHE (ours)	✓	✓	✓	✓	✓	✓	✓
General-purpose GP and BO libraries do not cater for molecular representations. Likewise, general-purpose molecular machine learning libraries do not consider GPS and BO. Here, we review existing libraries, highlighting the niche GAUCHE fills in bridging the GP and molecular machine learning communities. The closest work to ours is FlowMO [72], which introduces a basic molecular GP library in the GPflow framework. In this project, we extend the scope of the library to a broader class of molecular representations (graphs), problem settings (BO) and applications (reaction optimisation and protein engineering). An overview of how GAUCHE fits into the existing open-source GP, BO and molecular machine learning stack is presented in Table 2.
Gaussian Process Libraries GP libraries include GPy (Python) [73], GPflow (TensorFlow) [69, 70], GPyTorch (PyTorch) [26] and GPJax (Jax) [74] while examples of recent BO libraries include BoTorch (PyTorch) [27], Dragonfly (Python) [75], HEBO (PyTorch) [76] and Trieste (Tensorflow) [77]. The aforementioned libraries do not explicitly support molecular representations. Extending them to cover molecular representations, however, requires implementations of bespoke GP kernels for bit vector, string and graph inputs together with modifications to BO schemes to consider acquisition function evaluations over a discrete set of held-out molecules, a setting commonly encountered in virtual screening [78, 79].
Molecular Machine Learning Libraries Molecular machine learning libraries include DeepChem [71], DGL-LifeSci [80] and TorchDrug [81]. DeepChem features a broad range of model implementations and tasks, while DGL-LifeSci focuses on graph neural networks. TorchDrug caters for applications including property prediction, representation learning, retrosynthesis, biomedical knowledge graph reasoning and molecule generation. However, none of the aforementioned libraries includes GP implementations. In terms of atomistic systems, DScribe [82] features, amongst other methods, the Smooth Overlap of Atomic Positions (SOAP) representation [83], which is typically used in conjunction with a GP model to learn atomistic properties. Automatic Selection And Prediction (ASAP) [84] also principally focuses on atomistic properties as well as dimensionality reduction and visualisation techniques for materials and molecules. Lastly, the Graphein library focuses on graph representations of proteins [49].
Graph Kernel Libraries Graph kernel libraries include GraKel [59], graphkit-learn [85], graphkernels [86], graph-kernels [87], pykernels (https://github.com/gmum/pykernels) and ChemoKernel [88]. The aforementioned libraries focus on CPU implementations in Python. Extending graph kernel computation to GPUs has been noted as an important direction for future research [89]. In our work, we build on the GraKel library to construct GPyTorch-based GPs that can be trained on non-tensorial, graph-structured inputs. It is worth noting that GAUCHE extends the applicability of GPU-enabled GPs to general graph-structured inputs beyond just molecules and proteins.
Molecular Bayesian Optimisation BO over molecular space can be divided into two classes. In the first class, molecules are encoded into the latent space of a variational autoencoder (VAE) [4]. BO is then performed over the continuous latent space and queried molecules are decoded back to the original space. Much work on VAE-BO has focussed on improving the synergy between the surrogate model and the VAE [90, 5, 91, 92, 93, 94, 95, 96]. One of the defining characteristics of VAE-BO is that it enables the generation of new molecular structures. In the second class of methods, BO is performed directly over the original discrete space of molecules. In this setting it is not possible to generate new structures and so a candidate set of queryable molecules is defined. The inability to generate new structures however, is not a bottleneck to molecule discovery as the principle concern is how best to explore existing candidate sets. These candidate sets are also known as molecular libraries in the virtual screening literature [97]. To date, there has been little work on BO directly over discrete molecular spaces. In [56], the authors use a string kernel GP trained on SMILES to perform BO to select from a candidate set of molecules. In [98], an optimal transport kernel GP is used for BO over molecular graphs. In [99] a surrogate based on the Nadarya-Watson estimator is defined such that the kernel density estimates are inferred using BNNs. The model is then trained on molecular descriptors. Lastly, in [100] and [101] a BNN and a sparse GP respectively are trained on fingerprint representations of molecules. In the case of the sparse GP the authors select an ArcCosine kernel. It is a longstanding aim of the GAUCHE Project to compare the efficacy of VAE-BO against vanilla BO on real-world molecule discovery tasks.
Chemical Reaction Optimisation Chemical reactions describe how reactants transform into products. Reagents (catalysts, solvents, and additives) and reaction conditions heavily impact the outcome of chemical reactions. Typically the objective is to maximise the reaction yield (the amount of product compared to the theoretical maximum) [41], in asymmetric synthesis, where the reactions could result in different enantiomers, to maximise the enantiomeric excess [102], or to minimise the E-factor, which is the ratio between waste materials and the desired product [103]. A diverse set of studies have evaluated the optimisation of chemical reactions in single and multi-objective settings [103, 104]. [105] and [106] benchmarked reaction optimisation algorithms in low-dimensional settings including reaction conditions, such as time, temperature, and concentrations. [6] suggested BO as a general tool for chemical reaction optimisation and benchmarked their approach against human experts. [107] compared the yield prediction performance of different kernels and [108] the impact of various molecular representations. In all reaction optimisation studies above, the representations of the different categories of reactants and reagents are concatenated to generate the reaction input vector, which could lead to limitations if another type of reagent is suddenly considered. Moreover, most studies concluded that simple one-hot encodings (OHE) perform at least on par with more elaborate molecular representations in the low-data regime [6, 108, 109]. In GAUCHE, we introduce reaction fingerprint kernels, based on existing reaction fingerprints [46, 45] and work independently of the number of reactant and reagent categories.
One potential limitation of GAUCHE is the focus on core implementations that are likely to remain robust across as many applied problems as possible which will enable the library to have the most impact. As such, there is less of a focus on bespoke GP implementations for more targeted problems [110, 111, 112, 113]. Nonetheless, we hope that GAUCHE will function as an active development platform for such implementations.
We have introduced GAUCHE, a library for Gaussian Processes in Chemistry, with the aim of providing a user-friendly and robust library of state-of-the-art uncertainty quantification and Bayesian optimisation tools that may hopefully be deployed for screening in laboratory settings. Our aim is to maintain a lean, well-tested and up-to-date codebase and invite community-driven contributions principally as pull requests in the form of notebooks that reflect the needs and considerations that researchers come across in practice. In this fashion, we may support more advanced features without bloating the codebase and increasing maintenance requirements.
BR and PS acknowledge support from the NCCR Catalysis (grant number 180544), a National Centre of Competence in Research funded by the Swiss National Science Foundation. ARJ is funded by a Biotechnology and Biological Sciences Research Council (BBSRC) DTP studentship (BB/M011194/1). GT acknowledges the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), and the Vector Institute. FS-K is a postdoctoral fellow in the Eric and Wendy Schmidt AI in Science Postdoctoral Fellowship Program, a program by Schmidt Futures. AA-G acknowledges the generous support of Anders G. Frøseth, the Canadian Institute for Advanced Research (CIFAR), and the Canada 150 Research Chair program. LK acknowledges support from the University of Oxford's Clarendon Fund.
R-RG declares that this work was done exclusively at the University of Cambridge. Code/data was not created/accessed by Meta.
[1] Ying Zhang and Chen Ling. A strategy to apply machine learning to small datasets in materials science. Npj Computational Materials, 2018.
[2] Ryan-Rhys Griffiths, Jiachen Jiang, Douglas JK Buisson, Dan Wilkins, Luigi C Gallo, Adam Ingram, Dirk Grupe, Erin Kara, Michael L Parker, William Alston, et al. Modeling the multiwavelength variability of Mrk 335 using Gaussian processes. The Astrophysical Journal, 914(2):144, 2021.
[3] Ryan-Rhys Griffiths, Jake L Greenfield, Aditya R Thawani, Arian R Jamasb, Henry B Moss, Anthony Bourached, Penelope Jones, William McCorkindale, Alexander A Aldrick, Matthew J Fuchter, et al. Data-driven discovery of molecular photoswitches with multioutput Gaussian processes. Chemical Science, 2022.
[4] Rafael Gómez-Bombarelli, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamin Sánchez-Lengeling, Dennis Sheberla, Jorge Aguilera-Iparraguirre, Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS Central Science, 2018.
[5] Ryan-Rhys Griffiths and José Miguel Hernández-Lobato. Constrained Bayesian optimization for automatic chemical design using variational autoencoders. Chemical Science, 2020.
[6] Benjamin J Shields, Jason Stevens, Jun Li, Marvin Parasram, Farhan Damani, Jesus I Martinez Alvarado, Jacob M Janey, Ryan P Adams, and Abigail G Doyle. Bayesian reaction optimization as a tool for chemical synthesis. Nature, 2021.
[7] Yuanqi Du, Tianfan Fu, Jimeng Sun, and Shengchao Liu. Molgensurvey: A systematic survey in machine learning models for molecule design. arXiv preprint arXiv:2203.14500, 2022.
[8] Jasper Snoek, Oren Rippel, Kevin Swersky, Ryan Kiros, Nadathur Satish, Narayanan Sundaram, Mostofa Patwary, Mr Prabhat, and Ryan Adams. Scalable Bayesian optimization using deep neural networks. In International Conference on Machine Learning, 2015.
[9] Jost Tobias Springenberg, Aaron Klein, Stefan Falkner, and Frank Hutter. Bayesian optimization with robust Bayesian neural networks. Advances in Neural Information Processing Systems, 2016.
[10] Yao Zhang et al. Bayesian semi-supervised learning for uncertainty-calibrated prediction of molecular properties and active learning. Chemical Science, 2019.
[11] Kevin Maik Jablonka, Giriprasad Melpatti Jothiappan, Shefang Wang, Berend Smit, and Brian Yoo. Bias free multiobjective active learning for materials design and discovery. Nature Communications, 2021.
[12] Seongok Ryu, Yongchan Kwon, and Woo Youn Kim. A bayesian graph convolutional network for reliable prediction of molecular properties with uncertainty quantification. Chemical Science, 2019.
[138] Stephen A Forbes. Method of determining and measuring the associative relations of species. Science, 61(1585):518-524, 1925.
[139] Robert R Sokal, Peter HA Sneath, et al. Principles of numerical taxonomy. *Principles of numerical taxonomy*, 1963.
Supplementary to:
A Uncertainty Quantification Experiments
In Table 3 we present further uncertainty quantification results using the quantile coverage error (QCE) metric. Numerical errors were encountered with the WL kernel on the large lipophilicity dataset and the corresponding entry is left blank.
Table 3: UQ benchmark. QCE values (↓) for 80/20 train/test split across 20 random trials.
GP Model	Dataset
Kernel	Representation	Photoswitch	ESOL	FreeSolv	Lipophilicity
Tanimoto	fragprints	0.019 ± 0.003	0.023 ± 0.002	0.023 ± 0.002	0.006 ± 0.002
fingerprints	0.023 ± 0.003	0.022 ± 0.002	0.018 ± 0.003	0.006 ± 0.001
fragments	0.025 ± 0.005	0.012 ± 0.002	0.014 ± 0.002	0.009 ± 0.002
Scalar Product	fragprints	0.033 ± 0.006	0.010 ± 0.002	0.017 ± 0.003	0.010 ± 0.001
fingerprints	0.036 ± 0.006	0.014 ± 0.002	0.016 ± 0.002	0.009 ± 0.001
fragments	0.027 ± 0.004	0.012 ± 0.003	0.021 ± 0.003	0.010 ± 0.001
String	SMILES	0.024 ± 0.003	0.016 ± 0.002	0.019 ± 0.003	0.005 ± 0.001
WL Kernel (GraKel)	graph	0.025 ± 0.007	0.011 ± 0.004	0.019 ± 0.009	-
B Chemical Reaction Yield Prediction Experiments
Further regression and uncertainty quantification experiments are presented in Table 4. The differential reaction fingerprint in conjunction with the Tanimoto kernel is the best-performing reaction representation.
Table 4: Chemical reaction regression benchmark. 80/20 train/test split across 20 random trials.
GP Model	Buchwald-Hartwig
Kernel	Representation	RMSE↓	R2score ↑	QCE ↓
Tanimoto	OHE	7.94 ± 0.05	0.91 ± 0.001	0.011 ± 0.001
DRFP	6.48 ± 0.45	0.94 ± 0.015	0.027 ± 0.002
Scalar Product	OHE	15.23 ± 0.052	0.69 ± 0.002	0.008 ± 0.001
DRFP	14.63 ± 0.050	0.71 ± 0.002	0.010 ± 0.001
RBF	RXNFP	10.79 ± 0.049	0.84 ± 0.001	0.024 ± 0.001
Suzuki-Miyaura
Tanimoto	OHE	11.18 ± 0.036	0.83 ± 0.001	0.007 ± 0.001
DRFP	11.46 ± 0.038	0.83 ± 0.001	0.019 ± 0.000
Scalar Product	OHE	19.91 ± 0.042	0.47 ± 0.003	0.012 ± 0.001
DRFP	19.66 ± 0.042	0.52 ± 0.003	0.014 ± 0.001
RBF	RXNFP	13.83 ± 0.048	0.75 ± 0.002	0.007 ± 0.001
C Scaffold Split Experiments
To investigate how GP models behave on more challenging train/test splits, we have re-run parts of our experimental evaluation with 80-20 Bemis-Murcko [114] scaffold splits instead of random splits. As only the lipophilicity dataset exhibits sufficient scaffold diversity to perform this analysis (the skewness of the scaffold distribution in the others makes an 80-20 split impossible), the results in Table 5 focus on the predictive accuracy (RMSE) and calibration (NLPD) of GP models in this setting.
While this more challenging evaluation setup leads to slightly higher RMSEs and NLPDs, we note that one can observe the same trends as with random splits: Tanimoto-based GPs generally outperform Scalar Product ones, while string kernel-based GPs are better than both.
Table 5: Regression and uncertainty quantification experiments on scaffold splits.
Kernel	Representation	RMSE (↓)	NLPD (↓)
Tanimoto	Fragments	0.86±0.01	1.02±0.04
Fingerprints	0.88±0.01	1.12±0.04
Fragments	0.89±0.01	2.10±0.13
Scalar Product	Fragments	0.89±0.01	1.75±0.08
Fingerprints	0.95±0.01	1.99±0.09
Fragments	1.00±0.01	-
String	SMILES	0.82±0.01	1.08±0.04
D Deep Probabilistic Models
There is a growing body of work applying deep learning to molecular property prediction [115]. Therefore in addition to evaluating GPs with varying shallow kernel functions, we repeat the regression experiments with a range of deep Bayesian models, varying both the network architecture and the Bayesian inference procedure. We evaluate the following models:
- FC-BNN + VI is a fully connected neural network with a single variational inference (VI) Bayesian layer with 100 nodes, followed by the rectified linear unit (ReLU) activation, and a final output layer. Trained with early stopping.
- GNN-BNN + VI utilises the same network and graph features as used for the graph embeddings [116] followed by a final VI Bayesian layer. Trained with early stopping.
- CNN DKL GP is the same approach and architecture used by [96] to predict molecular properties for Bayesian optimisation. Using SELFFIES representations, a 1D CNN encoder is shared and trained jointly through a generative masked language model (MLM) head [47] and a discriminative deep kernel GP head [117].
- CNN Ensemble is a deep ensemble of 1D CNN networks, also implemented by [96], where each ensemble component uses the SELFIES molecule representation and is trained independently to minimize the MSE loss. Deep ensembles have been shown to provide high-fidelity approximations of Bayesian model averages relative to alternative approaches such as Laplace approximation or VI [118].
E Further Experiments
E.1 Protein Fitness Prediction
We consider the task of protein fitness prediction where the fitness function (target label) takes the form of the melting point in degrees Celsius. We collate a dataset of 151 PETase protein sequences from values reported in [119, 120, 121]. PETases, recently discovered in 2016 [122], are a class of esterase enzymes which, via hydrolysis, catalyse the breakdown of polyethylene terephthalate (PET) plastic to monomeric mono-2-hydroxyethyl terephthalate (MHET). While PET plastics can require hundreds of years to degrade naturally, PETases are capable of degrading PET in a matter of days but their melting point is a key property of interest for deployed applications. Each sequence consists of an amino acid chain of length 290. We use the 'bag of amino acids' representation with a max n-gram value of 5 which gives rise to a count vector where each component represents the number of a given n-gram contained in the sequence. A max n-gram value of 5 means that all n-grams up to and including length 5 are included in the featurisation. We subsequently train a Tanimoto kernel GP
on the featurisation. The results are provided in Table 6 and indicate that the GP model obtains low generalisation error on the melting point prediction task.
Table 6: PETase melting point prediction experiment. 80/20 train/test split across 20 random trials.
GP Model	PETase Test Set
Kernel	Representation	RMSE ↓	MAE ↓	R2 score ↑
Tanimoto	Bag of Amino Acids	3.68 ± 0.13	2.54 ± 0.10	0.81 ± 0.02
E.2 Molecular Preference Learning
In many optimisation problems, while it may be challenging to define a mathematical utility function, human feedback on pairwise comparisons can be leveraged to learn a latent utility function. In the machine learning literature, this observation has inspired work on preference learning [123, 124, 125, 126] where the goal is to learn a utility function  $g(\mathbf{x})$  using binary preference data  $r(\mathbf{x}_1,\mathbf{x}_2)$  obtained from human feedback on inputs  $\mathbf{x}$ . In the case of molecular design, there are many situations in which it is difficult to articulate a utility function [127, 128, 129] which has motivated recent work on human-in-the-loop preference learning [130, 131]. In this setting, pairwise preferences  $r(\mathbf{x}_1,\mathbf{x}_2)$  are collected from a human chemist by presenting with many choices over pairs of molecules. In this section, we highlight a use-case for GAUCHE in GP-based preference learning [132]. Using the photoswitch dataset as a case study, we simulate noiseless pairwise comparisons from a human chemist on the E isomer  $\pi - \pi^*$  transition wavelength (latent utility function). We train a Tanimoto kernel GP on the fingerprints representation of the molecules, using a probit likelihood and the Laplace approximation [132]. The results are provided in Table 7. We report the Kendall-Tau rank correlation as the preference model learns ordinal rankings in place of absolute values of the transition wavelength. The results indicate that it is possible to learn accurate models of molecular properties through binary feedback alone.
Table 7: Photoswitch preference learning experiment. 80/20 train/test split across 20 random trials with 5000 training set pairwise comparisons. Kendall-Tau rank correlation (K-T) values reported.
GP Model	Photoswitch Dataset
Kernel	Representation	K-T Train ↑	K-T Test ↑
Tanimoto	Fragments	0.95 ± 0.001	0.79 ± 0.01
E.3 Ablation over the ECFP Radius Parameter
In addition to the experimental results presented in Section 5, we have performed an ablation study over the radius parameter of extended-connectivity fingerprints (ECFPs). Specifically, we trained Tanimoto-kernel GPs on the Photoswitch dataset using a series of five increasing radius parameters. In Table 8, we report the mean and standard error of the RMSE and NLPD over 50 different 80-20 train-test-splits.
Table 8: RMSE and NLPD of Tanimoto-kernel GPs trained on the Photoswitch dataset using a series of five increasing ECFP radius parameters.
Fingerprint	RMSE (↓)	NLPD (↓)
ECFP4	22.65±0.55	0.41±0.05
ECFP6	23.50±0.55	0.47±0.05
ECFP8	24.43±0.54	0.52±0.05
ECFP10	25.17±0.54	0.56±0.04
ECFP12	25.70±0.53	0.58±0.04
Intriguingly, these results show a strong negative correlation between the fingerprint radius and predictive performance. We hypothesize that this is caused by the fact that an expanding feature
space leads to lower and less informative Tanimoto similarity scores, making it more difficult to train generalisable models.
E.4 Preferential Bayesian Optimisation