Harnessing Query Heterogeneity for Cost-Effective Proactive Caching in LLM Inference

We thank all reviewers for their valuable comments.

$\**Q1: Clarification on the Assumption 2.**$ Assumption 2 is primarily required for our theoretical analysis. The assumption that the sampling probability of each query is less than $1/2$ comes from Lemma 8 in [2], and it is used in our proof to theoretically bound the waiting time for Algorithm 2 to receive the required query. In practical applications, we treat certain hot topics as queries that are frequently repeated to the LLMs. The probability of these repeated queries typically does not exceed $1/2$. For example, in an e-commerce LLM, the proportion of the most frequently asked queries, such as refund requests, accounts for about 30% of all queries. While this represents a high probability in practice, it is still less than $1/2$. On the other hand, it is clear that there can be at most one query with a probability greater than $1/2$, since we have $\\sum_{q\\in \\mathcal{Q}}P(q)=1$. We can easily identify this query and handle it separately, allowing the cache to still focus only on queries with a sampling probability less than $1/2$. In fact, based on our experiments, we found that whether this assumption holds does not significantly affect the performance of our experiments.

$\**Q2: Regret depends on$1/p^*$.**$ In the online setting corresponding to Algorithm 2, when the Oracle provides the optimal cache solution, the required query may not be available immediately. As a result, the cache replacement cannot happen instantly; algorithm must wait for the required query to arrive before performing the replacement. Therefore, in the online setting, the term $1/p^*$ is difficult to eliminate due to the update delay. As we discussed in Section 4.3, in some special cases, such as the scenario in [1], $1/p^*$ can be replaced by $N$. On the other hand, if we assume popularity follows a uniform distribution, we have $p^*=1/N$ and the regret becomes $\\widetilde{O}(MN\\sqrt{T})$, which aligns with the bound presented in [1].

$\**Q3: Memory cost in Algorithm 2.**$ We apologize for the misunderstanding in our expression. The precise expression should be that there is no need to store the answer information corresponding to all queries. Algorithm 1 needs to store the textual information of all queries and their corresponding answers, while Algorithm 2 does not require storing the textual information of all answers. The query’s textual information is typically used to allow the algorithm to recognize the sampled query and determine whether it is a cache hit. The textual information of the query’s answer enables the cache to directly return the saved answer text in the case of a cache hit, without requiring processing by the LLMs. Algorithm 1 requires storing the textual information of all answers to immediately replace query-answer pairs when updates occur. In contrast, Algorithm 2 only store the answer corresponding to the current query in the cache, without requiring storage for answers to queries not in the cache. At the same time, the textual information is much larger compared to the statistical data used for computation and updates, such as the count and estimate of each query. Therefore, reducing the memory usage of the textual information can help decrease the overall memory overhead. As a result, we can conclude that Algorithm 2 reduces memory costs. For the query textual information, we found that strategies using string hashing, such as BKDRHash [3], can be employed to create a hash table for the textual information of each query. By using this string hash function, the input string can be converted into a numerical value based on a specific encoding scheme, which can then be used to identify the current query, thereby reducing the storage required for the textual information.

$\**Q4,Q5: typo.**$ We appreciate the reviewers for pointing out the typographical errors in our work. We will make the necessary corrections in the revised version.

$\**Q6: Cache length.**$ This is a typographical error; it should be 'cache size', not referring to the number of queries that can store. In our paper, both "length" and "size" refer to the length of the query's textual information, which is also the memory size occupied by the query's text.

