Gatekeeper to save COGS and improve efficiency of Text Prediction (2023)

Abstract: The text prediction (TP) workflow calls a Large Language Model (LLM), almost, after every character to get subsequent sequence of characters, till user accepts a suggestion. The confidence score of the prediction is commonly used for filtering the results to ensure that only correct predictions are shown to user. As LLMs require massive amounts of computation and storage, such an approach incurs network and high execution cost. So, we propose a Model gatekeeper (GK) to stop the LLM calls that will result in incorrect predictions at client application level itself. This way a GK can save cost of model inference and improve user experience by not showing the incorrect predictions. We demonstrate that use of a model gatekeeper saved approx 46.6% of COGS for TP, at the cost of approx 4.5% loss in character saving. Use of GK also improved the efficiency (suggestion rate) of TP model by 73%.

## introduction large language models (llms), such as generative pre-trained transformers (gpt-2, gpt-3) and turing natural language generation (t-nlg) models, have billions of parameters. these can be fine-tuned for various natural language processing (nlp) tasks, such as text classification, question answering and text prediction. our text editor application uses a distilled version of one such large text prediction model to provide text suggestions when user types in editor boxes. this improves users' writing productivity and reduces grammar and spelling errors. this application calls a large text prediction (tp) model after every keystroke (i.e. a character is typed) to show text completion suggestions. the last 256 character(s) typed by a user are sent to this model to get subsequent text prediction with confidence score. the editor application considers only the predictions that have a high confidence score. but, these confidence values are available only after model inference. as these models have a large number of parameters, they require large number of floating point operations (flops) for an inference. we perform another round of quantization aware distillation to reduce the latency and host it on cloud from where it is accessed by millions of users. such a large number of inferences incur high cost of goods sold (cogs). furthermore, it has been observed that they are typically overconfident in their predictions on out-of-distribution (ood) data @xcite . so, when the llm provides outputs for input examples that are far from distribution in training set (i.e. ood), their predictions can be arbitrarily bad. these false positives from model will reduce the reliability of application and result in poor user experience. to maintain user confidence, save unrewarding cogs and delay, we propose to have a model gatekeeper for llm. a model gatekeeper filters out the inputs for its large model. a model gatekeeper is small in size so that it can be used at edge to stop calls for model inferences for the inputs that may result in incorrect prediction. this provides average latency, performance and cost advantages for enterprises hosting large models. gatekeeper is a binary classification model trained using the large model's evaluation data. for a given input, gatekeeper predicts 0, if large model may return a valid prediction else predicts 1. we developed and evaluated gatekeeper for a large text prediction model using publicly-available data and internal data. we demonstrate that the model gatekeeper is capable of identifying relevant inputs for text prediction model. use of gatekeeper improved the suggestion rate (i.e. the percentage of times the server model was able to provide predictions with a confidence score higher than set threshold) by ≈ 70% and reduced the cogs by ≈ 47%. the reduction in inferences from large model resulted in better user experience and reduced cogs. ## related work multiple researchers @xcite , @xcite have established that softmax prediction probability is a good baseline for error and out-of-distribution (ood) detection across several architectures of deep neural networks (dnns). @xcite defined the confidence score as a maximum value of the predictive distribution. they demonstrated that, while these prediction probabilities create a consistently useful baseline, at times they are less effective. @xcite improved their performance by using temperature scaling (that uses a single scalar parameter @xmath0.e. raises the output entropy) with @xmath1. although such inference methods are computationally simple, they depend on how well the base model was trained. so, there has been a lot of effort in improving the training of base dnn models for better ood and uncertainty determination. @xcite utilizes temperature scaling with input perturbations using the ood validation dataset to tune hyper-parameters of base model. @xcite and @xcite proposed data augmentation methods to generate out-of-domain samples, then use them to train the base model for improved ood detection. @xcite proposed jointly training a generator and a classifier, the generator produces examples that appear to be at the boundary of the data manifold to serve as out-of-distribution examples, while the classifier is encouraged to assign these uniform class probabilities. @xcite and (devries and taylor, 2018) train neural networks that produce two outputs: a prediction and an uncertainty estimate. @xcite , @xcite proposed separate confidence estimation modules on top of end-to-end (e2e) models, which are classification models trained to minimize the binary cross entropy between the estimated confidence and the target. bayesian probabilistic models @xcite and ensembles of classifiers @xcite , learn a distribution over weights during model training for estimating the predictive uncertainty. however, these require significant modifications to the training procedure and are computationally expensive compared to standard (non-bayesian) neural networks (nns). above methods improve the estimation of prediction confidence, thus, require model execution, while our method stops model execution if output would be unreliable. most of them do uncertainty calibration with the base model training. however our method can be used for an existing model, without any information about their training data. our gatekeeper is similar to selective prediction approaches. selective prediction is also commonly used to increase the reliability of machine learning models @xcite . in selective prediction, a calibrator is used to preemptively filter out model inputs whose prediction score will not clear the system threshold. @xcite proposed a calibrator model trained using the difficulty level of the instances and confidence scores. this approach needs to be executed along with base model training for the difficulty level calculation. @xcite proposed a calibrator model for selective question answering under domain shift. the calibrator is trained using qa system prediction confidence scores on held-out source data and known ood data. thus it requires access and knowledge of base model training data. @xcite distill the knowledge of qa models into transformer-based question filtering model. to train such a student model knowledge of teacher qa model architecture is required. however, for a gk training the base model's architecture is not required, it can work for black-box base models. our approach uses the base model's test and/or execution data (input and output) to train the gatekeeper model. gatekeeper learns a relationship between inputs, in and/or out of model source domain, and outputs that are below confidence threshold. ## problem formulation the tp model (a llm) returns accurate predictions but is costly. current text prediction workflow calls this model, after the user types a character to get subsequent sequence of characters. the high call rate further increases the cogs (cost of goods sold) of text prediction workflow . the server-side model returns a confidence score (ranging from -1000 to 1), along with text prediction, to indicate the quality of the prediction. the predictions having score less than a rendering threshold (based on model evaluation study) are not shown to user. in a production environment, the suggestion rate was less than 10%. this means that more than 90% of requests to the server model result in a prediction that is not good enough to show to user. in offline analysis, we observed that tp model returned a very low confidence score for ≈ 18% of the tested records. figure 1 shows the percentage split of confidence scores for the key observed text characteristics. these observations made us think: can we identify such input texts which will get no or poor quality response from tp model? can we ensure that every call made to large model returns useful response and is actually worth its cost? ## gatekeeper we propose a gatekeeper (gk) to suppress only the unrewarding calls to large tp model, while maintaining the overall performance. gatekeeper is a light weight classifier on client-side that can predict the probability of getting a null from the tp model. figure 2 shows the text prediction workflow with a gk. here, for each user request, the client application invokes the local gk (passing up to the last 256 character) to find the probability of getting an incorrect prediction. if this probability is low, then client application calls large tp model and shows its prediction (based on rendering threshold); else doesn't show any suggestion (waits for the next character input and so on). this way the unrewarding executions of the server-model and cogs are reduced; user is not bothered with incorrect suggestions and delays. the output of gk is 1 when the probability of poor response from the large tp model for a given figure 3: gatekeeper development approach input is high, else it is 0. gk(x) = 1, if tp(x) < rendering threshold, 0, otherwise. ## results analysis and discussion we evaluated the performance of gk by executing a pipeline of gatekeeper and tp model on test set. the data in test set was not used during gatekeeper model training or validation. we used these results to determine the threshold for transformer-model gk and select the type of gatekeeper. ## conclusion and future work in this paper we presented a gatekeeper to improve the usage efficiency of llm. the model gk is designed to the reduce number of executions of llm without negatively impacting the overall performance of scenario. we developed a gatekeeper for large tp model using its evaluation results. we demonstrated that the model-based gatekeeper improves large tp model's efficiency (i.e. ratio of cogs increase to char saving decrease) by ≈ 10 times at a threshold. in production, we observed that gk (transformer + rule) provided ≈ 55% cogs saving with less than 1% loss in character saving (when 5% is acceptance criteria) for a set of web-client users. we plan to test and tune the transformer-based gatekeeper for a few large tp models, to establish the generality of the gk. we will develop gatekeepers for other text sequence models, such as grammar and sentence correction, for reducing their cogs without impacting the user experience. we need to ensure that model gatekeepers are developed and updated at the same pace as the llm are being released. so, we plan to use continuous learning algorithms to update the gatekeeper models in dynamic environments.