Paper: Memoization-Aware Bayesian Optimization for AI Pipelines with Unknown Costs (ICLR-2024)

Reviewer comment (Novelty):
The novelty of the proposed method is limited. The proposed method is a combination of cost-aware BO, unknown-cost modeling, cost cooling, and the memoization trick. All these techniques are not new and have been used in existing BO works (as shown in Table 1). The technical issues or challenges of combining them are not clearly shown in this work. The proposed EEIPU simply replaced the cost function \(c(x)\) of EIPU with a total cost function \(C(x)\) which is defined as the sum of \(c(x)\) over multiple stages. Although the author(s) claimed that computing the expected total inverse cost \(\mathbb{E}[1/C(x)]\) is not a straightforward task, the issue is resolved by conventional MC sampling which does not contribute to a non-trivial solution.

Author response:
**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. 

**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.

Reviewer comment (Experiments):
What's the value of \(\epsilon\) used in the experiments? Is the experimental result sensitive to the settings of \(\epsilon\), \(M\), or \(N\)?

Author response:
**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.

Reviewer comment (Novelty):
The assumption of a multi-stage BO process with independent stages and observable intermediates is a very strong one that can be exploited more than this work currently does. This work only exploits this structure to memoize previously evaluated candidates to save costs expended during the BO process. This structure has been exploited to improve the modelling and guide the search, for example, see "Bayesian Optimization of Composite Functions" (Astudillo and Frazier, 2019) and "Bayesian Optimization of Function Networks" (Astudillo and Frazier, 2021).

Author response:
**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.

Reviewer comment (Experiments):
Perhaps the $M$ prefixes can be sampled more intelligently than randomly, e.g., by weighting the prefixes based on the results of previous evaluations incorporating those prefixes.

Author response:
**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.

Reviewer comment (Writing):
Typos/inconsistencies: 1) Written as $EI \times \mathbb I^\eta$ in Algorithm 1, but $EI * \mathbb I$ and $EI * \mathbb I^\eta$ in Equations (5) and (6); 2) in Figures 3 and 4, left and center column plots are supposed to be best objective value, but the y-axis is written as $f(x^*)$ in the left plots and 'stage costs' in the center plots; 3) [()] described as 'empty prefix' and 'empty subset' in different parts of Sec. 3.2.

Author response:
**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.

Reviewer comment (Theory):
What is the rationale of introducing the $\epsilon$-cost for stages 1 to $\delta$? Why not set to $0$?

Author response:
**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.

Reviewer comment (Evaluation):
There are a few peculiarities with the results for EIPS and CArBO in Figures 3 and 4. In Figure 3 bottom row right plot and Figure 4 top row right plot, EIPS and CArBO consume the exact same cost in each iteration as the non-cost aware EI. This is very strange since EIPS and CArBO are supposed to be cost-aware. In addition, in all the results shown, the non-cost aware EI outperforms EIPS and CArBO in terms of best objective value achieved against cost incurred, which is again strange given that EIPS and CArBO were designed for this very setting. Could you investigate and explain these anomalies?

Author response:
**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.

Reviewer comment (Experiments):
More compelling experiments can be conducted. I think that the experiments tested in this paper seem interesting, but the scale of experiments is small compared to the common scale of the experiments in ICLR.

Author response:
**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.

Reviewer comment (Writing):
The reference by Mockus (1998) might be published in 1978, not 1998.

Author response:
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.

Reviewer comment (Theory):
Why should costs be positive? Is it a mandatory condition?

Author response:
**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.

Reviewer comment (Presentation):
In Figure 2, why are $c_1$ and $c_2$ zero in (c) and (d)? Are they correct? If they are correct, please add description in the rebuttal.

Author response:
**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.

Reviewer comment (Writing):
As the topic covered in this paper is quite new, it will aid more if the paper includes a problem statement to describe in detail the problem setting. It took me quite some efforts to go back and forth to understand the setup of the problem tackled in this paper.

Author response:
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).

Reviewer comment (Novelty):
Please correct me if I'm wrong but it seems the application for this memoization technique is limited, as the stored data in the previous stages can only be reused when a repeated input for a certain stage is queried again (this is related to the 2nd question in the Question section).

Author response:
**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.

Reviewer comment (Experiments):
The number of benchmarks is too few: only 2 synthetic and 2 real-world benchmarks, and with low dimensions.

Author response:
**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.

Reviewer comment (Evaluation):
In the synthetic benchmarks, it seems the improvement is only because the method manages to evaluate more data. Given the same number of iterations (middle column of Figure 3), even methods without cost-awareness like EI can find similar (or even some better) results. So this seems to me that the memoization only helps with increasing the evaluated data, but not really help much with the modelling of the surrogate model or the BO process.

Author response:
**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.

Reviewer comment (Reproducibility):
What are the specifications of GPs using for the cost modelling? Which priors, kernels and hyper-parameter settings are used?

Author response:
**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.

Reviewer comment (Reproducibility):
For synthetic benchmarks, what is the input space for each stage? Is it discrete? If it is continuous, I'm just wondering how exactly does EEIPU re-use the stored data? It is hard for the acquisition function to propose the exact same points in the continuous domain. From the right column of Figure 3, it seems that the amount of reusing data is approximately 40 times over a total of 80 function evaluations.

Author response:
**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.