Return of the Latent Space COWBOYS: Re-thinking the use of VAEs for Bayesian Optimisation of Structured Spaces

Thank you for your thoughtful assessment of our work. We appreciate your recognition of both the descriptive and empirical strengths of our approach. It is encouraging to hear that you find our decoupled strategy promising, and that our competitive results align well with other state-of-the-art methods.

We have addressed all your suggestions thoroughly and have implemented the requested revisions as detailed below. We hope that our efforts to address the comments from you and the other reviewers will be appreciated and could justify increasing your score.

1. Generalizability of COWBOYS Beyond Molecular Design and Drug Discovery. We appreciate your interest in the broader applicability of COWBOYS outside its current focus on molecular design and drug discovery. Our evaluation focuses on these tasks because they remain the most widespread use cases for latent space Bayesian optimization (LSBO), and well-established baseline suites exist to facilitate direct comparison. Fitting a Gaussian Process (GP) in the original, potentially high-dimensional and highly structured space can indeed be challenging with standard kernels. However, we see an ability to model directly in structure space as a key advantage of COWBOYS, rather than a limitation. By working in the raw structure space, our approach naturally supports specialized kernels tailored to particular domains—an area where there is a rich, yet underused, literature. For instance, the Tanimoto kernel, originally proposed by chemists, encodes prior knowledge of what sort of molecular attributes are important in a way that can be especially powerful in low-data regimes [1,2]. Similar structural kernels exist for various structured objects and, via COWBOYS, could now be used for molecular graph optimisation with graph kernels [3], engineering design with 3d mesh kernels [4], optimising computer code via tree kernels [5], or protein design with new protein kernels [6]. In contrast, many current LSBO strategies cannot leverage these specialized kernels because they rely on Euclidean latent spaces. By enabling the direct use of such kernels, COWBOYS allows practitioners to incorporate a wealth of domain-specific prior knowledge. We hope this work helps revitalize interest in harnessing these powerful kernels across a range of complex design problems and spurs further synergy between rich generative models and carefully structured discriminative models

2. We are pleased to hear that our discussion clarified the robustness of COWBOYS with respect to PCN parameters. However, In the revised manuscript, we have added details on how to interpret the corresponding table, emphasizing that no single parameter choice can yield uniformly optimal performance across different objective functions, owing to varying degrees of model mismatch between our GP and each specific problem objective and varying degrees of difficulty/locality of the optimisation problems.

3. “Is the search space problem also an issue for the initial design of COWBOYS”? In standard LSBO, a specific region of the latent space must be chosen, over which we sample an initial design and then perform BO. We argue in Section 3.3 that this clipping can lead to pathological performance issues. In contrast, COWBOYS just samples from the VAE in the standard way, i.e. we generate a random Gaussian (which has infinite support) and then push it through the decoder. There is no need to specify a search space for COWBOYS. We have greatly increased discussion around this point in the final version.

4. Also, thanks for the minor corrections, they have now been changed in the final version, including adding more four specific references to justify ”replacing the VAE’s probabilistic decoding with the most-likely mapping is already a common strategy”.

[1] Griffiths, Ryan-Rhys, et al. "GAUCHE: a library for Gaussian processes in chemistry." Advances in Neural Information Processing Systems 36 (2023): 76923-76946.

[2] Moss, Henry, et al. "Boss: Bayesian optimization over string spaces." Advances in neural information processing systems 33 (2020): 15476-15486.]

[3] Vishwanathan, S. Vichy N., et al. "Graph kernels." The Journal of Machine Learning Research 11 (2010): 1201-1242.

[4] Perez, Raphaël Carpintero, et al. "Gaussian process regression with Sliced Wasserstein Weisfeiler-Lehman graph kernels." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.

[5] Beck, Daniel, et al. "Learning structural kernels for natural language processing." Transactions of the Association for Computational Linguistics 3 (2015): 461-473.

[6] Groth, Peter Mørch, et al. "Kermut: Composite kernel regression for protein variant effects." Advances in Neural Information Processing Systems 37 (2024): 29514-29565.

