Trieste: Efficiently Exploring The Depths of Black-box Functions with TensorFlow

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-03 11:04:59

Abstract: We present Trieste, an open-source Python package for Bayesian optimization and active learning benefiting from the scalability and efficiency of TensorFlow. Our library enables the plug-and-play of popular TensorFlow-based models within sequential decision-making loops, e.g. Gaussian processes from GPflow or GPflux, or neural networks from Keras. This modular mindset is central to the package and extends to our acquisition functions and the internal dynamics of the decision-making loop, both of which can be tailored and extended by researchers or engineers when tackling custom use cases. Trieste is a research-friendly and production-ready toolkit backed by a comprehensive test suite, extensive documentation, and available at https://github.com/secondmind-labs/trieste.

Main paper: 
 key features and design: 
 We now present the modular structure of TRIESTE which contains four key building blocks, a choice of high-level interface (either AskTellOptimizer or BayesianOptimizer), a choice of ProbabilisticModel, and a pairing of AcquisitionRule and AcquisitionFunction. Although TRIESTE's structure allows a high level of customization, sensible defaults are provided throughout the library in order to give new users good starting points.A key design choice of TRIESTE is its AskTellOptimizer interface, which need not have direct access to the objective function. In many libraries, the objective must be a query-able function to be passed into the loop and called for each BO step, an assumption that is rarely suitable when performing BO in the real world. In contrast, an AskTellOptimizer outputs recommended query points (the ask) and then waits for the user to return new evaluations (the tell) (see Figure 3 ). This interface allows TRIESTE users to apply BO across a range of non-standard real-world settings, e.g., when evaluating the objective function requires laboratory [19] or distributed computing resources, such that users have only partial control over the environment [37] or batches of evaluations arrive asynchronously [20] . For settings where it is appropriate for the objective function to be passed into the BO, e.g., for experiments with synthetic problems, we also provide a more standard BayesianOptimizer interface that will run multiple BO steps.The models in the BO loop can be any model written in TENSORFLOW, and TRIESTE is designed to make it easy to add a new model, through a set of general model interfaces. We provide direct interfaces to import models from well-established TENSORFLOW modeling libraries, e.g., Gaussian processes from GPFLOW [27] and GPFLUX [10] , as well as neural networks from KERAS. TRIESTE users have a range of popular probabilistic models from these libraries available out of the box. They cover both regression and classification tasks and range from standard GPs [40] to alternatives like sparse variational GPs [14] , Deep GPs [43] or Deep Ensembles [25] , that scale much better with the number of function evaluations. Finally, we provide user-friendly model builders with sensible default setups, allowing users to get more quickly to good results and facilitate usage for those with less experience with probabilistic models. Hence, TRIESTE users can benefit from a large choice of models with a wide range of complexity, unlocking novel applications for BO.Importantly, TRIESTE's BO loops allow the modeling of multiple quantities, using either multiple separate models or a single multi-output model. This framework naturally supports common BO extensions like multi-objective optimization [23] , multi-fidelity optimization [46] , optimization with constraints [44] , and combinations thereof.Regardless of the interface and model choice, the recommendation of query points is controlled by an AcquisitionRule. While the vanilla BO rule is to query the point that optimizes a particular AcquisitionFunction, the AcquisitionRule is a useful abstraction to handle complex cases for which a high level of flexibility is needed. For example, Trieste includes variable optimization spaces for AcquisitionFunction [e.g., using trust regions, 9], a multi-step procedure for selecting query points [4] , and a greedy approach to build batches of query points [29] . A wide range of acquisition functions are already provided in TRIESTE to tackle most of the usual BO cases, with many based on the gold-standard Expected Improvement [EI; 18], including variants for batch [6, 12, 2] , noisy [16] , multi-objective [7] and constrained [11] optimization. We also include implementations of recent information-theoretic approaches for multi-fidelity optimization [30] , a scalable extension of batched Thompson sampling [47] , as well as popular active learning methods for improved classification [15] or contour line estimation [35] .TRIESTE is designed to make it straightforward for a user to specify a new acquisition function or rule.Automatic differentiation directly provides the function gradients, which are leveraged by TRIESTE's supported acquisition function optimizers, including an effective parallelized multi-start L-BFGS-B optimizer. Special care is taken to allow AcquisitionFunctions to be updated without expensive retracing of the computational graphs for each BO step.  
feature benchmark: 
 ✓ ✓ ✓ ✓ ✓ ✓ ✓ TENSORFLOW BOTORCH ✓ ✓ ✓ 2 ✓ ✓ ✓ ✓ TORCH EMUKIT × × × ✓ ✓ × ✓ NUMPYTable 1 : A comparison of feature support (according to their documentation) of relevant Python BO libraries still under active maintenance. 
