SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Baselines. We compare our methods with four BO baselines: LOL-BO [10], W-LBO [13], TuRBO [20] and LS-BO. LOL-BO proposed the joint loss to close the gap between the latent space and the discrete input space. Additionally, W-LBO weighted data points to emphasize important samples with regard to their objective values. To handle discrete and structured data, we leverage TuRBO in latent space, TuRBO- $L$  through encoder  $q_{\phi}$  of pretrained frozen VAE and reconstruct it by decoder  $p_{\theta}$  to be evaluated by the objective function. For more details about TuRBO- $L$ , see [10]. In our LS-BO approach, we adopt the standard LBO methodology. Specifically, we use a VAE with pre-trained parameters that are frozen. For sampling, the candidates are sampled from a standard normal Gaussian distribution,  $\mathcal{N}(0,I)$ .
3.1 Results on diverse benchmark tasks
Figure 2, 3 represent the graphs that depict the number of oracle calls, i.e., the number of the black-box function evaluations, and the corresponding mean and standard deviation of objective value. The maximum number of oracles is set to 70k with Gucamol benchmarks, 3k with the DRD3 benchmark, and 500 with the arithmetic expression task. As shown in Figure 2, our method outperformed all the baselines in all the tasks. Note that the goal of the arithmetic task in Figure 3a and the DRD3 task Figure 3b is minimizing the objective function, while others aim to maximize it. In the case of Figure 2a and Figure 2d, LS-BO and TuRBO-  $L$  markedly failed to search for the one that has the best score. Especially in Figure 2a, LS-BO and TuRBO-  $L$ , which only search with the fixed latent space, failed to update the initial best point despite searching 70k molecules. It implies that the pretrained VAE cannot generate better molecules unless the latent space is updated.
3.2 Ablation study
Table 1: An ablation of CoBO's components. Table 2: An ablation on trust region recoordi-nation.
Lz	λ(y)	L Lip	Score
✓	✓	✓	0.7921
×	✓	✓	0.7835
×	×	✓	0.7733
×	×	×	0.7504
Recoordination	Score
✓	0.7921
×	0.6844
(a) Median molecules 2 (med2)
(b) Zaleplon MPO (zale)
(c) Perindopril MPO (pdop)
(d) Osimertinib MPO (osmb)
Figure 2: Optimization results with four different tasks on the Guacamol benchmark. The lines and range are the mean and standard deviation of three repetitions with the same parameters.
(a) Arithmetic expression
(b) TDC DRD3
Figure 3: Optimization results with the arithmetic expression and TDC DRD3 benchmark. The lines and range are the mean and standard deviation of three repetitions with the same parameters.
In this section, we evaluate the main components of our model to analyze their contribution. We employ Perindopril MPO (pdop) task for our experiment and note that all scores of ablation studies are recorded when the oracle number is  $20k$  to compare each case in limited oracle numbers as the original BO intended.
We ablate the three main components of our CoBO and report the results in Table 1. As latent space regularization  $\mathcal{L}_{\mathrm{z}}$  comes from Lipschitz regularization and loss weighting schema aims to prioritize penalized input points, we conducted cascading experiments. Notably, we observed that performance decreased as the components were omitted, and the largest performance dropped (0.0229)
(a) Perindopril MPO (pdop)
(b) Arithmetic expression
Figure 4: Effects of  $\mathcal{L}_{\mathrm{align}}$ . The plot depicts the Pearson correlation between the distance of the latent vectors and the distance of the objective values over the Bayesian optimization process. Each line represents training with a different loss for the latent space during the optimization process: one with  $\mathcal{L}_{\mathrm{align}}$  (orange) and the other without  $\mathcal{L}_{\mathrm{align}}$  (blue), where  $\mathcal{L}_{\mathrm{align}} = \mathcal{L}_{\mathrm{Lip}} + \mathcal{L}_{\mathrm{z}}$ . We measure the correlation after every VAE update.
when Lipschitz regularization was not employed. Table 2 reports the effectiveness of trust region recoordination. We observe that applying trust region recoordination makes a considerable difference within a small oracle budget, as the score increases 0.6844 to 0.7921 with recoordination.
3.3 Analysis on proposed regularizations
All the analyses were conducted on the Perindopril MPO (pdop) task from the Guacamol benchmark. We include an additional analysis with the arithmetic data. For convenience, we define the  $\mathcal{L}_{\mathrm{align}}$  as follows:
$$
\mathcal {L} _ {\text {a l i g n}} = \mathcal {L} _ {\mathrm {L i p}} + \mathcal {L} _ {\mathrm {z}}, \tag {12}
$$
which is the loss that aligns the latent space with the black-box objective function.
Effects of  $\mathcal{L}_{\mathrm{align}}$  on correlation. In Section 2.2, we theoretically prove that  $\mathcal{L}_{\mathrm{align}}$  increases the correlation between the distance in the latent space and the distance within the objective function. Here, we demonstrate our theorem with further quantitative analysis to show correlation changes during the Bayesian optimization process. Figure 4 shows Pearson correlation value between the latent space distance  $\| \mathbf{z}_i - \mathbf{z}_j\| _2$  and the objective value distance  $|y_{i} - y_{j}|$ . The blue and orange lines indicate the models with and without our align loss  $\mathcal{L}_{\mathrm{align}}$  in Eq. 12, respectively. We measure the correlation with  $10^{3}$  data point and every  $10^{6}$  pair. The data is selected as the top  $10^{3}$  points with the highest value among  $10^{3}$  data, which are from training data for the VAE model and surrogate model. Over the training process, the Pearson correlation values with our align loss  $\mathcal{L}_{\mathrm{align}}$  were overall higher compared to the baseline. Moreover,  $\mathcal{L}_{\mathrm{align}}$  increases the Pearson correlation value high over 0.7 in Figure 4a which is normally regarded as the high correlation value. This means align loss increases the correlation between the distance of the latent vectors and the distance of the objective values effectively, leading to narrowing the gap between the latent space and the objective function.
Effects of  $\mathcal{L}_{\mathrm{align}}$  on smoothness. To validate that our model encourages the smooth latent space, we qualitatively analyze our model with  $\mathcal{L}_{\mathrm{align}}$  and without  $\mathcal{L}_{\mathrm{align}}$  by visualizing the landscape of the latent space after the training finishes. In Figure 5, we visualize the top- $k$  objective value and the corresponding latent vector in 2D space with the 2D scatter plot and the corresponding 3D plot for better understanding. We reduce the dimension of the latent vector to two dimensions using principal component analysis (PCA). The color means the normalized relative objective score value. The landscape of objective value according to the latent space with  $\mathcal{L}_{\mathrm{align}}$  (right) is smoother than the case without  $\mathcal{L}_{\mathrm{align}}$  (left). It demonstrates our  $\mathcal{L}_{\mathrm{align}}$  loss encourages smoothness of the latent space with respect to the objective function. Note that due to the inherent discreteness of the black-box objective function, it's expected to observe some spaces between the clusters in the 2D plot. Still, this does not detract from our purpose of effectively aligning latent vectors with respect to their discrete objective values.
(a) Without  $\mathcal{L}_{\mathrm{align}}$
(b) With  $\mathcal{L}_{\mathrm{align}}$
Figure 5: Visualizations on the landscape of latent space for  $\mathcal{L}_{\mathrm{align}}$  ablation. The scatter plot and corresponding 3D plot of the latent vectors with objective values. The landscape becomes much smoother with applying  $\mathcal{L}_{\mathrm{align}}$ . A colorbar indicates the normalized objective value, where yellow means higher value and purple means lower value.
4 Related Works
4.1 Latent space Bayesian optimization
Latent space Bayesian optimization [10, 12-14, 19, 21, 29, 30] aims to resolve the issues in optimization over high-dimensional, or structured input space by introducing Bayesian optimization over latent space. As the objective for structured inputs is usually defined over a large, complex space, the challenges can be alleviated by the lower-dimensional and continuous latent space. Variational autoencoders (VAEs) [15] are commonly leveraged to learn the continuous embeddings for the latent space Bayesian optimizers. Some prior works propose novel architectures for decoders [11, 12, 31-33], while others introduce loss functions to improve the surrogate for learning the objective function [13, 14, 18, 29]. Note that while the surrogate model (typically GP) is modeled based on the latent space, it is the input space which the objective value is obtained from. Although this results in an inherent gap in latent space Bayesian optimization, many prior works [12, 14, 19, 21, 29] do not update the generative model for the latent space. LOL-BO [10] seeks to address this gap by adapting the latent space to the GP prior, while [13] suggests periodic weighted retraining to update the latent space.
4.2 Latent space regularization
Several previous studies in latent space have focused on learning the appropriate latent space for their tasks by incorporating additional regularizations or constraints alongside the reconstruction loss. These approaches have been applied in a wide range of areas, including molecule design [34-36], domain adaptation [37, 38], semantic segmentation [39, 40], representation learning [41-43], and reinforcement learning [44]. Notably, the Lipschitz constraint is commonly employed to promote smoothness in diverse optimization problems. For example, [42] introduces Lipschitz regularization in learning implicit neural functions to encourage smooth latent space for 3D shapes while [45] penalizes the data pairs that violate the Lipschitz constraint in adversarial training. Additionally, CoFLO [46] leverages Lipschitz-like regularization in latent space Bayesian optimization. Our method proposes regularization to close the gap inherent in the latent space Bayesian optimization. Concretely, we introduce Lipschitz regularization to increase the correlation between the distance of latent space and the distance of objective value and give more weight to the loss in the promising areas.
In this paper, we addressed the problem of the inherent gap in the latent space Bayesian optimization and proposed Correlated latent space Bayesian Optimization. We introduce Lipschitz regularization which maximizes the correlation between the distance of latent space and the distance of objective value to close the gap between latent space and objective value. Also, we reduced the gap between latent space and input space with a loss weighting scheme, especially in the promising areas. Additionally, by trust region recoordination, we adjust the trust regions according to the updated latent space. Our experiments on various benchmarks with molecule generation and arithmetic fitting tasks demonstrate that our CoBO significantly improves state-of-the-art methods in LBO.
Limitations and broader impacts. Given the contribution of this work to molecular design optimization, careful consideration should be given to its potential impacts on the generation of toxic or harmful substances in the design of new chemicals. We believe that our work primarily has the positive potential of accelerating the chemical and drug development process, setting a new standard in latent space Bayesian optimization.
This work was partly supported by ICT Creative Consilience program (IITP-2023-2020-0-01819) supervised by the IITP, the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (NRF-2023R1A2C2005373), and Samsung Research Funding & Incubation Center of Samsung Electronics under Project Number SRFC-IT1701-51.
[1] Ksenia Korovina, Sailun Xu, Kirthevasan Kandasamy, Willie Neiswanger, Barnabas Poczos, Jeff Schneider, and Eric Xing. Chembo: Bayesian optimization of small organic molecules with synthesizable recommendations. In ICAIS, 2020.
[2] José Miguel Hernández-Lobato, James Requeima, Edward O Pyzer-Knapp, and Alán Aspuru-Guzik. Parallel and distributed thompson sampling for large-scale accelerated exploration of chemical space. In ICML, 2017.
[3] Ryan-Rhys Griffiths and José Miguel Hernández-Lobato. Constrained bayesian optimization for automatic chemical design using variational autoencoders. Chemical science, 2020.
[4] Ke Wang and Alexander W Dowling. Bayesian optimization for chemical products and functional materials. Current Opinion in Chemical Engineering, 2022.
[5] Roberto Calandra, André Seyfarth, Jan Peters, and Marc Peter Deisenroth. Bayesian optimization for learning gaits under uncertainty: An experimental comparison on a dynamic bipedal walker. Annals of Mathematics and Artificial Intelligence, 2016.
[6] Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. In NeurIPS, 2012.
[7] Eric Brochu, Vlad M Cora, and Nando De Freitas. A tutorial on bayesian optimization of expensive cost functions, with application to active user modeling and hierarchical reinforcement learning. Arxiv, 2010.
[8] Peter I Frazier. A tutorial on bayesian optimization. *Arxiv*, 2018.
[9] Bobak Shahriari, Kevin Swersky, Ziyu Wang, Ryan P Adams, and Nando De Freitas. Taking the human out of the loop: A review of bayesian optimization. Proceedings of the IEEE, 2015.
[10] Natalie Maus, Haydn Jones, Juston Moore, Matt J Kusner, John Bradshaw, and Jacob Gardner. Local latent space bayesian optimization over structured inputs. In NeurIPS, 2022.
[11] Matt J Kusner, Brooks Paige, and José Miguel Hernández-Lobato. Grammar variational autoencoder. In ICML, 2017.
[32] Hiroshi Kajino. Molecular hypergraph grammar with its application to molecular optimization. In ICML, 2019.
[33] Hanjun Dai, Yingtao Tian, Bo Dai, Steven Skiena, and Le Song. Syntax-directed variational autoencoder for structured data. In ICLR, 2018.
[34] Egbert Castro, Abhinav Godavarthi, Julian Rubinfien, Kevin Givechian, Dhananjay Bhaskar, and Smita Krishnaswamy. Transformer-based protein generation with regularized latent space optimization. Nature Machine Intelligence, 2022.
[35] Pascal Notin, José Miguel Hernández-Lobato, and Yarin Gal. Improving black-box optimization in vae latent space using decoder uncertainty. In NeurIPS, 2021.
[36] ANM Abeer, Nathan Urban, M Ryan Weil, Francis J Alexander, and Byung-Jun Yoon. Multi-objective latent space optimization of generative molecular design models. *Arxiv*, 2022.
[37] Guoliang Kang, Lu Jiang, Yi Yang, and Alexander G Hauptmann. Contrastive adaptation network for unsupervised domain adaptation. In CVPR, 2019.
[38] Lei Tian, Yongqiang Tang, Liangchen Hu, Zhida Ren, and Wensheng Zhang. Domain adaptation by class centroid matching and local manifold self-learning. IEEE Transactions on Image Processing, 2020.
[39] Francesco Barbato, Marco Toldo, Umberto Michieli, and Pietro Zanuttigh. Latent space regularization for unsupervised domain adaptation in semantic segmentation. In CVPRW, 2021.
[40] Umberto Michieli and Pietro Zanuttigh. Continual semantic segmentation via repulsion-attraction of sparse and disentangled latent representations. In CVPR, 2021.
[41] Samarth Sinha and Adji Bousso Dieng. Consistency regularization for variational auto-encoders. In NeurIPS, 2021.
[42] Hsueh-Ti Derek Liu, Francis Williams, Alec Jacobson, Sanja Fidler, and Or Litany. Learning smooth neural functions via lipschitz regularization. In SIGGRAPH, 2022.
[43] Ramana Subramanyam Sundararaman, Riccardo Marin, Emanuele Rodola, and Maks Ovsjanikov. Reduced representation of deformation fields for effective non-rigid shape matching. In NeurIPS, 2022.
[44] Mete Kemertas and Tristan Aumentado-Armstrong. Towards robust bisimulation metric learning. In NeurIPS, 2021.
[45] David Terjek. Adversarial lipschitz regularization. In ICLR, 2020.
[46] Fengxue Zhang, Yair Altas, Louise Fan, Kaustubh Vinchure, Brian Nord, and Yuxin Chen. Design of physical experiments via collision-free latent space optimization. In NeurIPSW, 2020.