Thank you very much for your thoughtful review and valuable suggestions. We sincerely appreciate your recognition of the novelty, clarity, rigorous empirical evaluation, comprehensive background coverage, and overall manuscript quality of our work. We have addressed all your suggestions thoroughly and have implemented the requested revisions as detailed below:

1. Generalizability of COWBOYS Beyond Molecular Design and Drug Discovery: We fully agree with your observation on the importance of clarifying the generalizability of our approach. Following your advice, we have expanded our discussion (now included prominently in the introduction and highlighted in the limitations section) regarding the applicability and significant potential of COWBOYS beyond molecular fingerprint-based problems. Please refer to our detailed response to reviewers vn27 and 1JvV for further context. Furthermore, your insightful suggestion about investigating the robustness of COWBOYS against degraded surrogate mode performance (to mimic settings where it may be hard to propose a structure-space kernel) was extremely helpful. As suggested, we have conducted additional experiments on the high-dimensional BO benchmark of Section 7.2, deliberately impairing the Tanimoto GP surrogate by randomly swapping ones to zeros in molecular fingerprints with varying probabilities. This experiment corresponds to adding significant input noise to the fingerprints by randomly removing the count of particular (hashed) molecular substructures. Our findings indicate that COWBOYS maintains robust performance except under the most severe perturbations and have been added to our ablation studies of Appendix B. In summary, perturbing Tanimoto entries with probability 0.01, 0.1, and 0.5 lead to a statistically significant drop in the best molecule found across 0, 3, and 8 tasks from the 25 molecular optimization tasks, respectively.
2. Small additional Points. We have clarified our notation by using boldface m for fingerprints to distinctly differentiate from molecules represented by boldface x. Following your recommendation, we have included a graphical model illustrating the dependencies within both COWBOYS and existing LSBO approaches. We appreciate this suggestion as it significantly enhances the manuscript's clarity. We have added the suggested recent reference, highlighting its relevance. The presence of this paper (and also [1] at the same conference in 2 months time), both using extensions of standard LSBO frameworks, underscores the timeliness and importance of our contributions.

Thank you once again for your constructive feedback, which has significantly strengthened our manuscript.

[1] Lee, Seunghun, et al. "Latent Bayesian Optimization via Autoregressive Normalizing Flows." The Thirteenth International Conference on Learning Representations. 2025.

Thank you very much for your thoughtful review and positive feedback regarding our experimental design and validations. We greatly appreciate your recognition of our method's strong performance. We have carefully addressed all of your suggestions and implemented the requested revisions, detailed below. We hope these efforts demonstrate our commitment to improving the manuscript and may justify an increased evaluation score.

1. Code Accessibility In response to your suggestion we have now made some of our code available here. The repository allows the recreation of the benchmarking results of Section 7.2. Indeed, incorporating COWBOYS (see cowboys.py) required only minimal additions, emphasizing both its practical usability and potential for broad community adoption. The fully-fledged codebase, allowing COWBOYS to be applied to generic pre-trained VAEs and thus applied to a range of downstream tasks (and recreate the results of Sections 7.1 and 7.2), will be provided upon acceptance as it cannot be shared without violating ICML’s anonymity rules.

2. Clarification of Equation (Line 267) We appreciate you pointing out the ambiguity in the equation on line 267. As requested, we have clarified this equation explicitly in the revised manuscript, clearly defining the delta function as indicating a deterministic decoding strategy that selects only the most-likely decoded molecule from the latent location z. The delta function here means that the probability is 0 unless we are taking the most likely prediction from the decoder (from the chosen latent location z).

