Memoization-Aware Bayesian Optimization for AI Pipelines with Unknown Costs

Q4: Why not adapt LaMBO and MS_BO to your setting?

A4: Regarding LaMBO and Multistage BO, the brief explanation is that both methods require additional information that are not available in the pipeline setting that EEIPU is designed for. In contrast, EIPS and CArBO do not require these additional pieces of information, and are runnable in the same setting as EEIPU. We explain further:

• LaMBO employs a slowly moving bandit (SMB) algorithm, utilizing a tree structure where each arm represents a candidate variable or hyperparameter in a tuning case. The cost of switching between arms is encoded in the tree, and the switching cost is defined as the distance between the leaves corresponding to the two arms. Importantly, for LaMBO to construct its tree, it needs to be provided with a user-constructed division of the hyperparameter space into disjoint partitions. In Section 3.3 of the LaMBO paper, the authors propose a simple “bisection” heuristic for creating these partitions, but this causes the number of LaMBO tree leaves to grow exponentially with the hyperparameter dimension, in turn causing LaMBO’s computational requirements to also grow exponentially (as it requires evaluating a local BO problem at each leaf). This makes LaMBO not as well-suited to high-dimensional problems. We had reached out to the LaMBO authors to request code and start a discussion, but did not receive a response. Due to these difficulties, we chose to prioritize baselines that could scale well with higher dimensions, such as CArBO and its multi-stage variation. We regret that we could not compare against LaMBO in this paper submission, owing to these circumstances. We hope to provide a LaMBO benchmark in a future work, as it will require substantial effort and time to reproduce (and validate that the reproduction is indeed correct) absent any author-provided code.
• We cited Multistage BO for its stock and resume mechanism that most resembles our memoization technique. There are two key differences from our EEIPU setting: (1) Multistage BO requires stronger assumptions on cost functions.- They assume an easier setting with prior knowledge of the cost of running every stage of their pipeline, and the tunable parameters do not affect costs at all. Modeling of the cost function, as EEIPU does, is unneeded in this easier setting. (2) Multi-stage BO requires each pipeline stage to have a separate objective function, whereas in EEIPU’s setting, there is only one objective function for the entire pipeline. To run Multistage BO in our setting, we would need to make additional strong assumptions, such as assuming that every stage has the same objective function. As a point of clarification, in our EEIPU setting, stage outputs are only used for memoization; we do not assume each stage (except the final one) produces an objective value that can be optimized for. This assumption is consistent with most Machine Learning pipelines, where the desired objective to be optimized, such as a validation accuracy or other metric, is only known after the entire pipeline finishes execution.

For these reasons, we hope we have clarified why the unimplemented techniques either could not be generalized to our setting, or presented unusual difficulties. We also hope that the new and improved experiments we ran in this rebuttal phase, summarized in the global rebuttal section, could portray our efforts to provide a fair and comprehensive assessment of our approach.

Q5: Is performance sensitive to $\epsilon, M,$ and $N$?

A5: $\epsilon$ is only set in our synthetic experiments and has no impact on performance. $M$ and $N$ are parameters used by all BO techniques in the optimization process. We set them in our experiments to their default values, whereas users can tune them freely.

Q6: What causes EIPS and CArBO to perform poorly sometimes?

A6: The relatively poor EIPS and CArBO performance compared to EI on the stacking pipeline is further analyzed in the global rebuttal section.

We would like to thank you again for your valuable review, and apologize for the lengthy reply we drafted. However, we hope that our detailed analysis of the existing methods provided you with a clearer view on our novelty and understanding of our chosen baselines.

Thank you for your detailed review of our paper, as well as the observations you made, all of which will be addressed below. First, we would like to direct you to the global rebuttal section, as well as the new appendix of our draft, for a summary of all our rebuttal efforts to clarify our proposal.

Q1: Novelty

A1: While the components of EEIPU, such as black-box cost modeling and cost-cooling, are inspired by existing cost-aware techniques, there is no method that generalizes cost-awareness to the multi-stage setting, which our results show reduces costs using memoization even in cases where existing cost-aware methods fail to do so. Existing memoization techniques also differ from our setting in several ways, which we elaborate on below:

