Title: Towards Optimal Caching and Model Selection for Large Model Inference
Venue: NeurIPS 2023
Accepted: yes
Keywords: caching

model selection

large language models

foundation models

inference

bandit

regret
Abstract: Large Language Models (LLMs) and other large foundation models have achieved impressive results, but their size exacerbates existing resource consumption and latency challenges. In particular, the large-scale deployment of these models is hindered by the significant resource requirements during inference. In this paper, we study  two  approaches for mitigating these challenges: employing a cache to store previous queries and learning a model selector to choose from an ensemble of models for query processing.

Theoretically, we provide an optimal algorithm for jointly optimizing both approaches  to reduce the inference cost in both offline and online tabular settings. 
By combining a caching algorithm, namely Greedy Dual Size with Frequency (GDSF) or Least Expected Cost (LEC), with a model selector, we achieve optimal rates in both offline and online settings. Empirically, simulations show that our caching and model selection algorithm greatly improves over the baselines, with up to $50\times$ improvement over the baseline when the ratio between the maximum cost and minimum cost is $100$.  Experiments on real datasets show a $4.3\times$ improvement in FLOPs over the baseline when the ratio for FLOPs is $10$, and a $1.8\times$ improvement in latency when the ratio for average latency is $1.85$.

On Optimal Caching and Model Multiplexing for Large Model Inference

Abstract

Large Language Models (LLMs) and other large foundation models have achieved noteworthy success, but their size exacerbates existing resource consumption and latency challenges. In particular, the large-scale deployment of these models is hindered by the significant resource requirements during inference. In this paper, we study two approaches for mitigating these challenges: employing a cache to store previous queries and learning a model multiplexer to choose from an ensemble of models for query processing.

Theoretically, we provide an optimal algorithm for jointly optimizing both approaches to reduce the inference cost in both offline and online tabular settings. By combining a caching algorithm, namely Greedy Dual Size with Frequency (GDSF) or Least Expected Cost (LEC), with a model multiplexer, we achieve optimal rates in both offline and online settings. Empirically, simulations show that the combination of our caching and model multiplexing algorithms greatly improves over the baselines, with up to 50 × improvement over the baseline when the ratio between the maximum cost and minimum cost is 100 . Experiments on real datasets show a 4 . 3 × improvement in FLOPs over the baseline when the ratio for FLOPs is 10 , and a 1 . 8 × improvement in latency when the ratio for average latency is 1 . 85 .

1 Introduction

The recent emergence of Large Language Models (LLMs) and foundation models has significantly increased the capabilities of AI systems (Bubeck et al., 2023; Nori et al., 2023; Ziegler et al., 2019; Ouyang et al., 2022; OpenAI, 2023; Beeching et al., 2023; Chowdhery et al., 2022; Wei et al., 2022a; Google, 2023). This progress comes at a cost, however, of increased resource consumption and latency during both training and inference, presenting challenges not only in real-world deployment but also in terms of environmental impact and energy usage (Sharir et al., 2020; Patterson et al., 2021; Bommasani et al., 2022). For instance, LLM-based chatbots typically consist of large transformerbased networks with parameter counts ranging from one to several hundred billion (Zhou et al., 2023). Moreover, the auto-regressive nature of LLMs exacerbates the issue of latency and resource consumption because the model can only generate one token at a time. Thus, compared to traditional AI-powered services, language model inference costs are much higher and the latency is significantly

Lianmin Zheng

Jiantao Jiao

longer, making it nearly impossible to process each query using LLMs in high-throughput query systems such as search engines.

In this paper, we explore two simple yet effective strategies to mitigate this problem: (1) employing a caching system to store previous queries, and (2) developing a model multiplexer to choose the most appropriate model from a set of models for processing the queries. The general workflow of our proposed LLM-based inference system is shown in Figure 1: upon receiving a query or prompt, we initially check if it can be retrieved from the cache. If the query is not found in the cache, we employ the model multiplexer to determine which model should be used for processing it first, based on the estimated cost for both models.

The choice of cost function and models can vary based on the goal. One measure of cost, for example, could be floating point operations (FLOPs). Other alternatives could include the number of API calls as a measure of resource consumption, latency as a measure of time consumption, or a score provided by a user as a measure of user satisfaction. The cost could also be a weighted sum of multiple factors. For the models, a natural choice would be to have a small and a large model, where the small model costs less and is also less accurate, and the large model has a higher cost and also provides higher accuracy. Another alternative would be to have models with expertise in different areas, i.e., each model has high accuracy in its own area of expertise. We provide more discussion in Appendix A.

Cache

query

hit

miss

Multiplexer

Large Model

There is a long history of existing literature on caching algorithms, with prominent applications including computer architecture and web retrieval (Smith, 1982; Wang, 1999; Kumar and Singh, 2016). Existing caching algorithms deal with queries with different frequencies and cost, and must also provide guidelines for choosing the cache size. In addition to these well-known difficulties, the use of caching for LLMs raises new challenges, including:

The need for fuzzy search. Since the prompt lies in a discrete space that is exponentially large with respect to the token size, it is impossible to match and save all distinct queries. Thus, to be at all useful, approximate matching and grouping is required when retrieving queries saved in the cache.

The randomness of the cost. The cost for processing each query is a random variable that depends on the query and has a large variance due to the auto-regressive generation procedure and the difference in the length and quality of generated responses. When combined with the long-tailed distribution of the query frequency, the estimation of the cost requires a non-trivial algorithm design.

The effect of model multiplexing. When the cache system is combined with the model multiplexer, the estimation of cost must change accordingly to take into consideration the different costs induced by various models.

For the fuzzy search problem, semantic search or vector-embedding-based ideas provide a systematic solution that includes embedding extraction and matching algorithms (Bast et al., 2016; Chang et al., 2020; Kamalloo et al., 2023). To simplify the problem, we assume that there exists some semantic search oracle that can group the prompts with the same semantic meaning and that the total cache size is limited by the number of queries, ignoring the difference in cache size between each individual query and response.

The remainder of this paper is organized as follows. In Section 2, we formally define the pipeline of caching and model multiplexing. In Section 3, we study the optimality of the Least Expected Cost (LEC) caching strategy, which estimates the frequency and cost of processing each query, and evicts the one with the least estimated expected cost when there is only one model to call. In section 4, we

consider the case when we have access to two models, and jointly design optimal caching and model multiplexer. In both sections, we start by assuming there are infinite samples and then analyze the offline and online learning cases where the cost and frequency need to be learned from data. The experimental results are presented in Section 5. We discuss the potential choices of cost, model, and output in the real world in Appendix A. We provide a brief discussion of the generalization to variable cache sizes in Appendix B and of the generalization to multi-model multiplexing in Appendix C.

1.1 Related work

Cache replacement algorithms Traditional cache replacement algorithms investigate optimal ways to cache queries with different frequencies, costs, and cache sizes. To address varying frequencies, a standard approach is to use a Least Frequently Used (LFU) or Least Recently Used (LRU) cache eviction strategy (Lee et al., 2001). These have been proven to be optimal for both adversarial and stochastic queries (Stallings and Paul, 2012; Bura et al., 2021). Caching has also been combined with machine learning advice and online learning analysis in the literature (Chang et al., 2018; Shuja et al., 2021; Jiang et al., 2019; He et al., 2017; Mukhopadhyay and Sinha, 2021; Faizal et al., 2023). When varying costs and varying frequencies exist simultaneously, Jin and Bestavros (2000); Arlitt et al. (2000) propose and study the Greedy Dual-Size with Frequency (GDSF) replacement algorithm, which takes both frequency and cost into consideration. Bahn (2005) proposes the Least Expected Cost (LEC) algorithm, which is similar to GDSF, except that it estimates frequency from data. Our work extends this idea by attempting to learn a model for both frequency and cost from data. Moreover we explore the statistical optimality of these algorithms in both offline and online settings. We also investigate combining caching algorithms with model multiplexing in order to boost performance.

Acceleration of LLM inference Much effort has been devoted to reducing the cost and latency of LLMs during inference. For example, post-training quantization-based approaches aim to compress the model size by using lower-precision arithmetic without losing too much accuracy (Gholami et al., 2021; Frantar et al., 2023). Early-exit frameworks aim to utilize the output in the middle decoder blocks so that only a small fraction of decoder blocks are called when processing a query (Bakhtiarnia et al., 2022; Schuster et al., 2022). The Mixture of Experts approach designs a gating function that only assigns a small fraction of the network for each query (Fedus et al., 2022). Embedding recycling caches activations from an intermediate layer of a pre-trained model to accelerate the training and inference procedure (Du et al., 2020; Wei et al., 2022b; Saad-Falcon et al., 2023). LLM cascade starts with the smallest model and continues to call larger models if the output is not acceptable (Chen et al., 2023b). The big little transformer decoder framework uses a smaller model to generate a draft response and calls the large model to identify the unreliable tokens and perform correction (Kim et al., 2023). Similar ideas have been combined with speculative sampling to guarantee that the output remains the same in distribution as that of the large models (Chen et al., 2023a; Leviathan et al., 2022).

2 Formulation

We formalize the workflow in Figure 1. Consider the set of (finite) prompts / queries Q ⊂ R d . In the t -th round, a query q t ∈ Q is sampled from a fixed population distribution P ∈ ∆( Q ) . We maintain a small set of cache L t ⊂ Q with |L t | ≤ L . We say the query hits the cache if the query satisfies q t ∈ L t . When the query hits the cache, the incurred cost is zero. When the query does not hit the cache, we choose among the existing models to process the query.

