Cache & Distil: Optimising API Calls to Large Language Models (2024)

Abstract: Large-scale deployment of generative AI tools often depends on costly API calls to a Large Language Model (LLM) to fulfil user queries, a process that also exposes the request stream to external providers. To curtail the frequency of these calls, one can employ a local smaller language model -a student- which is continuously trained on the responses of the LLM. This student gradually gains proficiency in independently handling an increasing number of user requests, a process we term neural caching. The crucial element in neural caching is a policy that decides which requests should be processed by the student alone and which should be redirected to the LLM, subsequently aiding the student’s learning. In this study, we focus on classification tasks, and we consider a range of classic Active Learning-based selection criteria as the policy. Our experiments suggest that Margin Sampling and Query by Committee bring consistent benefits over other policies and baselines across tasks and budgets.

## introduction large language models (llms) offer unique capabilities in understanding and generating human-like text. they have gained widespread use in a wide range of applications, such as assistive tools and entertainment bots. however, large models are often very challenging for all but a few companies and institutions to run on their infrastructure @xcite . meanwhile, smaller models typically under-perform in these applications, at least without additional fine-tuning on task-specific labelled data. consequently, many applications access llms via commercial apis despite the costs involved and the exposure of their entire request stream to the api providers. to minimise the costs and data exposure associated with calling the api, we propose to locally train a smaller private language model, which we refer to as student, on the llm's predictions and, as the student gets more accurate, it handles an increasing number of requests. the knowledge of the llm gets continuously distilled into the smaller model. we refer to this scenario as neural caching (see figure 1 ), as the student can be thought of as a smart cache. note though that the student not only remembers what the llm predicted but also generalises beyond these examples. the goal of this paper is to formalise the neural caching problem and investigate simple ways of approaching it. the key element in the neural caching scenario is the policy determining which requests the student processes independently. a good policy should weigh the expected immediate user benefit (i.e., if the llm is substantially more likely to make a correct prediction than the student) and the anticipated benefit for the student (i.e., whether the llm's prediction will aid in training the student). the latter underscores its relationship with active learning @xcite , although al is typically associated with soliciting human annotations. in particular, there is a similarity to online al @xcite , where new unlabelled data points arrive in a stream and are discarded immediately or sent to an annotator. however, online al tends to focus on maximising the accuracy of the final model (i.e. student in our terminology). in contrast, what matters in neural caching is the accuracy of the joint system (student, teacher, along with the policy) over its lifetime since this online accuracy reflects the average level of service offered to a user. despite the aforementioned differences with al, evaluating the existing al algorithms -specifically the example selection criteria -remains valuable given the maturity of the al field and the ease of implementation of some of the al methods. this study aims to achieve this, as well as to investigate the potential shortcomings of these methods. for instance, will the al methods end up selecting examples that are too challenging even for the llm? would learning from these noisy examples be detrimental to the student? answering these questions can inform future research on this practically significant scenario. in this work, our focus is specifically on classification tasks, as opposed to free text generation. many practical problems, such as routing user requests to relevant departments or answering questions about factual knowledge, can be framed as classification tasks. by confining our focus to classification, we can apply methods developed in al without modification. this also allows us to circumvent additional challenges tied to the automatic evaluation of text generation @xcite . our findings reveal the benefits of using albased policies such as margin sampling @xcite and query by committee @xcite . across datasets and budgets, these methods consistently outperform baselines, such as routing examples randomly or training the student at the very start. our analysis also reveals that the student appears robust to the noise introduced by an llm. we also analyse a simplified practical scenario where the student is not retrained and observe even greater improvements in online accuracy from using al-based policies. we release our code to encourage further work on this problem. 1 the key contributions of this work are: in neural caching, llm calls are optimised, while the student model is periodically retrained on the labels. we believe online knowledge distillation could play a key role in saving calls to expensive models. ## the neural caching problem the objective of neural caching is to optimise the usage of an llm in a scenario where labels need to be generated for a stream of inputs. as we get more predictions from the llm, a student model is trained on them. our goal is to achieve the highest level of service possible within a set budget of llm calls; hence, calling the llm serves both to attain high accuracy for the incoming input as well as to train a student model. to put it formally, our goal is to establish a mapping between elements in the input space x and the corresponding labels in the space y. we start with a student model s 0 , and we can access a teacher model t on demand. our task is to predict labels for a sequence of n examples (x 1 , . . we retrain the student model on the labels obtained from the llm every f processed requests. this simulates the situation where the number of requests is uniform in time, and there is a set time to retrain the model, e.g. at night. for simplicity and to follow the convention in al to retrain the model from scratch @xcite , every time we retrain the student model, we reset it to the original pre-trained model and then use parameter-efficient fine-tuning. although continual learning methods could be employed @xcite , we believe this is largely orthogonal to our primary focus on policies and resetting enhances the reproducibility of our analysis. importantly, we do not assume access to ground truth (or human annotation) at any point in learning to simulate a fully automatic scenario. for every new input x i , we use the student model s i/f to obtain the predicted label ŷs i . then, we have the option to request the label ŷt i from the teacher model (llm), which incurs a cost of c(x i ). finally, we return the label ŷi for x i : the teacher's label if requested or the student's otherwise. the processing of the n examples is subject to a budget constraint, where the total cost must not exceed a fixed budget b. we assess the effectiveness of our querying strategy based on the accuracy of our predicted label ŷi compared to the actual label y i (online accuracy) on the online examples. additionally, we measure the accuracy of the final student model s n/f on a test dataset (final accuracy). algorithm 1 describes the process. algorithm 1: pseudo-code for the neural caching algorithm with budget b, retraining frequency f , cost per query c, data from the llm d llm and an initial student ## experiments we first present our results and then their analysis. to report accuracy across budgets, we use the corresponding area under the curve (auc) divided by the budget range, thus obtaining an average accuracy. ## conclusions in this work, we have studied how instance selection criteria from al behave when they are used to decide in real time whether we should perform an llm call or use a student model that has been trained on previous llm predictions. in the scenario where we are not retraining the student model, margin sampling performs the best, across different datasets. in the scenario where we retrain the student model with some time periodicity, query by committee is the most robust option. while margin sampling outperforms the front-loading baseline on harder tasks, it is more sensitive to the initial budget spent to train the student model s 0 . we find that the embedding-based strategy (coreset) consistently performs poorly across different encoders; to our knowledge, it is the only llm caching approach that has been adopted by practitioners to date (e.g., gptcache @xcite ). we believe these types of strategies could be useful in certain contexts, e.g. multiple near-identical calls to an llm, a scenario which has not been the focus of this work. 3 our results suggest that (i) there is room for smart llm query allocation in the context of continuously distilling an llm into a student model and (ii) previous literature in active learning can transfer well to this setup. we believe that online knowledge distillation could play a key role in caching llms, saving unnecessary calls to expensive models, and training smaller private models.