• LaMBO (Lin et al, 2021) is a technique that aims to minimize the switching costs from one hyperparameter combination to another across two consecutive iterations. They model the hyperparameter search space as a multi-armed bandit, choosing a suffix of stages to update every iteration with the aim of minimizing the number of updated stage parameters from the previous iteration. This method, however, has no knowledge of the cost function of each space, and makes no use of it during the optimization process.
• The Multistage BO (Kusakawa et al, 2021) paper explores the potential of multi-stage pipeline optimization in an industrial setting rather than a theoretical one. The scope of this paper covers pipelines with stagewise intermediate outputs, each of which is modeled by a separate GP, which aids in suspending the pipeline at a stage $k$ if a stage’s output is too far off its modeled expectation. This allows the storing of the $K-1$ first stages, to be used by future iterations for a potentially better final objective than the suspended process. This technique implements a memoization variant, but makes no use of the costs incurred by the corresponding stages.
• While these techniques introduce a memoization-like mechanism, EEIPU is the first method to incorporate memoization-based cost efficiency into the querying process. Using uncertainty modeling, EEIPU is not only able to maximize the exploitation of cheap-to-evaluate areas of the search space through cost awareness, but is also able to use memoized observations to explore high-cost regions without having to incur the cost of rerunning most of the pipeline stages for a second time. This combination allows for faster (less costly) exploration of the search space. Thus, with EEIPU we observe a much higher number of evaluations compared to other methods, when all methods are given a similar budget. Considering that almost all BO methods are able to reach the global optimum as the number of iterations increases, EEIPU is more likely to reach a higher objective value under budget constraints.

Q2: Why not use simple MC Sampling to compute expected inverse cost?

A2: Due to the nature of our stagewise cost-modeling and discounting strategy, we are unable to rely on MC sampling alone to model the expectation of the total inverse cost, as explained in section 3.1 of our paper. However, we acknowledge that our current approach may not be the most accurate way to represent this expectation, and we are actively working towards a more theoretically grounded technique, such as a Taylor expansion.

Q3: Typos

A3: We apologize for the typos in Table 1, Figures 3 and 4, and the last sentence of the Background, all of which have been rectified.

Thank you for the in-depth comments on our paper, and the valuable insights you have provided. To answer your question about the performance of CArBO and EIPS compared to EI, we would like to refer you to our global rebuttal section, where we introduce a new acquisition function, MS_CArBO, that is designed to help understand the true power of memoization, in the events where hyperparameters have little to no impact on the cost function. Additional clarifications and experiments are also summarized in the new appendix of our draft. We will now respond to each of the points you raised in your review.

Q1: The setting of independent stages and intermediate observations can be exploited more. There may be a smarter way to sample the $M$ candidates

A1: The papers you referenced provide a greater view on the potentials of multi-GP modeling approaches, and suggest potentially valuable expansions to our method.

Our cost modeling strategy could especially benefit from the setting explored by the paper “Bayesian Optimization of Function Networks” (Astudillo and Frazier, 2021), their concept could be applied to our cost GPs for identifying a number of stages to run in order to minimize costs, while conditioning this process on the search space of the corresponding objective value, which could further maximize the EEIPU value.

The current scope of this paper serves as a proof of concept, dedicated to showcasing the potential of a memoization-based approach in the cost-aware BO setting. Further improvements are being explored for EEIPU, including your suggestion to implement a smart sampling technique, following this initial proposal. Thank you for making the valuable observation.

Q2: Memoization fails when encountering a stage with noisy outputs

A2: We acknowledge that our method is not applicable to pipelines which require a true (i.e. non-pseudorandom) or external source of randomness. Our setting considers AI applications where pseudorandom number generators are given fixed seeds, which is a common approach in research and production AI because it ensures reproducibility of results, and makes debugging easier.

Q3: Why set costs to $\epsilon$ instead of zero?

A3: We set the cost of running memoized stages to $\epsilon$, because there is always an overhead cost of retrieving outputs from the cache, which can never be zero.

Q4: Why do EIPS and CArBO have similar cost to EI in some cases?

A4: This is a nuanced issue and we would like to explain as best we can. Our experiments show two differences in the behavior of EIPS/CArBO vs EEIPU, which explain why multi-stage cost GPs (as seen in EEIPU but not the other baselines) are important for the pipeline setting: (1) EIPS/CArBO use a single GP to model the cost of the entire pipeline, but for some (though not all) of our synthetic and real experiments, this single-cost-GP assumption turns out to be too coarse-grained: EIPS/CArBO are unable to discover low-cost regions in these multi-stage pipelines; rather, they chose hyperparameters whose cost only differs very slightly from EI (we apologize that this is difficult to see on the Figures; we’ve doubled checked our numbers and confirm the EIPS/CArBO costs are not identical to EI). (2) Even for experiments where EIPS/CarBO incur visibly different cumulative costs versus EI, our proposed EEIPU still has lower cumulative costs. This is due to the substantial cost savings from re-using memoized results. We also note that EEIPU's use of stage-wise cost GPs is a prerequisite for effective memoization, because we cannot model the cost of partially executing a pipeline with the EIPS/CarBO single-cost-GP approach. We will put this discussion into the paper if it is accepted.