In the processing stage, we first describe the setting of caching without model multiplexing, and extend it to the case of caching with model multiplexing.

2.1 Caching without model multiplexing

In the case when we only have one model, let C l ( q ) denote the random variable of the cost when processing the query with the model. Assume that C l ( q ) is supported on [ B 1 , B 2 ] with B 2 > B 1 > 0 being the upper and lower bounds for the cost. Let c ⋆ l ( q ) = E [ C l ( q )] be the expected true cost of

processing the query q . The cost for a given query q and cache L can be written as:

cost ( q, L ) = ✶ ( q ̸∈ L ) E [ C l ( q )] = ✶ ( q ̸∈ L ) c ⋆ l ( q ) . By taking the expectation over the distribution q , we have the expected cost as

In the offline learning setting, we collect an offline dataset and hope to learn a caching policy ̂ L such that cost ( ̂ L ) is minimized.

In the online setting, the query comes in a streaming fashion. At the beginning of each round, we receive a query q t . If the query misses the current cache L t , we let the model process the query and receive a cost c t ∼ P C l . Then we can choose to update the cache L t by adding the current query and response to the cache, and replacing one of the existing cached items if the cache L t is full. If the query hits the cache q t ∈ L t , then the cost for this round is set to zero with no more observations. In this case, we are interested in characterizing the average difference in the cost throughout the execution of the online learning process. This can be characterized by the regret:

2.2 Caching with model multiplexing

For the simplicity of the notation, we focus on the case of selecting from a small model and a large model, 1 and discuss how it can be generalized to the case of selecting from multiple models in Appendix C. Let C s ( q ) denote the random variable of the cost when processing the query with the small model, and C l ( q ) denote the random variable of the cost when processing the query with the large model. We assume that both random variables are supported on [ B 1 , B 2 ] . We observe i.i.d. draws of the random variables C s ( q ) when executing the small model, and C l ( q ) when executing the large model. Denote the expected cost as c ⋆ s ( q ) = E [ C s ( q )] and c ⋆ l ( q ) = E [ C l ( q )] .

Let π : Q ↦→ [0 , 1] be the (possibly random) model multiplexing policy that maps the query q to values in [0 , 1] , where π ( q ) = 1 represents that the query is always sent to the small model, and π ( q ) = 0 represents the query is always sent to the large model. The randomness in the policy π is independent of the cost C s ( q ) , C l ( q ) . The total cost can be written as the following function of the query q , cache L and policy π :

✶ By taking the expectation over q , we have the expected cost as

In the offline learning setting, we collect an offline dataset and hope to learn a caching policy ̂ L and a multiplexer ˆ π such that cost ( ̂ L , ˆ π ) is minimized. In the online setting, we get to update the cache in each round by adding the current query into the cache and evicting the ones in the cache if full. When the query q t misses the cache in round t , we will observe a sample from C s ( q t ) if it is processed by the small model, or a sample from C l ( q t ) if it is processed by the large model. There will be no observations of cost if q t hits the cache. We aim at minimizing the regret:

3 Optimal Caching without Model multiplexing

3.1 Population setting

We start with the population setting where the probability distribution P and the cost c ⋆ l are both known. In the case with only one model, the optimal caching strategy is the Least Expected Cost (LEC) or Greedy Dual Size with Frequency (GDSF) algorithm:

The traditional frequency-based caching strategy, including Least Recent Used (LRU) and Least Frequently Used (LFU), aims at caching the most frequent queries:

Weshow in Appendix D that the ratio between the cost of LFU and LEC can be as high as max q ∈Q c ⋆ l ( q ) min q ∈Q c ⋆ l ( q ) in the worst case, which shows that LFU can be highly suboptimal when the cost varies significantly.

3.2 Finite sample setting: Offline learning

The previous section characterizes the optimal caching strategy in the population setting. We now consider the finite-sample offline learning setting, where we hope to produce a cache L based on prior data such that the introduced cost is minimized. Denote D N = { ( q 1 , c 1 ) , · · · , ( q N , c N ) } , where q i is sampled from the distribution P ( · ) , and c i is a sample from random variable C l ( q i ) . We consider estimating P, c ⋆ l from oracles ˆ P = DenEstOracle ( q 1 , · · · , q N ) , ˆ c l ( q ) = RegressionOracle ( D N ) . In practice, one may remove the last layer of the pre-trained language model and concatenate it with a linear head and fine-tune the model as the estimator. For theoretical analysis, we focus on the tabular case, where we set both ˆ P and ˆ c l ( q ) to be the plug-in estimator:

✶ In practice, the distribution of q may have a long tail. Although the estimation of P ( q ) is uniformly good for all q , the estimation of c ⋆ ( q ) can be bad for the queries that are visited less. To select the maximum L elements from the imbalanced samples, we compensate the plug-in estimator by introducing pessimism (Rashidinejad et al., 2021; Jin et al., 2021) 2 . As we show in Lemma 1, the true frequency for any query q ∈ L ⋆ is lower bounded by some constant that depends on B 1 , B 2 , |Q| . Thus the pessimism helps eliminate those less visited queries in the long tail of the distribution and encourages caching the queries in L ⋆ . The lower-confidence-bound based estimator is:

Theorem 1. Assume that N ≥ 8 B 2 |Q| log(3 L/δ ) B 1 and taking δ = 1 /N . We have

The proof is deferred to Appendix E, where we prove a stronger high-probability bound rather than a bound in expectation. From the theorem, we know that the cost of the finite-sample caching policy converges to the cost of the optimal policy at a rate of 1 / √ N , which achieves the optimal dependence on N . The insights from the tabular case also indicate that the cost needs to be estimated in a conservative fashion when considered for the cache replacement algorithm.

3.3 Finite sample setting: Online learning

We summarize the caching algorithm pipeline in 1, which relies on the two estimation oracles, DenEstOracle and RegressionOracle , which estimate both the frequency and cost of models from data.

Algorithm 1 Caching in Online Learning 1: Initialize the set of cache L 1 = {} , past observations H 1 = {} , ˆ c l, 0 ( q ) = B 1 , ∀ q ∈ Q . 2: For iteration t = 1 , 2 · · · , T 3: Receive query q t . 4: Update the density estimation ˆ P t = DenEstOracle ( q 1 , · · · , q t ) . 5: If q t ∈ L t : 6: Output the cached result, set ˆ c l,t = ˆ c l,t -1 , update the past observation H t = H t -1 ⋃ ( q t , × ) , and continue. 7: Use the large model to process the query, and observe a cost c t ∼ P C l ( q ) . 8: Update the past observation H t = H t -1 ⋃ ( q t , c t ) . 9: Update ˆ c l,t = RegressionOracle ( H t ) . 10: If |L t | < L : 11: Let L t +1 be the union of L t and q t . 12: Else if ˆ P t ( q t ) · ˆ c l,t ( q t ) > min q ∈L t ˆ P t ( q ) · ˆ c l,t ( q ) : ˆ

13: Replace the minimizer element of P t ( q ) · ˆ c l,t ( q ) in the cache L t with q t to get L t +1 .

For theoretical analysis, we focus on the tabular case and define the oracles as follows:

̸

̸

̸

For the estimation of density, we use plug-in estimator since there is no imbalance in the sampling process. For the estimation of the cost, we subtract the confidence bound to include pessimism. We have the following regret guarantee.

Theorem 2. When substituting the DenEstOracle and RegressionOracle with Equation (3) and (4) and set δ = 1 /T , we have for some universal constant C :

On the other hand, for any caching policy {L t } T t =1 , there exist some cases of P ( q ) , c ⋆ l ( q ) such that for some universal constant C ′ ,

The proof is deferred to Appendix F. Different from the offline case, one interesting feature of the online case is the partial observation phenomenon : when the query hits the cache, it will not be processed by the model, and thus we cannot observe the sample from C l ( q ) in this round. This is different from the traditional bandit literature where the selected arm is always observed in each round. Thus the partial observation thus requires new upper and lower bound analysis.

4 Optimal Caching and Model multiplexing

4.1 Population setting

In the case when we have access to two models, we need to design a good caching and model multiplexing strategy jointly. We can compute the optimal caching and model multiplexing policy as

̸

L ⋆ , π ⋆ = arg min cost ( L , π )

Such optimal strategies are straightforward: π ⋆ always assigns the query to the model with a smaller cost, and L ⋆ saves the L queries with the largest P ( q ) · min( c ⋆ s ( q ) , c ⋆ l ( q )) .

For the model multiplexing algorithm, we consider two baselines: (a) one always uses large model π l ( q ) ≡ 0 ; (b) one always uses the small model π s ( q ) ≡ 0 . This is related to the LLM cascade idea in the concurrent work of Chen et al. (2023b). We provide more discussion in Appendix A, and present comparisons between baselines and π ⋆ in Appendix D.

4.2 Finite sample setting: Offline learning

We now consider the finite sample case. Let D N = { ( q 1 , c s, 1 , c l, 1 ) , · · · , ( q N , c s,N , c l,N ) } , where c s,n is a sample from random variable C s ( q n ) , the observed cost for processing query q n with the small model in round n . And c l,n is a sample from random variable C l ( q n ) , the observed cost for processing query q n with the large model in round n . We consider estimating P, c ⋆ s , c ⋆ t with some oracles ˆ P = DenEstOracle ( q 1 , · · · , q N ) , ˆ c s ( q ) , ˆ c t ( q ) = RegressionOracle ( D N ) . We focus on the tabular case for theoretical analysis, where we set ˆ P , ˆ c s ( q ) and ˆ c l ( q ) to be the plug-in estimator:

✶ Similar to the case of caching without model multiplexing, for a long-tailed distribution P ( q ) , the estimation of c ⋆ s ( q ) , c ⋆ l ( q ) can be bad for the queries that are visited less. To select the maximum L elements from the plug-in estimator, we introduce pessimism to the estimate of ˆ c l and ˆ c s . This leads to the following design of caching and model multiplexer ˆ L and ˆ π :

✶ We now show the cost for the caching and model multiplexer obtained from the empirical estimate is close to the optimal cost. The proof is deferred to Appendix G.

Theorem 3. Assume that N ≥ 8 B 2 |Q| log(4 L/δ ) B and take δ = 1 /N

4.3 Finite sample setting: Online learning

We turn to the online case. We first propose a meta-algorithm in Algorithm 2. We provide a theoretical analysis of the meta-algorithm for the tabular case, with DenEstOracle ˆ P t ( q ) = ∑ t i =1 ✶ ( q i = q ) t , and the RegressionOracle defined as follows: t

✶ ✶ We provide the following theorem on the regret of the overall algorithm.

Algorithm 2 Joint Design of Caching and Model multiplexing

1: Initialize the set of cache L 1 = {} , past observations H 1 = {} , ˆ c l, 0 ( q ) = B 1 , ˆ c s, 0 ( q ) = B 1 , model multiplexing policy π 0 ( q ) = 1 , ∀ q ∈ Q .

2: For iteration t = 1 , 2 · · · , T

3: Receive query q t .

4: Update the density estimation ˆ P t = DenEstOracle ( q 1 , · · · , q t ) .

5: If q t ∈ L t : output the cached result, set ˆ c s,t = ˆ c s,t -1 , ˆ c l,t = ˆ c l,t -1 , π t = π t -1 , update the past observation H t = H t -1 ⋃ ( q t , × , × ) , and continue.

6: Select the models according to s t = π t ( q t ) .

7: Update the past observation H t = H t -1 ⋃ ( q t , s t , c t ) .

8: Update ˆ c l,t , ˆ c s,t = RegressionOracle ( H t ) . Set π t +1 ( q ) = ✶ (ˆ c s,t ( q ) < ˆ c l,t ( q )) . 9: If |L t | < L : let L t +1 be the union of L t and q t . 10: Else if ˆ P t ( q t ) · min(ˆ c s,t ( q t ) , ˆ c l,t ( q t )) > min q ∈L t ˆ P t ( q ) · min(ˆ c s,t ( q ) , ˆ c l,t ( q )) :

11: replace the minimizer element in the cache L t on the RHS with q t to get L t +1 .

Theorem 4. Substituting the oracles in Algorithm 2 with the oracles above and δ = 1 /T , we have

The proof is deferred to Appendix H. Compared with the lower bound in Theorem 2, we see that the dependency on T is tight. The pessimism plays two different roles here: on the one hand, it encourages the exploration for model multiplexing to choose the ones with more uncertainty in the cost; on the other hand, it encourages the exploitation to be conservative about which query to save into the cache.

For the model multiplexer to work well, one needs to have a small yet accurate model multiplexer. In the case when the model multiplexer is not accurate, the small model always comes with a much smaller cost, and we are allowed to regenerate the responses and make corrections for the output, one may combine LEC with cascade (Chen et al., 2023b) to achieve better performance.

5 Experiments

We conduct both simulations and real-world experiments with our proposed methods. The code is available at https://github.com/Ying1123/llm-caching-multiplexing .

5.1 Simulations for algorithm analysis

We conduct synthetic online and offline experiments for joint optimization of caching and model switching. In Figure 2, we plot the cumulative cost and regret in online learning for LFU and LEC caching algorithms. For LFU, we consider model switchers which always select the small or large models as the baselines. We consider 20 distinct prompts and set the cache size to be 10 . We set the frequency distribution as power distribution with α = 0 . 9 . The ground truth cost for each query processed by both models is set as a sample from 100 X +1 , where X is a random variable generated from a Bernoulli distribution with the parameter 0 . 5 . We repeat the simulation 100 times and plot the mean and standard deviation in the figure. Our simulation suggests that LEC with model switcher greatly improves the two baselines by a factor of 50 × when the cost ratio is 100 . We include additional results on the synthetic datasets for both online and offline settings with different α values, cost ratios, and switcher accuracy in Appendix I.1.

5.2 Experiments on real datasets

We evaluate our algorithms on two tasks: next-token prediction on the Lambada (Paperno et al., 2016) dataset and chat assistant on the OpenAssistant (Köpf et al., 2023) dataset.

For the next-token prediction task, we run the offline algorithm with two models: OPT-1.3B and OPT-13B (Zhang et al., 2022) and use FLOPs as the cost. The target performance metric is the number of correct tokens predicted, where we get the ground-truth token from the Lambada dataset.

10 0

10 1

10 2

10 3

10 4

10

2

10

1

10 0

10 1

10 2

10 3

10 4

10 5

Cumulative Cost

LFU+small

LFU+large

LEC+selector

Optimal

10 0

10 1

10 2

10 3

10 4

10 1

10 2

10 3

10 4

10 5

Cumulative Regret

LFU+small

LFU+large

LEC+selector

steps steps

For a given query, an algorithm can choose to run the small model or the large model. If the small model is chosen but its result is wrong, the large model must be run and it will incur an additional penalty. We fine-tune a BERT base model with 2000 samples as the model switcher by predicting whether the small model can give the correct result and achieve 80.2% accuracy. We work with 100 unseen distinct prompts in the offline setting with total queries 10000 and cache size 40 . We compare our offline caching and switcher algorithms against LFU, large-model-only, and cascade (which always calls the small model first). As shown in Table 1, LEC is better than LFU in all cases. Combining LEC and switcher brings up to 4 . 3 × cost reduction compared to the baseline 'LFU + Large.' However, as the predictor accuracy is limited, the model switcher may not be as good as the cascade algorithm in some cases. We leave the training of a better switcher as future work.

On the chat assistant task, we run the online algorithm with two models: FastChat-T5-3B and Vicuna-13B (Chiang et al., 2023), and use the inference latency as the cost. The quality of response is evaluated by GPT4 evaluation (Liu et al., 2023). We say a response is satisfying if the score is larger than 6 out of 10, and unsatisfying otherwise. If the response from the small model is unsatisfying, we will call the large model again and incur an additional cost in latency. The ratio between the average latency of the large model and the small model is 1.85. We work with 100 distinct prompts in the online setting with total queries 10000 and cache size 40 . After a sufficient number of online learning steps, the switcher learns the accurate costs of two models on this finite prompts set, so 'LEC + switcher' outperforms other algorithms in all cases on Table 2 with up to 1 . 8 × latency reduction compared to "LFU + large" baseline.

We provide more experiments in Appendix I, where we evaluate both FLOPs and latency for both offline and online setting on both synthetic and real dataset, with varying cache size, query size and distinct prompts.

6 Conclusions

We have studied the joint optimization of caching and model multiplexing and proposed an optimal algorithm for the tabular case. There are a variety of further work that can be pursued in this vein, including:

Designing the optimal caching and model multiplexing algorithm when there is a query queue, such that the query arrives at a random interval rather than a fixed interval. A more complicated serving pattern also needs to take batching strategies into consideration.

Understanding the scaling law of the predictors. We hope to use a small yet accurate model for prediction to reduce overhead introduced by the predictor. It is important to understand the trade-off between prediction accuracy, model size, and training data size.

Designing optimal caching algorithm when the responses generated in each round have diverse qualities.

References

Arlitt, L. Cherkasova, J. Dilley, R. Friedrich, and T. Jin. Evaluating content management techniques for web proxy caches. ACM SIGMETRICS Performance Evaluation Review , 27(4): 3-11, 2000.

Bahn. Web cache management based on the expected cost of web objects. Information and Software Technology , 47(9):609-621, 2005.

Bakhtiarnia, Q. Zhang, and A. Iosifidis. Single-layer vision transformers for more accurate early exits with less overhead, 2022.

Bast, B. Buchhold, E. Haussmann, et al. Semantic search on text and knowledge bases. Foundations and Trends® in Information Retrieval , 10(2-3):119-271, 2016.

Beeching, Y. Belkada, K. Rasul, L. Tunstall, L. von Werra, N. Rajani, and N. Lambert. StackLLaMA: An RL fine-tuned LLaMA model for Stack Exchange question and answering, 2023. URL https://huggingface.co/blog/stackllama .

Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein, J. Bohg, A. Bosselut, E. Brunskill, E. Brynjolfsson, S. Buch, D. Card, R. Castellon, N. Chatterji, A. Chen, K. Creel, J. Q. Davis, D. Demszky, C. Donahue, M. Doumbouya, E. Durmus, S. Ermon, J. Etchemendy, K. Ethayarajh, L. Fei-Fei, C. Finn, T. Gale, L. Gillespie, K. Goel, N. Goodman, S. Grossman, N. Guha, T. Hashimoto, P. Henderson, J. Hewitt, D. E. Ho, J. Hong, K. Hsu, J. Huang, T. Icard, S. Jain, D. Jurafsky, P. Kalluri, S. Karamcheti, G. Keeling, F. Khani, O. Khattab, P. W. Koh, M. Krass, R. Krishna, R. Kuditipudi, A. Kumar, F. Ladhak, M. Lee, T. Lee, J. Leskovec, I. Levent,