[1] Banghua Zhu, Ying Sheng, Lianmin Zheng, Clark Barrett, Michael I. Jordan, and Jiantao Jiao. 2024. On optimal caching and model multiplexing for large model inference. In Proceedings of the 37th International Conference on Neural Information Processing Systems (NIPS '23). Curran Associates Inc., Red Hook, NY, USA, Article 2577, 59062–59094.

[2] Li, S. and Zhang, S. (2018). Online Clustering of Contextual Cascading Bandits. Proceedings of the AAAI Conference on Artificial Intelligence, 32(1).

[3] J. Wang, "Github," 9 8 2015. https://github.com/JamzyWang/HashCollector.

Thanks for your valuable suggestions, this is a very meaningful question. Our problem shares similarities with OIM, particularly in treating the cache $M_t$ as a super arm. However, we identify key differences in action (super arm) selection between our online setting and OIM. In our online setting, the cache can only be updated using the currently arriving query and the queries already stored in the cache, imposing constraints on action selection. In contrast, OIM allows actions to utilize all base arms. Specifically, the key differences in action (super arm) selection are as follows:

$

M_{t+1} \subseteq M_{t} \cup \{q_t\} \text{ (Ours)}

$

$

M_{t+1} \subseteq \mathcal{S} \text{ (OIM)}

$

where $\\mathcal{S}$ is a fixed set of valid super arms for all queries.

Namely, in our problem, action across different rounds exhibit local dependence, which makes the problem more challenging. As a result, the analytical methods and conclusions from [1,2] cannot be directly applied to our setting. Currently, it remains challenging for us to develop a clear method to eliminate the influence of $1/p^*$.

[1] Combinatorial Multi-Armed Bandit and Its Extension to Probabilistically Triggered Arms. JMLR'16

[2] Improving Regret Bounds for Combinatorial Semi-Bandits with Probabilistically Triggered Arms and Its Applications. NeurIPS'17

We thank all reviewers for their valuable comments.

$\**Q1: Dependency on$1/p^*$(or$1/p_{min**).} Please refer to the detail discussion in our response to Q2 for Reviewer AHQS.

$\**Q2: Oracle for$M_t$and its relation to the knapsack problem.**$ We set the Oracle as a solver for the 0-1 knapsack problem. The Oracle is required to not only output the maximum value corresponding to the optimal knapsack but also provide the optimal knapsack itself. The algorithm then uses the optimal knapsack output by the Oracle to determine the current cache $M_t$ that should be selected. In our implementation, we adopt the classical dynamic programming approach and record the optimal knapsack found. The Oracle in our work provides the exact solution. While an approximate algorithm could be used to reduce computational overhead, it is crucial that the Oracle's solution remains stable. This is because, once the cache converges to an optimal state, any change in the Oracle's solution would trigger a cache update, forcing the system to converge to a new optimal cache. Such instability could prevent our algorithm from ever converging.

Thank you for your feedback. Analyzing of the influence of $1/p^*$ remains a challenge. In the online setting of Algorithm 2, the cache can only be updated using the currently arriving query and the queries already stored in the cache, imposing constraint on action selection, i.e. $M_{t+1} \subseteq M_{t}\cup \\{q_t\\}$. Due to this constraint, the algorithm faces challenges in exploring other queries, and it becomes more difficult for the cache to recover from a suboptimal solution. Our lower bound is derived directly from previous work [1], which can be considered a special case of our model. If the query popularity distribution is uniform, then $p^*=1/N$, and our regret upper bound aligns with the result in [1]. Addressing the challenges associated with the term $1/p^*$ and providing a tighter lower bound is part of our future work.

[1] Banghua Zhu, Ying Sheng, Lianmin Zheng, Clark Barrett, Michael I. Jordan, and Jiantao Jiao. 2024. On optimal caching and model multiplexing for large model inference. In Proceedings of the 37th International Conference on Neural Information Processing Systems (NIPS '23). Curran Associates Inc., Red Hook, NY, USA, Article 2577, 59062–59094.

We thank all reviewers for their valuable comments.

$\**Clarification on problem reduction to knapsack problem.**$ We appreciate the valuable feedback from the reviewers. Our problem setting can indeed be reduced to the 0-1 knapsack problem. We will correct this statement in the revised version.

$\**Q1: More detailed explanation of pessimism method.**$ As mentioned in Section 1, we use the lower confidence bound (LCB) method for estimation to prevent our cache policy from getting trapped in a suboptimal solution, a method also used in previous work [1]. The key idea behind LCB is to penalize queries that remain in the cache for too long, making them more likely to be replaced. For example, consider a query set $\\mathcal{Q}=\\{1, 2\\}$ with a cache size of one. If the optimal cache is $\\{1\\}$, and the algorithm inserts $2$ into the cache, the algorithm will always replace the query with the higher estimated value. In this case, only queries not in the cache have their cost estimates updated due to a miss, while those in the cache retain their estimates. Under LCB, when $1$ experiences a miss, its estimation increases as the confidence interval narrows, eventually surpassing $2$ and leading to its replacement. In contrast, with traditional UCB (optimism method), $1$'s estimate would decrease with each miss, always remaining lower than $2$, and thus never being replaced. This situation is referred to as the algorithm getting trapped in a suboptimal solution. A detailed discussion can be found in Section 4.1.

$\**Q2: Utility of averaging over multiple reward observations.**$ In our work, we primarily considered the impact of cost rather than reward, aiming to enhance the overall reward by reducing cost, so we did not consider the observation of reward here. Since cost and reward are complementary, we believe that optimizing for reward can also be applied to the problem we propose.

$\**Q3: The practicality of maintaining query count and cost.**$ Compared to Algorithm 1, Algorithm 2 primarily reduces memory usage associated with the textual information of the query's answer. In practice, the memory required to store the textual information of each query's answer is much larger than the memory needed to record the observation count of each query and store statistical information, such as its estimation and count information. Therefore, we assert that Algorithm 2 reduces memory overhead compared to Algorithm 1, making it more practical for real-world applications. On the other hand, storing the count and statistical information of queries is essential in bandit strategies, as this information is needed for the policy to compute and update effectively. Moreover, the memory overhead for this part is small.

$\**Q4: Guidance on cache size.**$ The size of the cache can be decided based on the user's actual needs, but it should be much smaller than the total size of all queries, as in practical applications, the query set can be very large. A larger cache can reduce the overall cost, but a larger cache also incurs greater additional storage overhead. On the other hand, if the cache is too small, it may result in insufficient improvement for reward. In our simulation experiments, we set the cache size to be twice the maximum value of the sum of the query length and its corresponding answer length, i.e. $2\\max_q{S(q)}$.

[1] Banghua Zhu, Ying Sheng, Lianmin Zheng, Clark Barrett, Michael I. Jordan, and Jiantao Jiao. 2024. On optimal caching and model multiplexing for large model inference. In Proceedings of the 37th International Conference on Neural Information Processing Systems (NIPS '23). Curran Associates Inc., Red Hook, NY, USA, Article 2577, 59062–59094.

We thank all reviewers for their valuable comments.

$\**Clarification on contributions.**$ Our main contribution is that we extend [1] to capture query size, which enables us to design an algorithm that addresses this problem and analyze its regret with variable query sizes. The involving of variable query size presents several challenges in our theoretical analysis, such as we can not determine how many queries are currently stored in the cache or to find the regular of value of the queries stored in the current cache. Despite these challenges, we demonstrate that our Online Stream Cache Bandits algorithm achieves a $\\tilde{O}(\\sqrt{T})$ regret bound with a $O(\\log T)$ additional computational cost. Since our work is an extension of [1] and the problem setting in [1] can be considered as a special case of ours, we directly used the conclusion of [1] as our lower bound.

$\**Q1: Notation mistake.**$ We apologize for the notation error in our work; it was a notation mistake, and the correct notation should be $\\Omega$. We will make the necessary correction in the revised version.

$\**Q2: Handling Randomness in$A(q)$due to Non-Deterministic LLM Outputs.**$ 1) In our setting, we assume that the length of each response to a query is fixed, as we focus on tasks with stable answer lengths for LLMs, such as machine translation and text classification. Generally, the answers from LLMs do not change significantly within a over time, so we primarily consider the case where the length of the answer remains fixed. 2) In cases where the size of answer to a query is random, [1] provides a brief discussion but no detailed analysis. We believe that the expected value of the answer size can be estimated and that this expected length can be used to define the optimal cache. To address random size, a good method is to use the expected value of the query length as the input to the Oracle and update the cache accordingly.

$\**Q3: Tightness with respect to parameters other than$T$.**$ The analysis for the tightness of other parameters in Algorithm 2 remains a challenge. In the online setting, the cache can replace at most one query per round, leading to delays when updating the cache to the recommended version. These delays are difficult to bound, so we have only conducted a rough analysis of the lower bound. As a result, we are currently unable to provide precise lower bound, but we plan to address this issue in future work.

[1] Banghua Zhu, Ying Sheng, Lianmin Zheng, Clark Barrett, Michael I. Jordan, and Jiantao Jiao. 2024. On optimal caching and model multiplexing for large model inference. In Proceedings of the 37th International Conference on Neural Information Processing Systems (NIPS '23). Curran Associates Inc., Red Hook, NY, USA, Article 2577, 59062–59094.