Q5: Typos

A5:

• For the EEIPU formula: In section 3.2, we described our cost-discounting process, before showing how the mechanism contributes to EEIPU. The $\eta$ factor was not yet introduced, and was kept for section 3.3, where we defined cost-cooling, and proceeded to include $\eta$ in the equation.
• Thank you for pointing out the typos in our plot labels. Since we commonly use $f(x*)$ to denote the best objective value, we have updated our plots to display it.
• In our paper, we refer to the empty prefix by an empty list [()]. However, when referred to as a subset, we mean the subset {[()]}, which is part of the global set of stored prefixes $P_X$. We apologize for that oversight.

Thank you again for your observations, all of which have been rectified in the new draft, and some of which have provided us with meaningful insights on possible expansions of our method. We hope that our replies, as well as the expansion to our appendix, have provided enough clarity to establish the merit of our paper.

Thank you for your valuable review, which directed our attention to rectifying the details you pointed out, including the EI reference, which we found to be indeed originally published in 1978, the captions on top of the tables, and the mathematical expressions in Figure 1. We have implemented a new acquisition function baseline, MS_CArBO, to further establish the merit and practical advantages of memoization. We also added two different pipelines (1 real and 1 synthetic) as additional sets of experiments, all of which are included in the appendix of our paper (to be moved up to the main content should we be accepted). Please allow us to expand on your points below, one at a time.

Q1: More compelling experiments can be conducted

A1: We designed experiments that showcase a proof of concept on the synthetic pipelines, before testing practicality by optimizing hyperparameters for heavy models, in order to show the advantages of memoization-discounted costs in a practical setting. However, we understand the need for more in-depth experimental analysis, and have run two additional experiments for that reason.

Synthetic

We ran a high-dimensional synthetic experiment, where the objective-GP fits an ensemble of 30 hyperparameters against the final objective value, and 5 different cost-GPs are trained, with a potential of memoizing up to 4 stages of the pipeline.

Real

This experiment optimizes a model consisting of a three-staged T5 transformer fine-tuning and distillation pipeline whose first stage is data preprocessing and has a single parameter, batch size. The second stage is a fine-tuning stage with three hyperparameters, namely learning rate, weight decay and number of epochs. This stage fine-tunes a pre-trained Google t5-small model with 60 million parameters. The final stage is a distillation step with four hyperparameters – learning rate, temperature, weight decay and number of epochs. This stage distills the 60 million-parameter model to a 7 million-parameter model.

The results of this experiment are plotted in the new appendix of our paper, and are also displayed on a table in the global rebuttal section.

Q2: I don’t fully agree with the use of memoization

A2: During this rebuttal period, we have dedicated a portion of our experiments to showing the true impact of memoization in discounting costs, especially when hyperparameters have little to no impact on the runtime costs. We designed a new acquisition function, referred to as MS_CArBO, it is a multi-stage version of the CArBO method with stagewise cost-GPs with no use of memoization, which most accurately portrays the advantage of memoization in reducing evaluation costs, therefore leading to a higher number of iterations. In our synthetic experiment, EEIPU manages to run for 44 iterations more than MS_CArBO, which led to a higher objective value. In the real experiment, EEIPU ran for 24 more iterations than MS_CArBO in the stacking pipeline where hyperparameters have no impact on the costs, and 1 extra iterations in the T5 distillation pipeline. These results, while not consistently staggering in difference, prove a degree of robustness in our proposal in challenging settings where existing BO approaches continue to fail.

Q3: Why should costs always be positive?

A3: In the context of our paper, cost-GPs are designed to model the cost of running a parameter combination through an AI pipeline, which typically refers to the wall-clock time, but could also be set as GPU usage, energy consumption… We, as well as existing cost-aware BO techniques, assume all measures of incurred costs to be positive, hence the modeling of log cost.

Q4: Figure 2 Ambiguities

A4: In Figure 2, we introduce memoization by example, and communicate the idea of discounted costs by showing how the cost of running a parameter combination through a stage of the pipeline goes from a certain value the first time it appears, down to zero when chosen again as a candidate by the memoization algorithm. In reality, there is always a negligible cost incurred even if we retrieve stage outputs from the cache memory. Therefore, memoization costs would always have to take an epsilon value, but that is not currently shown in this introductory Figure. We will update the Figure to show the epsilon term.