3. Addressing concerns about discarding the stochasticity of the decoder. We understand your concerns, as we initially were also worried. We agree that there is indeed a theoretical risk of being unable to generate as wide a range of molecules as the stochastic VAE (although not one we observed in practice). We believe the risks of allocating 100% confidence to implausible molecules, which would arise when decoding areas of the latent space where the decoder is highly uncertain, are strongly mitigated by the fact that we are doing conditional VAE sampling (i.e. staying on the well-supported areas of the latent space) rather than exploring the whole latent space like standard LSBO methods. Nevertheless, future work will seek more sophisticated sampling strategies to perform COWBOYS-like strategies for generative models (see our Discussion Section). We like your idea of running the MCMC as described but then sampling from the decoder in order to properly diagnose if the restrictions of the deterministic decoder are limiting the flexibility of our proposed molecules. We have implemented and ran this suggestion for the high-dimensional BO benchmark of Section 7.2, with the full results added to the ablation study of Appendix B. Unfortunately, this lead to a significant drop in performance for 10 of the 25 tasks (and returning statistically similar scores for the other tasks), which we hypothesise is due to the fact that the resulting “sampled” molecules are now not the same as the ones passed to the GP to see if the considered latent code will yield an improvement in score – i.e. we return to a problem similar to that of the alignment issues that we discuss in Section 3.2. Based on your comments, we have greatly increased our discussion of our decision to use a deterministic simplification of the decoder, stressing the points above and that (1) the VAE is still initially trained with a stochastic decoder (it is only at inference/BO time that we take the most likely mapping), and (2) we have added four additional specific references of previous work that also using this deterministic simplification.

4. Clarifying Differing Comparisons in Experiments. The diversity in initial budgets, degrees of parallelism, and subsets of molecular optimization problems stems directly from our intention to replicate and extend the experiments used in influential prior work in LSBO precisely. We believe this breadth provides a robust and comprehensive empirical evaluation of COWBOYS' performance. For peace of mind, please see the comments of the other reviewers, e.g. “Empirical evaluations use the standard benchmarks in the field and replicates and extends the evaluations from previously published papers” and “experiments are conducted thoroughly by evaluating a number of baselines on a good number of test problems”. However, acknowledging your valid concern regarding transparency—especially for readers less familiar with molecular design benchmarks—we have now included an additional section in the appendix. This section explicitly details the structure and interpretation of these benchmark suites, clearly explaining the objectives, settings, and significance of the various optimization tasks we considered.

Thank you for your strongly positive review and for highlighting the novelty, clarity, and thorough experimental design of our work. We have carefully considered your two main comments, both of which guided additional discussion and improvements in the final version of the paper. These insights have helped us better refine our manuscript. We look forward to seeing how this work inspires new lines of inquiry and contributes to the continued advancement of Bayesian optimization.

1. On Incorporating Arbitrary Acquisition Functions. You raise an important question regarding the integration of an arbitrary acquisition function into COWBOYS. In essence, our proposed approach is fundamentally different from traditional BO settings that rely on standard utility-based acquisition functions. Here, we move away from maximizing an acquisition function over a constrained set of equally likely structures. Instead, we explore searching over an entire distribution (as modeled by the VAE) of likely structures. Because of this distribution-based view, conventional acquisition functions (e.g., Expected Improvement or Probability of Improvement) cannot be directly applied without inheriting the same limitations found in current LSBO methods, most notably their disregard for the distribution provided by the VAE. One of our paper’s main contributions is precisely to handle this distribution by replacing maximization-based strategies with a sampling-based approach (Equation 3). While Equation 3 does resemble Probability of Improvement (a point now greatly elaborated in the final version), adapting more sophisticated acquisition schemes is indeed an exciting direction for future research. For instance, we could condition the VAE to generate samples for which an alternative acquisition criterion (e.g., EI) exceeds a certain threshold. Determining this threshold adaptively would open up further avenues, possibly guided by theoretical regret analyses to explain COWBOYS’ strong empirical performance. We believe these avenues will become fertile ground for future studies that aim to bridge existing BO theory with generative models in complex, high-dimensional design spaces.
2. Generalizability of COWBOYS Beyond Molecular Design and Drug Discovery. Please see our response to vn27 for a discussion of the significant potential of COWBOYS in different domains.