X. L. Li, X. Li, T. Ma, A. Malik, C. D. Manning, S. Mirchandani, E. Mitchell, Z. Munyikwa, S. Nair, A. Narayan, D. Narayanan, B. Newman, A. Nie, J. C. Niebles, H. Nilforoshan, J. Nyarko, G. Ogut, L. Orr, I. Papadimitriou, J. S. Park, C. Piech, E. Portelance, C. Potts, A. Raghunathan, R. Reich, H. Ren, F. Rong, Y. Roohani, C. Ruiz, J. Ryan, C. Ré, D. Sadigh, S. Sagawa, K. Santhanam, A. Shih, K. Srinivasan, A. Tamkin, R. Taori, A. W. Thomas, F. Tramèr, R. E. Wang, W. Wang, B. Wu, J. Wu, Y. Wu, S. M. Xie, M. Yasunaga, J. You, M. Zaharia, M. Zhang, T. Zhang, X. Zhang, Y. Zhang, L. Zheng, K. Zhou, and P. Liang. On the opportunities and risks of foundation models, 2022.

Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al. Sparks of artificial general intelligence: Early experiments with GPT-4. arXiv preprint arXiv:2303.12712 , 2023.

Bura, D. Rengarajan, D. Kalathil, S. Shakkottai, and J.-F. Chamberland. Learning to cache and caching to learn: Regret analysis of caching algorithms. IEEE/ACM Transactions on Networking , 30(1):18-31, 2021.

W.-C. Chang, F. X. Yu, Y.-W. Chang, Y. Yang, and S. Kumar. Pre-training tasks for embedding-based large-scale retrieval, 2020.

Chang, L. Lei, Z. Zhou, S. Mao, and T. Ristaniemi. Learn to cache: Machine learning for network edge caching in the big data era. IEEE Wireless Communications , 25(3):28-35, 2018.

Chen, S. Borgeaud, G. Irving, J.-B. Lespiau, L. Sifre, and J. Jumper. Accelerating large language model decoding with speculative sampling, 2023a.

Chen, M. Zaharia, and J. Zou. FrugalGPT: How to use large language models while reducing cost and improving performance, 2023b.

W.-L. Chiang, Z. Li, Z. Lin, Y. Sheng, Z. Wu, H. Zhang, L. Zheng, S. Zhuang, Y. Zhuang, J. E. Gonzalez, I. Stoica, and E. P. Xing. Vicuna: An open-source chatbot impressing GPT-4 with 90%* ChatGPT quality, March 2023. URL https://lmsys.org/blog/2023-03-30-vicuna/ .

Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al. PaLM: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

Du, M. Ott, H. Li, X. Zhou, and V. Stoyanov. General purpose text embeddings from pre-trained language models for scalable inference, 2020.

Z. Faizal, P. Singh, N. Karamchandani, and S. Moharir. Regret-optimal online caching for adversarial and stochastic arrivals. In Performance Evaluation Methodologies and Tools: 15th EAI International Conference, VALUETOOLS 2022, Virtual Event, November 2022, Proceedings , pages 147-163. Springer, 2023.

Fedus, J. Dean, and B. Zoph. A review of sparse expert models in deep learning. arXiv preprint arXiv:2209.01667 , 2022.

Frantar, S. Ashkboos, T. Hoefler, and D. Alistarh. GPTQ: Accurate post-training quantization for generative pre-trained transformers, 2023.

Gholami, S. Kim, Z. Dong, Z. Yao, M. W. Mahoney, and K. Keutzer. A survey of quantization methods for efficient neural network inference, 2021.

Google. PaLM-2 technical report. 2023.

He, Z. Zhang, F. R. Yu, N. Zhao, H. Yin, V. C. Leung, and Y. Zhang. Deep-reinforcement-learningbased optimization for cache-enabled opportunistic interference alignment wireless networks. IEEE Transactions on Vehicular Technology , 66(11):10433-10445, 2017.

Jiang, G. Feng, S. Qin, T. S. P. Yum, and G. Cao. Multi-agent reinforcement learning for efficient content caching in mobile d2d networks. IEEE Transactions on Wireless Communications , 18(3): 1610-1622, 2019.

Jin and A. Bestavros. Popularity-aware greedy dual-size web proxy caching algorithms. In Proceedings 20th IEEE International Conference on Distributed Computing Systems , pages 254261. IEEE, 2000.

Jin, Z. Yang, and Z. Wang. Is pessimism provably efficient for offline RL? In International Conference on Machine Learning , pages 5084-5096. PMLR, 2021.

Kamalloo, X. Zhang, O. Ogundepo, N. Thakur, D. Alfonso-Hermelo, M. Rezagholizadeh, and J. Lin. Evaluating embedding APIs for information retrieval, 2023.

Kim, K. Mangalam, J. Malik, M. W. Mahoney, A. Gholami, and K. Keutzer. Big little transformer decoder, 2023.

Köpf, Y. Kilcher, D. von Rütte, S. Anagnostidis, Z.-R. Tam, K. Stevens, A. Barhoum, N. M. Duc, O. Stanley, R. Nagyfi, et al. Openassistant conversations-democratizing large language model alignment. arXiv preprint arXiv:2304.07327 , 2023.

Kumar and P. Singh. An overview of modern cache memory and performance analysis of replacement policies. In 2016 IEEE International Conference on Engineering and Technology (ICETECH) , pages 210-214. IEEE, 2016.

Lee, J. Choi, J.-H. Kim, S. H. Noh, S. L. Min, Y. Cho, and C. S. Kim. LRFU: A spectrum of policies that subsumes the least recently used and least frequently used policies. IEEE Transactions on Computers , 50(12):1352-1361, 2001.

Leviathan, M. Kalman, and Y. Matias. Fast inference from transformers via speculative decoding. arXiv preprint arXiv:2211.17192 , 2022.

Liu, D. Iter, Y . Xu, S. Wang, R. Xu, and C. Zhu. Gpteval: Nlg evaluation using gpt-4 with better human alignment. arXiv preprint arXiv:2303.16634 , 2023.

Mukhopadhyay and A. Sinha. Online caching with optimal switching regret. In 2021 IEEE International Symposium on Information Theory (ISIT) , pages 1546-1551. IEEE, 2021.

Nori, N. King, S. M. McKinney, D. Carignan, and E. Horvitz. Capabilities of GPT-4 on medical challenge problems. arXiv preprint arXiv:2303.13375 , 2023.

OpenAI. GPT-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

Ouyang, J. Wu, X. Jiang, D. Almeida, C. L. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. Training language models to follow instructions with human feedback. arXiv preprint arXiv:2203.02155 , 2022.

Paperno, G. Kruszewski, A. Lazaridou, Q. N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv:1606.06031 , 2016.

Patterson, J. Gonzalez, Q. Le, C. Liang, L.-M. Munguia, D. Rothchild, D. So, M. Texier, and J. Dean. Carbon emissions and large neural network training. arXiv preprint arXiv:2104.10350 , 2021.

Rashidinejad, B. Zhu, C. Ma, J. Jiao, and S. Russell. Bridging offline reinforcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems , 34: 11702-11716, 2021.

Rigollet and J.-C. Hütter. High dimensional statistics. Lecture notes for course 18S997 , 813(814): 46, 2015.

Saad-Falcon, A. Singh, L. Soldaini, M. D'Arcy, A. Cohan, and D. Downey. Embedding recycling for language models, 2023.

Schuster, A. Fisch, J. Gupta, M. Dehghani, D. Bahri, V. Q. Tran, Y. Tay, and D. Metzler. Confident adaptive language modeling, 2022.

Sharir, B. Peleg, and Y. Shoham. The cost of training NLP models: A concise overview. arXiv preprint arXiv:2004.08900 , 2020.

Shuja, K. Bilal, W. Alasmary, H. Sinky, and E. Alanazi. Applying machine learning techniques for caching in next-generation edge networks: A comprehensive survey. Journal of Network and Computer Applications , 181:103005, 2021.

J. Smith. Cache memories. ACM Computing Surveys (CSUR) , 14(3):473-530, 1982.

Stallings and G. K. Paul. Operating Systems: Internals and Design Principles , volume 9. Pearson New York, 2012.

Wang. A survey of web caching schemes for the internet. ACM SIGCOMM Computer Communication Review , 29(5):36-46, 1999.

Wei, Y. Tay, R. Bommasani, C. Raffel, B. Zoph, S. Borgeaud, D. Yogatama, M. Bosma, D. Zhou, D. Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682 , 2022a.

Wei, J. Qi, and S. He. A flexible multi-task model for BERT serving, 2022b.

Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 , 2022.

Zhou, Q. Li, C. Li, J. Yu, Y. Liu, G. Wang, K. Zhang, C. Ji, Q. Yan, L. He, H. Peng, J. Li, J. Wu, Z. Liu, P. Xie, C. Xiong, J. Pei, P. S. Yu, and L. Sun. A comprehensive survey on pretrained foundation models: A history from BERT to ChatGPT, 2023.

M. Ziegler, N. Stiennon, J. Wu, T. B. Brown, A. Radford, D. Amodei, P. Christiano, and G. Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.08593 , 2019.