Q5: Different $f(x_0)$ for Different Acquisition Functions

A5: The initial data is common across all acquisition functions, which leads to a common best objective value at the start of the optimization process. However, our logging process starts at the end of iteration 0, which increases the chances of different acquisition functions finding a better objective value in the first iteration.

We would like to thank you again for your valuable review, and we hope that our response and additional experiments have successfully addressed your concerns.

Thank you for the great amount of attention and analysis put into understanding our proposal, as well as valid suggestions for improving our paper. We have defined our problem setting, as well as the uniqueness of our proposal compared to existing memoization-based approaches in a new Problem Setup section. We have also carried out a number of new experiments (while we are unable to fully reformat the paper for this rebuttal due to the new experiments that will require moving some plots to the appendix, we have included all of our efforts in the appendix section for the this rebuttal).

Q1: The number of benchmarks is too few

A1: Our synthetic experiments were designed as a proof of concept of EEIPU, before testing its practicality by optimizing hyperparameters for heavy models, in order to show the advantages of memoization-discounted costs in a practical setting. However, we understand the need for more in-depth experimental analysis, and have run two additional experiments for that reason, described below.

Synthetic

We ran a high-dimensional synthetic experiment, where the objective-GP fits an ensemble of 30 hyperparameters against the final objective value, and 5 different cost-GPs are trained, with a potential of memoizing up to 4 stages of the pipeline.

Real

This experiment optimizes a model consisting of a three-staged T5 transformer fine-tuning and distillation pipeline whose first stage is data preprocessing and has a single parameter, batch size. The second stage is a fine-tuning stage with three hyperparameters, namely learning rate, weight decay and number of epochs. This stage fine-tunes a pre-trained Google t5-small model with 60 million parameters. The final stage is a distillation step with four hyperparameters – learning rate, temperature, weight decay and number of epochs. This stage distills the 60 million-parameter model to a 7 million-parameter model.

The results of this experiment are plotted in the new appendix of our paper, and are also displayed on a table in the global rebuttal section.

Q2: Memoization needs to re-query the same inputs to be useful

A2: Memoization is indeed only useful when re-querying an input from a previous iteration. In section 3.2 of our paper, we describe how we guarantee the resampling of previously queried inputs by fixing some candidates’ prefix values from previous evaluations, while randomizing the values of the remaining parameters. This approach guarantees the inclusion of re-queried inputs.

Q3: EEIPU is only useful because it runs for more iterations

A3: Bayesian Optimization guarantees the optimization of the search space with high sample efficiency. Given enough iterations, most BO techniques would be able to reach a similar objective value. Our method leverages this ability, and further improves it by increasing the cost efficiency per evaluation, leading to a higher number of iterations under the same budget, therefore increasing our chances of achieving a higher objective value compared to existing methods.

Q4: The novelty seems to be limited

A4: The idea of memoization is indeed the main novelty of our proposal, whereas cost-cooling is an existing approach to motivate the exploration of high-cost areas. While caching techniques are common and explored in the context of BO, EEIPU is also the first method that supports black-box costs in a multistage setting through multi-GP modeling. By using memoization to minimize the cost of every iteration, our method not only allows the sampling of cheap observations, but also reduces the cost of previously expensive observations with potentially high objective values, by using cache-retrieved outputs to avoid incurring the costs of their corresponding prefix stages. EEIPU is therefore more likely to evaluate more areas of the search space, including expensive ones, than any existing BO technique in the multistage pipeline setting.

Q5: What are the specifications of the GP models?

A5: In our experiments, we trained all GPs under similar specifications, for a fair comparison of relative performance across all acquisition functions. The specifications we used are:

• The hyperparameters $X$ are normalized to unit cube ($X \in [0,1]$), while objective and log-costs are standardized to zero-mean and unit variance.
• The prior is set to constant mean, which defaults to a zero-mean value.
• The covariance is set to the MaternKernel, commonly used to model smooth functions.

Q6: Input space

A6: The input space can be discrete or continuous. In our synthetic benchmarks, we optimize our hyperparameters in a continuous search space, where EEIPU makes use of memoization by using the process defined in the Memoization-based Sampling section above.

Thank you again for your detailed review and observations. We hope that we were able to address your concerns, and that we successfully communicated the merit of our proposal.