conclusions and future plans: 
 TRIESTE is an open-source project that allows TENSORFLOW practitioners to easily use BO in their systems. It is a highly flexible library designed with modularity in mind, easy to extend with custom models and acquisition functions. Backed by continuous integration and comprehensive unit tests (97% coverage), TRIESTE is a reliable and robust framework used for both real-world deployment and research has recently been taken up by researchers to develop new BO methodology [47, 36, 5, 13, 29, 31, 32, 37, 38, 39] whilst also being used across a range of applications including designing heat exchangers [34] and improving adhesive bonding [19] .TRIESTE relies on features added and improved by the community and so we gladly welcome feature requests and code contributions. We plan to continue to add new functionality orientated to supporting the application of BO in the real world. In the near term, this will include high-dimensional objective functions [3] and non-Euclidean search spaces [28, 42, 8] . 
: 
 Many open-source libraries have been built to support the recent increase in the use and development of BO methodology. For example, Python users with models written in TORCH or NUMPY can easily find compatible libraries such as BOTORCH [2] , GPYOPT [1] , ROBO [22] , EMUKIT [33] or DRAGONFLY [21] . Similarly, those running R, C++ or JAVA can use DICEOPTIM [41] , BAYESOPT [26] or SMAC [17] . The library GPFLOWOPT [24] was built on TENSORFLOW 1 with an intent similar to TRIESTE, but is not actively maintained anymore and does not support the fundamentally different TENSORFLOW 2. 
 
 Appendix: Trieste was the first crewed vessel to reach the bottom of the Mariana trench -the literal global minimum. Submitted to Workshop on Bayesian Decision-making and Uncertainty, 38th Conference on Neural Information Processing Systems (BDU at NeurIPS2024). Do not distribute.The AX library (https://ax.dev/) is a high-level interface for BOTORCH somewhat similar to an ask-tell. 



Reviews: review: title: title: Official Review
main_review: main_review: Summary: Trieste is a powerful and flexible Bayesian optimization toolkit for users within the TensorFlow ecosystem. Its modular design, comprehensive feature set, and strong community support make it a valuable tool for complex and large-scale optimization tasks. However, its TensorFlow-centric approach and potential complexity may pose challenges for beginners and users who prefer other machine learning frameworks.

The paper centers at machine learning system implementation, and more engineering practices are needed to evaluate its quality. The paper is clearly written, but would further benefit from attaching detailed appendix, even if some contents are released in repos. There are some existing open-source Python packages for Bayesian optimization, and Trieste is based on tensorflow, which gives it some strength in originality. However, this may also limit its application, especially for users who prefer other deep learning frameworks like PyTorch. The paper also lacks empirical experiments and lacks discussions on (i) scalability of the framework; (ii) support for modern deep learning architectures in addition to those simple implementations from KERAS.
title: title: A commendable open source contribution, but surprisingly unaware of Vizier
main_review: main_review: This paper presents Trieste, a modular open source library for Bayesian optimization.  The library is designed so that users may mix-and-match modules of four types, providing flexibility and good defaults.  The library supports many standard Bayesian optimization features including batched trials, multi-fidelity, multi-objective optimization, and outcome constraints.  The library compares favorably in feature offering to others such as Ax+BoTorch and Emukit.

I appreciate the naming of the library after a vessel that reached the literal global minimum altitude of Earth. :)

The second sentence of the paper claims that there are no other Bayesian optimization libraries for TensorFlow 2.  However, OSS Vizier was published in early 2023, which is an open source stand-alone version of Google's Bayesian optimization framework.  OSS Vizier seems to have comparable features to Trieste including batched trials and multi-objective optimization.  OSS Vizier may be missing multi-fidelity optimization and support for certain types of constraints.  Could the authors update the paper to include references and comparisons with Vizier?

The paper is well written, and the subject is clearly and explicitly aligned with the theme of the BDU workshop.  The open source contributions are valuable to the community.  So, this paper should be accepted.

Nit-picks:
1. Since Figure 1 and Figure 2 are both part of Figure 3, it'd be appropriate to combine them all as "Figure 1" with a part A and B.
2. Consider adding a reference for L-BFGS-B.
3. Consider tweaking the sentence on lines 104-109.  Simply deleting the word "has" might suffice.  Or adding "which" or "and" after "research", changing "used for both real-world deployment and research has" -> "used for both real-world deployment and research which has".