Appendix

A Discussions on the Choice of Output, Model and Cost

The proposed framework is flexible in the choice of outputs, models and costs. Below we discuss several possible choices and combinations of output, models and costs that are most practically relevant.

Per-token Output and Per-sentence Output. We have two design choices of the desired output in each round, namely per-token output and per-sentence output.

For per-token output, we aim at generating one token at each round as a response of the queries. In this case, we only cache the next token for a given query and estimate the cost for generating next token. We also have the flexibility of choosing different models to generate each token in each round.

For per-sentence output, we aim at generating a complete response at each round. In this case, we cache the whole responses for a given query, and estimate the cost for generating the whole responses. This may introduce more variance in the cost due to the variation and randomness in the length of the generated responses.

Choices of Costs The cost can be chosen as FLOPS, latency of the model, the price for API calls, user satisfaction of the results, or a combination of all the four factors.

Model multiplexing A common choice of model ensembles is a pair of small and large models. The cost for small model C s ( q ) can be written as C s ( q ) = C s, 0 ( q ) + Y ( q ) C s, 1 ( q ) . Here Y ( q ) is a binary random variable, indicating whether the small model outputs satisfying results ( Y ( q ) = 0 ) or not ( Y ( q ) = 1 ). In the case when the small model outputs a satisfying response, the incurred cost is C s, 0 ( q ) . In the case when the small model outputs a bad response, the incurred cost is C s, 0 ( q ) + C s, 1 ( q ) . We discuss two possible choices of Y ( q ) , C s, 0 ( q ) and C s, 1 ( q ) based on two different evaluation pipeline as below.

One-time evaluation pipeline. For the one-time evaluation pipeline, we can only call one of the models once and the generated content cannot be changed. In this case, C s, 0 ( q ) can be set as the cost for running the small model to generate responses, Y ( q ) is set to be 1 if the user is not satisfied with the response, and C s, 1 ( q ) is the incurred cost for unsatisfactory of the user. One can similarly set the same cost for the large model.

Correction-based evaluation pipeline . For correction-based evaluation, we may regenerate the content with a different model if it is unsatisfying, and get an extra cost for fixing the content. Such evaluation can be easily combined with LLM Cascade (Frantar et al., 2023) or the idea from Big Little Transformer Decoder (Kim et al., 2023) and Spectulative sampling (Chen et al., 2023a). For example, after running the small model, we run the large model once to infer all the log probabilities of the small model output in parallel, and reject its output if the log probabilities are low. If the small model output is rejected, we will set Y ( q ) = 1 and run large model to re-generate the responses. In this case, C s, 0 ( q ) is the cost of running the small model for generating responses, and running the large model once for checking the probability. And C s, 1 ( q ) is the cost of running the large model to generate the response.

We also remark here that in the special case when the cost for the small model is much smaller than that of the large model under the correction-based evaluation pipeline, the cascade selector which always runs the small model first may give better performance than the model multiplexer if the accuracy of the model multiplexer is low, since running small model does not introduce too much cost compared to running large model. In this situation, the cascade selector can also be combined with LEC caching to further improve the performance.

On the other hand, we may also choose among models with similar size but different expertise, including coding, summarization and chat etc. In this case, we also expect to see different qualities and cost of responses for specific queries.

B Generalization to Variable Size Cache

For the variable-size caching problem, assume that the cache size of q is a deterministic scalar, denoted as S ( q ) . In the population case we design the cache as follows:

In the case when all S, P, c ⋆ s , c ⋆ l are known, one may solve the above constrained optimization problem for the optimal caching. When S ( q ) ≪ L , a good cache replacement algorithm is GDSF itself, which replaces the query with the smallest expected cost per-size P ( q ) min ( c ⋆ s ( q ) , c ⋆ l ( q )) /S ( q ) rather than expected cost per-query.

A more practical setting is the case when the cache size for each query S ( q ) is a random variable. Due to the randomness in the generation procedure, we expect to see responses of different lengths even when we use the same model to process the same query. In each round, we will have a generated response with size s ( q ) that is sampled from the random variable S ( q t ) . We conjecture that the optimal cache replacement algorithm is to replace the query with the smallest expected cost per-size P ( q ) min ( c ⋆ s ( q ) , c ⋆ l ( q )) /s ( q ) as well, where s ( q ) is the size of the cached queries and responses.

C Generalization to Multiplexing of Multiple Models

The proposed algorithm can be generalized to model multiplexing with multiple models. Assume that we have K models, and each model has a random cost function C k ( q ) with expectation c ⋆ k ( q ) . In this case, the optimal population algorithm is

And the finite sample algorithm is natural to follow. In practice, one may train a neural network with K dimensional output to predict the cost for each of the models.

D Differences Between the Optimal Policy and the Baseline

Consider the population setting in Section 3, where we optimize caching without model multiplexing. We show via a simple example below that without considering the cost for individual query, LFU can be highly sub-optimal compared to the optimal caching strategy in the population. The ratio

Proposition 1. For any fixed cost function c ⋆ l , one can design some distribution of queries P such that for any ϵ > 0 ,

The construction can be seen from a two-query example. Let c ⋆ l ( q 1 ) = c 1 , c ⋆ l ( q 2 ) = c 2 with c 1 < c 2 . Let P ( q 1 ) = , P ( q 2 ) = This shows that when the individual cost varies drastically for different queries, the total expected cost for LFU can be highly sub-optimal compared with the cost-aware caching strategy.

To compare the performance of the model multiplexing in Section 4, we take the cache size L = 0 . We have the following proposition for the performance improvement of the model multiplexer.

Proposition 2. Let L = 0 . The difference in cost between the baseline and the model multiplexer can be written as

The proof is a direct result of plugging in the cost definition. We see that the gap between π s and the optimal model multiplexer becomes larger when a large fraction of the queries have smaller cost when processed by the large models, and vice versa.

E Proof of Theorem 1

Proof. We first prove the following lemma on the lower bound of P ( q ) for any q ∈ L ⋆ .

Lemma 1. For any q ∈ L ⋆ , we have P ( q ) ≥ B 1 / ( B 2 |Q| ) .

Proof. From the fact that ∑ q ∈Q P ( q ) = 1 and for any q ∈ L ⋆ and any q ′ ̸∈ L ⋆ , P ( q ) ≥ P ( q ′ ) c ⋆ l ( q ′ ) /c ⋆ l ( q ) ≥ P ( q ′ ) B 1 /B 2 , we know that for any q ∈ L ⋆ , P ( q ) ≥ B 1 / ( B 2 |Q| ) .

We define the following three events:

We know that the first two events hold simutaneously with probability at least 1 -2 δ/ 3 from Lemma 3. For the third event, from the Chernoff bound, we know that for any q ∈ Q , we have

From Lemma 1 we know that for any q ∈ L ⋆ , P ( q ) ≥ B 1 / ( B 2 |Q| ) . Thus the above inequality further implies

The last inequality is due to our assumption that N ≥ 8 B 2 |Q| log(3 L/δ ) B 1 .

We condition on the three events from now on. The last two events imply that for any q ∈ L ⋆ ,

We have

Let ˆ c l,pes ( q ) = ˆ c l ( q ) -( B 2 -B 1 ) √ log(6 |Q| /δ ) 2 ∑ N n =1 ( q n = q ) . Note that for any q ∈ L ⋆ , we know that

And similarly, for any q ̸∈ L ⋆ , we know that

Now consider any q ∈ ˆ L but q ̸∈ L ⋆ , and any other q ′ ∈ L ⋆ but q ′ ̸∈ ˆ L . We have

Overall, we know that conditioned on E 1 ∩ E 2 ∩ E 3 , we have

And this implies that

Taking δ = 1 /N finishes the proof.

F Proof of Theorem 2

Proof. Upper Bound. We start with the upper bound by the following lemma.

Lemma 2. In each round t ∈ [ T ] , we always have

Proof. We prove this lemma by induction. First, consider the case when |L t +1 | < L . In this scenario, we always put the query into the cache. And L t +1 contains all queries with non-zero ˆ P t . Thus such L t +1 is always one of the minimizers.

Now consider the case when |L t +1 | = L . Assume that the conclusion holds for time step t . Now consider the case of t +2 . When the new query is in the cache q t +1 ∈ L t +1 , the cache will remain unchanged L t +2 = L t +1 . In this case, the estimated probability for q t +1 is increased, while the others are decreased, and ˆ c l,t +1 is not changed for any query. Thus L t +2 is still the minimizer. When the new query does not hit the cache, the estimated probability times costs for all other queries except for q t +1 are decreased proportionally since ˆ P t +1 is decreased proportionally while ˆ c l,t +1 is not changed for all other queries. Thus the only potential change in the relative order of costs is that of q t +1 . Since we can add q t +1 at the end of query, we know that after this round L t +2 is still the minimizer.

Let g k ( q ) be the length of the interval between the k -th and k +1 -th arrival of query q in the sequence of received queries (we set g k ( q ) = 0 if k exceeds the total number of times q is queried.). Define the following three events:

̸

We prove that the three events hold simultaneously with probability at least 1 -δ :

Lemma 3. We have

Proof. From the Dvoretzky-Kiefer-Wolfowitz inequality, we have

