SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ying Sheng
Ying.Sheng@stanford.edu
Lianmin Zheng
lmzheng@berkeley.edu
Clark Barrett
barrett@cs.stanford.edu
jordan@berkeley.edu
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 \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.
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 transformer-based 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
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.
Figure 1: A workflow for LLM-based inference with caching and model multiplexing.
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.
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  $\mathcal{Q} \subset \mathbb{R}^d$ . In the  $t$ -th round, a query  $q_t \in \mathcal{Q}$  is sampled from a fixed population distribution  $P \in \Delta(\mathcal{Q})$ . We maintain a small set of cache  $\mathcal{L}_t \subset \mathcal{Q}$  with  $|\mathcal{L}_t| \leq L$ . We say the query hits the cache if the query satisfies  $q_t \in \mathcal{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) = \mathbb{E}[C_l(q)]$  be the expected true cost of
processing the query  $q$ . The cost for a given query  $q$  and cache  $\mathcal{L}$  can be written as:
$$
\operatorname {c o s t} (q, \mathcal {L}) = \mathbb {1} (q \notin \mathcal {L}) \mathbb {E} [ C _ {l} (q) ] = \mathbb {1} (q \notin \mathcal {L}) c _ {l} ^ {\star} (q).
$$
By taking the expectation over the distribution  $q$ , we have the expected cost as
$$
\operatorname {c o s t} (\mathcal {L}) = \sum_ {q} P (q) \mathbb {1} (q \notin \mathcal {L}) c _ {l} ^ {\star} (q).
$$
In the offline learning setting, we collect an offline dataset and hope to learn a caching policy  $\widehat{\mathcal{L}}$  such that  $\mathrm{cost}(\widehat{\mathcal{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  $\mathcal{L}_t$ , we let the model process the query and receive a cost  $c_{t} \sim \mathbb{P}_{C_{l}}$ . Then we can choose to update the cache  $\mathcal{L}_t$  by adding the current query and response to the cache, and replacing one of the existing cached items if the cache  $\mathcal{L}_t$  is full. If the query hits the cache  $q_{t} \in \mathcal{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:
$$
\operatorname {R e g r e t} _ {\text {c a c h e}} (T) = \sum_ {t = 1} ^ {T} \mathbb {E} \left[ \operatorname {c o s t} \left(q _ {t}, \mathcal {L} _ {t}\right) - \operatorname {c o s t} \left(q _ {t}, \mathcal {L} ^ {\star}\right) \right].
$$
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^\star(q) = \mathbb{E}[C_s(q)]$  and  $c_l^\star(q) = \mathbb{E}[C_l(q)]$ .
Let  $\pi : \mathcal{Q} \mapsto [0,1]$  be the (possibly random) model multiplexing policy that maps the query  $q$  to values in  $[0,1]$ , where  $\pi(q) = 1$  represents that the query is always sent to the small model, and  $\pi(q) = 0$  represents the query is always sent to the large model. The randomness in the policy  $\pi$  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  $\mathcal{L}$  and policy  $\pi$ :
$$
\begin{array}{l} \operatorname {c o s t} (q, \mathcal {L}, \pi) = \mathbb {1} (q \notin \mathcal {L}) \mathbb {E} \left[ C _ {s} (q) \pi (q) + C _ {l} (q) (1 - \pi (q)) \right] \\ = \mathbb {1} (q \notin \mathcal {L}) \left(c _ {s} ^ {\star} (q) \pi (q) + c _ {l} ^ {\star} (q) (1 - \pi (q))\right). \\ \end{array}
$$
By taking the expectation over  $q$ , we have the expected cost as
$$
\operatorname {c o s t} (\mathcal {L}, \pi) = \sum_ {q} P (q) \mathbb {1} (q \notin \mathcal {L}) \left(c _ {s} ^ {\star} (q) \pi (q) + c _ {l} ^ {\star} (q) (1 - \pi (q))\right).
$$
In the offline learning setting, we collect an offline dataset and hope to learn a caching policy  $\widehat{\mathcal{L}}$  and a multiplexer  $\hat{\pi}$  such that  $\mathrm{cost}(\widehat{\mathcal{L}},\hat{\pi})$  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:
$$
\operatorname {R e g r e t} _ {\mathsf {s e l}} (T) = \sum_ {t = 1} ^ {T} \mathbb {E} \left[ \operatorname {c o s t} \left(q _ {t}, \mathcal {L} _ {t}, \pi_ {t}\right) - \operatorname {c o s t} \left(q _ {t}, \mathcal {L} ^ {\star}, \pi^ {\star}\right) \right].
$$
<sup>1</sup>Note that although we name the models as small and large models, we do not impose any assumption on the relationship between their costs. Moreover, the model size and cost function can be arbitrary for both models.
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}^{\star}$  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:
$$
\mathcal{L}^{\star} = \mathcal{L}_{\text{LEC}} = \operatorname *{arg  min}_{\mathcal{L}:|\mathcal{L}|\leq L}\operatorname {cost}(\mathcal{L}) = \operatorname *{arg  min}_{\mathcal{L}:|\mathcal{L}|\leq L}\sum_{q\in \mathcal{Q}}P(q)\mathbb{1}(q\notin \mathcal{L})c_{l}^{\star}(q).
$$
The traditional frequency-based caching strategy, including Least Recent Used (LRU) and Least Frequently Used (LFU), aims at caching the most frequent queries:
$$
\mathcal{L}_{\text{LFU}} = \operatorname *{arg  min}_{\mathcal{L}:|\mathcal{L}|\leq L}\sum_{q\in \mathcal{Q}}P(q)\mathbb{1}(q\notin \mathcal{L}).
$$
We show in Appendix D that the ratio between the cost of LFU and LEC can be as high as  $\frac{\max_{q\in\mathcal{Q}}c_l^* (q)}{\min_{q\in\mathcal{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  $\mathcal{L}$  based on prior data such that the introduced cost is minimized. Denote  $\mathcal{D}_N = \{(q_1,c_1),\dots ,(q_N,c_N)\}$ , where  $q_{i}$  is sampled from the distribution  $P(\cdot)$ , and  $c_{i}$  is a sample from random variable  $C_l(q_i)$ . We consider estimating  $P,c_l^{\star}$  from oracles  $\hat{P} = \mathrm{DenEstOracle}(q_1,\dots ,q_N)$ ,  $\hat{c}_l(q) = \mathrm{RegressionOracle}(\mathcal{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  $\hat{P}$  and  $\hat{c}_l(q)$  to be the plug-in estimator:
$$
\hat {P} (q) = \frac {\sum_ {i = 1} ^ {N} \mathbb {1} \left(q _ {i} = q\right)}{N}, \tag {1}
$$
$$
\hat {c} _ {l} (q) = \left\{ \begin{array}{l l} \frac {\sum_ {i = 1} ^ {N} \mathbb {1} \left(q _ {i} = q\right) c _ {i}}{\sum_ {i = 1} ^ {N} \mathbb {1} \left(q _ {i} = q\right)}, & \text {i f} \sum_ {i = 1} ^ {N} \mathbb {1} \left(q _ {i} = q\right) > 0 \\ B _ {1}, & \text {i f} \sum_ {i = 1} ^ {N} \mathbb {1} \left(q _ {i} = q\right) = 0. \end{array} \right. \tag {2}
$$
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^{\star}(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)<sup>2</sup>. As we show in Lemma 1, the true frequency for any query  $q \in \mathcal{L}^{\star}$  is lower bounded by some constant that depends on  $B_{1}, B_{2}, |\mathcal{Q}|$ . Thus the pessimism helps eliminate those less visited queries in the long tail of the distribution and encourages caching the queries in  $\mathcal{L}^{\star}$ . The lower-confidence-bound based estimator is:
$$
\hat {\mathcal {L}} = \underset {\mathcal {L}: | \mathcal {L} | \leq L} {\arg \min } \sum_ {q \in \mathcal {Q}} \mathbb {1} (q \notin \mathcal {L}) \hat {P} (q) \cdot \max  \left(B _ {1}, \left(\hat {c} _ {l} (q) - (B _ {2} - B _ {1}) \sqrt {\frac {\log (6 N | \mathcal {Q} | / \delta)}{2 \sum_ {n = 1} ^ {N} \mathbb {1} (q _ {n} = q)}}\right)\right).
$$
We show how the cost for the caching from the empirical estimate differs from the optimal cost.
Theorem 1. Assume that  $N \geq \frac{8B_2|\mathcal{Q}|\log(3L / \delta)}{B_1}$  and taking  $\delta = 1 / N$ . We have
$$
\mathbb {E} \left[ \operatorname {c o s t} \left(\hat {\mathcal {L}}\right) - \operatorname {c o s t} \left(\mathcal {L} ^ {\star}\right) \right] \leq C \left(B _ {2} - B _ {1}\right) L \cdot \sqrt {\frac {B _ {2} \left| \mathcal {Q} \right| \log (N \mid \mathcal {Q} |)}{N B _ {1}}}.
$$
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 / \sqrt{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.
2If we impose a uniform lower bound on the probability  $P(q)$ , then the pessimism can be replaced with the plug-in estimator. However, it is usually not the case in practice since  $P(q)$  usually comes with a long tail.
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  $\mathcal{L}_1 = \{\}$ , past observations  $\mathcal{H}_1 = \{\}$ ,  $\hat{c}_{l,0}(q) = B_1$ ,  $\forall q \in \mathcal{Q}$ .  
2: For iteration  $t = 1,2\dots ,T$   
3: Receive query  $q_t$ .  
4: Update the density estimation  $\hat{P}_t = \text{DenEstOracle}(q_1,\dots,q_t)$ .  
5: If  $q_t \in \mathcal{L}_t$ :  
6: Output the cached result, set  $\hat{c}_{l,t} = \hat{c}_{l,t-1}$ , update the past observation  $\mathcal{H}_t = \mathcal{H}_{t-1} \bigcup (q_t, \times)$ , and continue.  
7: Use the large model to process the query, and observe a cost  $c_t \sim \mathbb{P}_{C_l(q)}$ .  
8: Update the past observation  $\mathcal{H}_t = \mathcal{H}_{t-1} \bigcup (q_t, c_t)$ .  
9: Update  $\hat{c}_{l,t} = \text{RegressionOracle}(\mathcal{H}_t)$ .  
10: If  $|\mathcal{L}_t| < L$ :  
11: Let  $\mathcal{L}_{t+1}$  be the union of  $\mathcal{L}_t$  and  $q_t$ .  
12: Else if  $\hat{P}_t(q_t) \cdot \hat{c}_{l,t}(q_t) > \min_{q \in \mathcal{L}_t} \hat{P}_t(q) \cdot \hat{c}_{l,t}(q)$ :  
13: Replace the minimizer element of  $\hat{P}_t(q) \cdot \hat{c}_{l,t}(q)$  in the cache  $\mathcal{L}_t$  with  $q_t$  to get  $\mathcal{L}_{t+1}$ .
For theoretical analysis, we focus on the tabular case and define the oracles as follows:
$$
\hat {P} _ {t} (q) = \frac {\sum_ {i = 1} ^ {t} \mathbb {1} \left(q _ {i} = q\right)}{t}, \tag {3}
$$
$$
\hat {c} _ {l, t} (q) = \left\{ \begin{array}{l l} B _ {1}, & \text {i f} \sum_ {i = 1} ^ {t} \mathbb {1} (c _ {i} \neq \times , q _ {i} = q) c _ {i} \\ \max  \left(B _ {1}, \frac {\sum_ {i = 1} ^ {t} \mathbb {1} (c _ {i} \neq \times , q _ {i} = q) c _ {i}}{\sum_ {i = 1} ^ {t} \mathbb {1} (c _ {i} \neq \times , q _ {i} = q)} - (B _ {2} - B _ {1}) \sqrt {\frac {\log (6 T | \Omega | / \delta)}{2 \sum_ {i = 1} ^ {t} \mathbb {1} (c _ {i} \neq \times , q _ {i} = q)}}\right), & \text {o t h e r w i s e} \end{array} \right. \tag {4}
$$
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  $\delta = 1 / T$ , we have for some universal constant  $C$ :
$$
\operatorname {R e g r e t} _ {\text {c a c h e}} (T) \leq \frac {C L \left(B _ {2} - B _ {1}\right) B _ {2} | \mathcal {Q} | L \log^ {2} (T | \mathcal {Q} |)}{B _ {1}} \cdot \sqrt {T}.
$$
On the other hand, for any caching policy  $\{\mathcal{L}_t\}_{t=1}^T$ , there exist some cases of  $P(q)$ ,  $c_l^\star(q)$  such that for some universal constant  $C'$ ,
$$
\operatorname {R e g r e t} _ {\text {c a c h e}} (T) \geq C ^ {\prime} \sqrt {T}.
$$
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