By taking ϵ = √ 2 log(6 T/δ ) t , we see that max q ∈Q | ˆ P t ( q ) -P ( q ) | ≤ ϵ holds with probability at least 1 -δ/ (3 T ) for any fixed t ∈ [ T ] . Now by taking a union bound over all t ∈ [ T ] , we know that ⋂ T t =1 E 1 ,t holds with probability at least 1 -δ/ 3 .

For the second event, from Hoeffding's inequality, we have for any q ∈ Q ,

̸

✶ Now taking union bound over t ∈ [ T ] and q ∈ Q gives that ⋂ T t =1 E 2 ,t holds with probability at least 1 -δ/ 3 .

For the third event, we know that the interval probability

P

(

q

)

. For any

q

∈ L

⋆

, we have

By taking s = B 2 |Q| log( TL/δ ) B 1 , we know that

By taking union bounds over all q ∈ L ⋆ and k we get the result.

Let E t = ⋂ t s =1 E 1 ,s ∩ E 2 ,s . We can write the regret as follows.

Let T t ( q ) = ∑ t -1 i =1 ✶ ( q i ̸∈ L i , q i = q ) . Note that the event c i = × is equivalent to that q i ∈ L i . Now at each round t , conditioned on event E t , we know that for any q ∈ L ⋆ ,

g

k

(

q

)

satisfies a geometric distribution with success

And similarly, for any q ̸∈ L ⋆ , we know that

Now consider any q ∈ L t but q ̸∈ L ⋆ , and any other q ′ ∈ L ⋆ but q ′ ̸∈ L t . We have

Thus we know that

Thus we have

Now for each q ∈ L ⋆ , we look at the term ∑ T t =1 E [ ✶ ( q ̸∈ L t ) · min ( 1 , √ 1 T t ( q ) ) | E t ∩ E 3 ] . We prove the following lemma:

Proof. Let t k ( q ) = ∑ k -1 l =1 g l ( q ) be the step that the k -th query of q arrives, with t 0 ( q ) = 0 . And let E = ( ⋂ T t =1 E t ) ∩ E 3 . The summation can be written as

The last inequality is due to (a) T t ( q ) does not change if at round t the query is not q ; (b) if q ∈ L t k +1 ( q ) , we will have q ∈ L t for any t ∈ [ t k ( q ) + 1 , t k +1 ( q )] since q never arrives in the middle and must remain in the cache set until t k +1 ( q ) . Now from event E 3 , we know that

We know that T t k +1 ( q )+1 ( q ) = T t k +1 ( q ) ( q ) + 1 = T t k ( q )+1 ( q ) + 1 if q ̸∈ L t k +1 ( q ) since the query q missing the cache will be sent to the model. Thus overall, we know that we have either T t k +1 ( q )+1 ( q ) = T t k ( q )+1 ( q ) + 1 , or ✶ ( q ̸∈ L t k +1 ( q ) ) · √ 1 T t k ( q )+1 ( q ) = 0 and T t k +1 ( q )+1 ( q ) = T t k ( q )+1 ( q ) . Thus overall, we have

By taking δ = 1 /T , we know the final regret can be bounded by

Lower bound. Now we turn to the lower bound. We apply Le Cam's two point lemma for the regret. Consider any family of algorithm {L t } T t =1 , where L t can be dependent on observations prior to time step t . We aim to design two instances with the same P ( q ) and different random variable C l ( q ) such that for any algorithm, the incurred cost for one of the instance is at least Ω( √ T ) . Consider the case when we only have two candidate queries Q = { q 1 , q 2 } . Set P ( q 1 ) = P ( q 2 ) = 1 / 2 for both instances and the cache size L = 1 . For instance one, we let C (1) l ( q 1 ) ∼ Bern (1 / 2) , C (1) l ( q 2 ) ∼ Bern (1 / 2 + ∆) . For instance two, we let C (2) l ( q 1 ) ∼ Bern (1 / 2) , C (2) l ( q 2 ) ∼ Bern (1 / 2 -∆) . We have

Let Regret (1) ( T ) be the total regret when C l = C (1) l , and Regret (2) ( T ) be the total regret when C l = C (2) l . Then we can verify that for any sequence of L t ,

Thus from Le Cam's Lemma, we have

Here E 1 [ T 2 ] is the expected times of observing the cost of q 2 under instance one. Taking ∆ = T -1 / 2 and minimizing the above equation with E 1 [ T 2 ] gives the desired bound.

G Proof of Theorem 3

Proof. We define the following four events:

We know that the above events hold simultaneously with probability at least 1 -δ from Lemma 3. We condition on the four events from now on. We first decompose the cost difference as

The first difference can be further written as

Note that if ˆ c s ( q ) -ˆ c l ( q ) has the same sign as c ⋆ s ( q ) -c ⋆ l ( q ) , the difference c ⋆ s ( q ) ✶ (ˆ c s ( q ) ≤ ˆ c l ( q )) + c ⋆ l ( q ) ✶ (ˆ c s ( q ) > ˆ c l ( q )) -min( c ⋆ s ( q ) , c ⋆ l ( q )) becomes 0 . Otherwise, if c ⋆ s ( q ) -c ⋆ l ( q ) > 0 , we know that

And similarly if c ⋆ s ( q ) -c ⋆ l ( q ) ≤ 0 , we know that c ⋆ l ( q ) -c ⋆ s ( q ) ≤ | ˆ c s ( q ) -c ⋆ s ( q ) | + | ˆ c l ( q ) -c ⋆ l ( q ) | . Overall, we have

Here ( i ) is due to Cauchy-Schwarz, and (ii) is from Jensen's inequality, and (iii) is the standard rate of the least squared estimator (Rigollet and Hütter, 2015).

For the second difference, we have

The last inequality uses event E 3 and Lemma 1. And similarly, for any q ̸∈ L ⋆ , we know that

Now consider any q ∈ L t but q ̸∈ L ⋆ , and any other q ′ ∈ L ⋆ but q ′ ̸∈ L t . We have

Here the last inequality uses the fact that q is inside L t and thus the difference between the first two terms are upper bounded by 0 . Finally, we know that conditioned on E 1 ∩ E 2 ∩ E 3 ∩ E 4 , we have

Overall, we know that

Taking δ = 1 /N finishes the proof.

H Proof of Theorem 4

Proof. Let g k ( q ) be the length of the interval between the k -th and ( k +1) -th arrival of query q in the sequence of received queries (we set g k ( q ) = 0 if k exceeds the total number of times q is queried.). Define the following four events:

From the same analysis as Lemma 3, we know that the four events ( ⋂ t s =1 E 1 ,s ∩ E 2 ,s ∩ E 3 ,s ) ∩ E 4 hold simultaneously with probability at least 1 -δ .

Let E t = ⋂ t s =1 E 1 ,s ∩ E 2 ,s ∩ E 3 ,s . The regret can be decomposed as follows. Regret ( T ) = T ∑ t =1 E [ cost ( q t , L t , π t ) -cost ( q t , L ⋆ , π ⋆ )] ≤ T ∑ t =1 E [ cost ( q t , L t , π t ) -cost ( q t , L ⋆ , π ⋆ ) ✶ ( E t )] + E [ cost ( q t , L t , π t ) -cost ( q t , L ⋆ , π ⋆ ) ✶ ( ¯ E t )] ≤ T ∑ t =1 E [ cost ( q t , L t , π t ) -cost ( q t , L ⋆ , π ⋆ ) ✶ ( E t )] + δTB 2 = T ∑ t =1 E [( cost ( q t , L t , π t ) -cost ( q t , L t , π ⋆ ) + cost ( q t , L t , π ⋆ ) -cost ( q t , L ⋆ , π ⋆ )) ✶ ( E t )] + δTB 2 . The first difference can be further written as

Note that if ˆ c s,t ( q t ) -ˆ c l,t ( q t ) has the same sign as c ⋆ s ( q t ) -c ⋆ l ( q t ) , the difference c ⋆ s ( q t ) ✶ (ˆ c s,t ( q t ) ≤ ˆ c l,t ( q t )) + c ⋆ l ( q t ) ✶ (ˆ c s,t ( q t ) > ˆ c l,t ( q t )) -min( c ⋆ s ( q t ) , c ⋆ l ( q t )) becomes 0 . Otherwise, if c ⋆ s ( q t ) -c ⋆ l ( q t ) > 0 and ˆ c s,t ( q t ) -ˆ c l,t ( q t ) ≤ 0 , we know that s t = 1 and

Here the last inequality uses the fact that for each q ∈ Q , the summation over time step is upper bounded by 2 ∑ T ( q ) i =1 √ log(8 T |Q| /δ ) / 2 i ≤ 2 √ log(8 T |Q| /δ ) T ( q ) , where T ( q ) is the number of steps of receiving query q in total T steps. Optimizing over T ( q ) gives the final bound.

For the second difference, we have

Below we justify the last inequality. First, note that π t ( q ) = 1 is equivalent to that ˆ c s,t ( q ) ≤ ˆ c l,t ( q ) . Thus if ˆ c s,t ( q t ) -ˆ c l,t ( q t ) has the same sign as c ⋆ s ( q t ) -c ⋆ l ( q t ) , the above inequality holds. Now consider the case when ˆ c s,t ( q t ) -ˆ c l,t ( q t ) has a different sign as c ⋆ s ( q t ) -c ⋆ l ( q t ) . Assume that ˆ c s,t ( q t ) > ˆ c l,t ( q t ) and c ⋆ s ( q t ) < c ⋆ l ( q t ) . We know that π t ( q ) = 0 , and

Now consider any q ∈ L t but q ̸∈ L ⋆ , and any other q ′ ∈ L ⋆ but q ′ ̸∈ L t . We have

P ( q ′ ) min ( c ⋆ s ( q ′ ) , c ⋆ l ( q ′ )) -P ( q ) min ( c ⋆ s ( q ) , c ⋆ l ( q )) ≤ ˆ P t ( q ′ ) min (ˆ c s,t ( q ′ ) , ˆ c l,t ( q ′ )) -ˆ P t ( q ) min (ˆ c s,t ( q ) , ˆ c l,t ( q )) + C ( B 2 -B 1 ) · min ( 1 , √ log(8 |Q| /δ ) 2 ∑ t -1 i =1 ✶ ( s i = π t ( q ) , q i = q ) ) + B 2 √ 2 log(8 /δ ) t ≤ C ( B 2 -B 1 ) · min ( 1 , √ log(8 |Q| /δ ) 2 ∑ t -1 i =1 ✶ ( s i = π t ( q ) , q i = q ) ) + B 2 √ 2 log(8 /δ ) t . Here the last inequality uses the fact that q is inside L t . Thus we have

✶ ✶

Lemma 5. We have

Proof. Let t k ( q ) = ∑ k -1 l =1 g l ( q ) be the step that the k -th query of q arrives. And let E = ( ⋂ T t =1 E t ) ∩ E 4 . The summation can be written as

The last inequality is due to (a). T t ( q, π ) does not change if at round t the query is not q ; (b). if q ∈ L t k +1 ( q ) , we will have q ∈ L t for any t ∈ [ t k ( q ) + 1 , t k +1 ( q )] since q never arrives in the middle and must remain in the cache set until t k +1 ( q ) ; (c) For t ∈ [ t k ( q ) + 1 , t k +1 ( q )] , π t does not change since both the frequency and cost estimator does not change for q . From the definition of t k ( q ) , we know that T t k ( q )+1 ( q, π ) ≥ 1 and thus we can drop the the minimum in the above equation. Now from event E 4 , we know that

We know that T t k +1 ( q )+1 ( q, π ) = T t k +1 ( q ) ( q, π ) + 1 = T t k ( q )+1 ( q, π ) + 1 if q ̸∈ L t k +1 ( q ) and π t k +1 ( q ) = π since the query q missing the cache will be sent to one of the models, and only the one selected will observe the cost. Thus overall, we know that we have either T t k +1 ( q )+1 ( q, π ) =

Thus we know that the second difference satisfies

Overall by taking δ = 1 /T , we know that

I Additional Experiments

I.1 Synthetic Datasets

We conduct synthetic online and offline experiments for joint optimization of caching and model multiplexing. We use i.i.d. Bernoulli distributions for two models because we want to mimic the model ensemble use case and give a large penalty to the wrong output. In Figure 2, we plot the cumulative cost and regret in online learning for LFU and LEC caching algorithms. We present more data points in Table 3 and Table 4, under the same setting of 10000 requests with 20 different queries and cache size 10. Similar to the real dataset setting, we compare all combinations of caching strategy choices and model multiplexer choices. We consider the frequency distribution as power distribution with α = 0 . 5 and 0 . 8 . The ground truth cost for each query processed by both models is set as a sample from r · X +1 , where r is called as cost ratio and X is a random variable generated from a Bernoulli distribution with the parameter 0 . 5 . We consider the model multiplexer accuracy with 0 . 8 and 1 . We repeat the simulation 1000 times and take the mean. Consistent with Figure 2, our simulation suggests that LEC with a perfect model multiplexer significantly improves the baselines when the cost ratio is large. Simulation 1000 times cannot remove all randomness so we can observe some fluctuations. Theoretically, the columns of choosing model 1 and the columns of choosing model 2 should behave similarly.

I.2 Real Datasets

In this section, we provide additional experiments on real-world dataset. We run both offline and online algorithms on Lambada dataset and OpenAssistant dataset, with OPT-1.3B vs OPT-13B or FastChat-T5-3B vs Vicuna-13B. We mainly consider three settings:

In Table 5-12, we consider distinct prompt size 100 , total query size 10000 , cache size 40 , same as the experiments in the main text.

In Table 13-20, we consider distinct prompt size 1000 , total query size 2000 , cache size 100 .

In Table 21-28, we consider distinct prompt size 1000 , total query size 2000 , cache size 0 , same as the experiments in the main text.

Our proposed algorithm mostly gives dominant results over other baselines. In Table 21-28, the cache size is set to be 0 so there is no difference between LFU and LEC.

α

0.2

0.5

0.8

0.2

0.5

0.8

α

0.2

0.5

0.8

LFU+

large

9.31

20.04

28.24

LFU+

cascade

13.88

29.88

42.12

LFU+

selector

7.24

15.11

21.14

LEC+

large

8.74

18.68

26.07

LEC+

cascade

8.82

16.90

20.31

LEC+

selector

5.93

11.87

15.49

selector accuracy

0.8

0.8

0.8

1

1

1

LFU+

large

3.49

7.71

10.81

3.49

7.71

10.81

LFU+

cascade

3.81

8.43

11.80

3.81

8.43

11.80

LFU+

selector

2.60

5.76

8.06

1.91

4.22

5.90

LEC+

large

3.44

7.54

10.36

3.44

7.54

10.36

LEC+

cascade

1.50

3.09

4.11

1.50

3.09

4.11

LEC+

selector

2.00

3.94

4.76

0.99

1.97

2.50

α

0.2

0.5

0.8

0.2

0.5

0.8

α

0.2

0.5

0.8

0.2

0.5

0.8

selector accuracy

0.8

0.8

0.8

1

1

1

LFU+

large

1.54

3.39

4.74

1.54

3.39

4.74

LFU+

cascade

1.55

3.42

4.82

1.55

3.42

4.82

LFU+

selector

1.09

2.41

3.38

0.79

1.75

2.45

LEC+

large

1.01

2.08

2.61

1.01

2.08

2.61

LEC+

cascade

0.56

1.10

1.36

0.56

1.10

1.36

selector accuracy

0.8

0.8

0.8

1

1

1

LFU+

large

8.64

19.07

26.70

8.64

19.07

26.70

LFU+

cascade

12.86

28.42

39.91

12.86

28.42

39.91

LFU+

selector

8.06

17.81

24.98

6.28

13.86

19.44

LEC+

large

8.03

17.66

24.53

8.03

17.66

24.53

LEC+

cascade

7.76

15.63

19.01

7.76

15.63

19.01

LEC+

selector

0.66

1.31

1.62

0.38

0.76

0.93

LEC+

selector

6.73

14.14

18.12

4.86

10.43

13.80

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents an innovative approach to mitigating the challenges posed by large language models (LLMs), namely high resource consumption and latency, by employing a cache to store previous queries and a model selector to choose the most efficient model for query processing. The authors propose an optimal algorithm that jointly optimizes both these approaches in offline and online tabular settings, and demonstrate its effectiveness through a series of simulations and real-world dataset experiments. The paper concludes by suggesting potential future research directions in caching and model selection optimization.\n\nSTRENGTHS:\n1) This paper introduces a novel approach to dealing with the resource-intensive nature of LLMs by jointly optimizing the usage of a cache and a model selector. It creatively merges established concepts of caching with model selection.\n\n2) The theoretical grounding and empirical validation of the proposed algorithm underline the quality of the research. The authors have effectively leveraged the Greedy Dual Size with Frequency (GDSF) and Least Expected Cost (LEC) caching algorithms alongside a model selector to achieve optimal rates.\n\n3) The paper is well-structured and effectively communicates complex ideas and methodologies in a clear and comprehensive manner. The stepwise progression from basic principles to specific techniques is particularly commendable.\n\n4) Addressing the issues of resource consumption and latency in LLM deployment is of critical significance in modern AI. The proposed solutions have a potentially broad impact, improving efficiency and feasibility of real-world AI applications.\n\nWEAKNESSES:\n1) Model Selection Scope: The experiments seem to consider a limited number of model options. The approach may encounter difficulties when scaling up to situations where selection is to be made from thousands or millions of models, which is a realistic scenario in complex AI systems.\n\n2) Prompt Representation in Caching: While the paper does a good job of exploring how to insert into a cache, it doesn't delve into the challenge of representing each prompt in the cache. This is an important aspect of LLM serving, as a comprehensive strategy for hashing, differentiating knowledge, and specifying unit information is necessary for efficient caching.\n\n3) The idea and each problem is well known. I am not sure about the novelty of the approach. The application and problem is critical though.\n\nQUESTIONS:\n1)  How does the proposed algorithm scale when selecting from a large number of models, possibly in the order of thousands or millions? Can you provide any theoretical or experimental insight on this scenario?\n\n2) While you discuss inserting prompts into a cache, the paper does not fully address how prompts are represented in the cache. Could you elaborate on this aspect of your approach?\n\nLIMITATIONS:\nThe authors have touched upon potential future work related to this study, highlighting open problems that require further investigation. However, a discussion on the limitations of their approach and its potential negative societal impacts seems to be missing. An exploration of the potential risks associated with the approach, such as the possible amplification of biases if the model selector disproportionately selects certain models, would be beneficial. Further, a discussion on how this approach could contribute to increasing centralization and monopolization in AI due to the computational requirements of maintaining and selecting from large model ensembles could offer a balanced perspective.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-17"}

comments: SUMMARY:
* this paper proposes a theoretically optimal algorithm for efficient large-scale deployment of LLMs
* particularly; they combine caching and model selection to reduce the inference cost
* they provide theoretical guarantees for both caching w and w/o model selection

STRENGTHS:
* large scale models are deployed everywhere and the motivation to make inference more efficient is very strong
* the paper comes with strong theoretical guarantees
* the method's efficiency gains are demonstrated in relevant real-world settings using models up to 13B
* the idea to perform model selection on the fly in an online fashion seems clever

WEAKNESSES:
I have to admit that this paper is out of my depth. The considered (problem, solution) seems very convincing and relevant to me, but I'm not familiar with its literature and cannot identify any obvious weaknesses.

QUESTIONS:
* how can we assess whether caching is relevant in a particular setup? for example, when using publicly available datasets like OpenAssistant; what if there are no redundant queries?

LIMITATIONS:
yes, the authors discussed limitations and future work directions
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents an LLM inference framework design that aims to reduce inference cost by caching and model selection. This work first provides an optimal formulation for jointly optimizing both caching and model selection in both offline and online settings. Then, evaluations on simulations and two tasks show that proposed work achieves cost saving compared to baseline.\n\nSTRENGTHS:\n(1) The high-level idea of caching + model selection sounds reasonable for LLM inference.\n\n(2) Detailed formulation of the optimization problem.\n\nWEAKNESSES:\n(1) My most unsatisfying part of this work is that the problem that authors try to solve should not be simplified to “cost saving” only, but it should be formulated as “a cost-accuracy trade-off”. I do see that the authors acknowledge that it is a trade-off (e.g., “The need for fuzzy search.” In introduction, and in section 5.2 “If the small model is chosen but its result is wrong, the large model must be run and it will incur an additional penalty.”), but I really want the authors to investigate more about this instead of simply comparing on the cost. For example, for fuzzy search during caching, how much will it actually affect inference accuracy? For model selection, what if there are more than two models, and what if there is no way to know whether the inference result is correct or not (so that you cannot use multiple models). Ultimately, the results should be presented as a cost-accuracy Pareto curve, and the proposed work has to show how it advances the Pareto frontier compared to baseline (achieving better accuracy under same cost, or achieving same accuracy under less cost). The unbalance between formulation (5 pages) and experiments (1.5 pages) exacerbates this issue.\n\n(2) In introduction it says “For the fuzzy search problem, semantic search or vector-embedding-based ideas provide a systematic solution that includes embedding extraction and matching algorithms. To simplify the problem, we assume that there exists some semantic search oracle that can group the prompts with the same semantic meaning…”. To me this assumption oversimplify the caching problem. In order to demonstrate the feasibility of caching, it’s needed to actually build yourself a semantic search-based cache, and actually evaluate how it affects the accuracy-cost tradeoff.\n\n(3) In section 5.2 it says “We fine-tune a BERT-base model as the model selector by predicting whether the small model can give the correct result and achieve 80.2% accuracy.”. But I can’t find any further information such as the fine-tune hyperparameters, and whether this 80.2% accuracy is training accuracy or validation accuracy.\n\nQUESTIONS:\n(1) In section 5.2, the next-token prediction task uses FLOPs as the cost, and the chat assistant task uses latency as the cost. Is there a reason to use different cost definition here? Or why not evaluate both cost definitions in both cases?\n\nLIMITATIONS:\nThe authors acknowledged some limitations, such as “We leave the training of a better selector as future work.”.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the computation costs of large language models. Specifically, the paper proposes a framework where previous queries are cached and retrieved and where given a query to process, an appropriate model is selected to answer based on the query. The framework flow is the following: given a query in test time, we check whether the response can be retrieved from the cache. If so, we return from the cache. If not, we select one of two models based on the query, where a natural configuration is to have two models - a large more accurate model with large computation cost and a small less accurate model with smaller computation costs (the com. cost can be measured in various ways). \n\n\nCache:\nThe method estimates two oracle - the first oracle DenEstOracle estimates the probability distribution that a query q will be observed. This is useful in defining the cache, sense we want to cache the most frequent queries.  The second oracle is RegressionOracle which estimates the cost of processing a query. This is useful since we want to cache the queries which cost the most to process. Then, in the online setting, in time step q we can have a cache which in stores the most frequent queries which also have the most cost. Remember that the cache is finite, so formally, in time step t the cache which is denoted L_t holds the L caches such that for query q in L_t the value P_t(q) * Cl_t(q) is larger than all previously seen queries that are not in the cache. This basically means that the set of queries in the cache is optimal, in the sense that the estimated cost of encountering those queries is the largest given both the probability that we will encounter them and the cost of processing them. \n\nModel selector:\nNow, for the model selector, the method assume that the cost of a query is weighed sum of the C_0(q) + Y(q) * C_1(q), where C_0(q) is the compute cost of the small model assuming the user is satisfied with the results, C_1(q) is the cost if the user is not satisfied with the results and Y(q) is a random binary variable denoting whether or not the user is satisfied. Appendix A discusses the costs and model selection and considers options such as C_1(q) == cost of user dissatisfaction or C_1(q) == cost of re-running the large model on q where the latter is used in the online experiment presented in the paper.\n\nNow, in the online setting the model selector simply chooses the model with the smaller cost and the cache stores the queries with largest P(q) * min(C*_s(q), C*_l(q)) where C* is the estimated cost for the RegressionOracle.\n\nThe paper provides experimental results in synthetic, offline and online setting (online setting is done using the open assist dataset). The paper compares against baseline such as simple Least Frequently Used (LFU) cache as baseline cache and simple cascade (calling the large model if the result of the small model is un-satisfactory) as model selector. The paper compares the baselines against the proposed cache and model selector, separately and jointly.\n\nSTRENGTHS:\n* The paper addresses an important real-world problem in a novel manner with both a detailed theoretical discussion and practical experiments and results.\n* The presented framework is general and can be extended/adjusted to various use cases and practical applications.  \n* The paper is well written and easy to follow.\n\nWEAKNESSES:\nIMO, the experiments can benefit from further analysis, e.g.:\n* Performance/accuracy analysis as function of number of steps in the online setting.\n* What is the cost of running the model selector and how does that affect the performance of the system?\n* Can you provide some accuracy measure of the oracles? for example, in the online setting, how many steps are required so that the oracles would be properly estimated? What is the accuracy as a function of the step number?\n\nQUESTIONS:\nCan you discuss how to ensure the cache diversity? the top most probable queries might be very similar to each other. I understand that the method does not cacher queries, but \"query groups\" where similar queries are cached as one group (it is assumed in the paper that there exists some semantic search oracle that can group the prompts with the same semantic meaning). Still, due to the long tail distribution of queries, the bottom most frequent queries in the cache might be less similar to each other than the top ones, which might motivates grouping the queries differently, i.e. taking into account P(q) and varying the minimal distance to group two queries together by semantic search oracle that groups prompts close embedding. If the cache is not diversified enough, less common queries would be processed slower. This is especially true if we want to ensure the system is not biased towards specific user population.\n\nLIMITATIONS:\nThe authors did not specifically addressed limitations or potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates the combination of caching and model selection strategies to reduce the inference costs of large models (LLMs in particular). In the result, the authors propose a theoretically grounded algorithm that demonstrates promising results in practice.\n\nSTRENGTHS:\n1. The problem of the cost-effective LLM inference is in high demand.\n2. The paper is well written and easy to follow.\n3. Theoretical contribution seems novel and complete. \n4. Synthetic experiments are promising, especially when the ratio between min/max query costs is large.\n\nWEAKNESSES:\nIn my opinion, Section 5.2 lacks some important details and clarifications. It makes it difficult to estimate the practical contribution of the proposed approach. I highly recommend adding the detailed evaluation protocols for both tasks to the appendix. \n\n**Minor**\n* In Figure 2, please specify which plot corresponds to the offline/online setting.\n* In L132 and L152, one can introduce $\\mathcal{L}^{\\star}$ and $\\pi^{\\star}$ that appear in the corresponding equations.\n* Consider adding the cost measures to the captions in Tables 1,2 \n* Consider specifying the hardware specs used for the latency evaluation and corresponding std values in Table 2.\n* Typos: L168 \"last year\" -> \"last layer\"\n\nQUESTIONS:\nThe following questions will help to address my concerns about the empirical contribution in Section 5.2:\n\n* How is the performance measured for each task?\n* For a fair comparison, all methods within each row (Tables 1,2) should provide the same model quality. \n    * Does this hold true in both cases? If yes, how is it guaranteed for the chat assistant task?\n    * What is the target performance for each task?\n    * What is the criterion to quantify if the output of the small model satisfies or not (especially in the chat assistant task)?\n* In the offline setting, the predictor is not accurate enough to outperform \"cascade\" while it is not the case in the online setting. I am curious why it happens.\n    * Is the chat assistant task easier for the selector? Or is it caused by the difference in the online/offline settings?\n    * What selector is used in the online setting? How accurate is it at convergence? What portion of queries does the predictor observe until convergence?\n   * What if one tries the online and offline settings for the Lambada and chat assistant tasks, respectively?\n* How do the costs and gains depend on the cache size? What are the cache sizes in both settings? How are they selected?\n* It might be useful to report the \"large\", \"cascade,\" and \"selector\" costs without caching to understand the LFU/LEC gains better.\n\nLIMITATIONS:\nLimitations are well discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-22"}