QUILL : Query Intent with Large Language Models using Retrieval Augmentation and Multi-stage Distillation (2022)

Abstract: Large Language Models (LLMs) have shown impressive results on a variety of text understanding tasks. Search queries though pose a unique challenge, given their short-length and lack of nuance or context. Complicated feature engineering efforts do not always lead to downstream improvements as their performance benefits may be offset by increased complexity of knowledge distillation. Thus, in this paper we make the following contributions: (1) We demonstrate that Retrieval Augmentation of queries provides LLMs with valuable additional context enabling improved understanding. While Retrieval Augmentation typically increases latency of LMs (thus hurting distillation efficacy), (2) we provide a practical and effective way of distilling Retrieval Augmentation LLMs. Specifically, we use a novel two-stage distillation approach that allows us to carry over the gains of retrieval augmentation, without suffering the increased compute typically associated with it. (3) We demonstrate the benefits of the proposed approach (QUILL) on a billion-scale, real-world query understanding system resulting in huge gains. Via extensive experiments, including on public benchmarks, we believe this work offers a recipe for practical use of retrieval-augmented query understanding.

## introduction the recent advent of billion+ parameter large language models (llms) -such as t5 @xcite , mt5 @xcite , gpt-3 @xcite and most recently palm @xcite -has disrupted many language understanding tasks -with new benchmarks set or eclipsed routinely by these transformer models and their variants. queries -especially keyword search onespresent a unique challenge though. their short length, inherent ambiguity and lack of grammar mean query understanding tasks typically require more memorization and world knowledge than other nlp tasks @xcite . consequently, despite llms leading performance on language and query understanding tasks -like intent classification, query parsing and relevance prediction -there is significant room for further improvement. in this paper we leverage retrieval-augmentation to provide llms more context and grounding for search queries. we show that the titles and urls of documents retrieved for the query, greatly help improve llms query understanding capabilities. while different retrieval augmentation models exist, we show that even simple concatenation of these titles / urls with the query can help improve llm performance considerably. however, the use of retrieval augmentation leads to a new challenge: increased complexity of llm inference. more specifically, the quadratic complexity of self-attention in transformer models means that the latency of llms blows up given these (often 10x+) longer input sequences. this presents a significant problem as llms are impractical for online use and thus need to be distilled into smaller, more efficient models to be served online. however knowledge distillation @xcite into these student models requires a lot of distillation data annotated by these llms -which may not be feasible for these retrieval augmented models. thus as a remedy we introduce a new two-stage distillation approach. in the first stage of this approach we distill the retrieval-augmented (long input) llm (the professor) into a non-retrieval augmented (short input) llm (the teacher) using a small distillation set. this second llm teacher is in turn distilled into the final student using a large set. via extensive experiments on a large-scale, realworld problem and data we demonstrate that the resulting quill system provides for an efficient and effective way of retaining the performance gains of retrieval augmented llms on query understanding tasks. ## related work large language models (llms) such as mt5 @xcite demonstrated significant performance improvements on a variety of natural language understanding (nlu) tasks. specifically in the context of query understanding, researchers found that (a) model size significantly effects the quality of the resulting models @xcite , and (b) using additional context in the form of query-associated documents is crucial to the model performance due to the paucity of context available in the query itself @xcite . retrieval augmentation of the query with the search results retrieved by it is a proven way to incorporate such context in llm training for nlu tasks, as has been shown recently by models such as rag @xcite , realm @xcite , and retro @xcite . in this paper, we leverage this insight to improve performance of query intent prediction @xcite -a crucial query understanding task that is at the heart of modern search engines -using llms. prior work by @xcite found importance of retrieval augmentation using statistical methods for this task. statistical retrieval augmentation has also been found critical for other query understanding tasks including query expansion @xcite and query tagging @xcite . we demonstrate similar benefits when using retrieval-augmented llms as well. we also leverage knowledge distillation @xcite @xcite techniques to create a student model that retains the llms gains. ## query intent understanding while the techniques described in this paper could be applied to any query understanding task, for the sake of brevity we focus on the task of query intent classification. query intent (qi) classification is a classical ir task studied for over two decades @xcite 2006; @xcite @xcite @xcite . this task is particularly important in practice, as it is at the top of the search funnel, and the entire search engine behavior may vary based on the predicted query intent. given the centrality of this task on overall retrieval, models for this task need to be both fast (i.e., low latency) and high efficacy. thus even a single percentage point quality gain on the qi task can be considered a major accomplishment. in this paper we tackle the qi task using llms. in particular we use two datasets in our study whose details are provided in ## quill methodology the keyword nature of queries and lack of context make the qi task (like other query understanding tasks) challenging for llms. thus we propose quill as a solution. as seen in figure 1 , quill consists of two stages: (a) retrieval augmented llm training, (b) multi-stage distillation into efficient student. ## experiments and results experimental setup: our experiments were all conducted using the mt5 @xcite ) checkpoints. we validate performance across three learning rates (1e-3, 1e-4, 5e-5) -selecting the best checkpoint using the validation set loss. for mod-els trained from the provided training sets, we used a batch size of 64 in our experiments and trained for 4k steps (ecomm) / 20k steps (orcas-i). for distilled models, we used a batch size of 128 for teacher models and 1024 for student models. we use different batch sizes because of the model architectures, mt5 for the teacher vs a bertbased model for the student. in both cases, we trained for 1 epoch, unless mentioned otherwise. we only use the encoder of the mt5 model with an additional layer added on top to predict the classification scores. the professor, the teacher and the student fine-tuning experiments are all set up as a query intent classification task. given that the teacher and student models are trained on millions of examples and this itself is a time and resource intensive step, we restrict our experiments to only one epoch. we demonstrate performance gains even with one epoch via the techniques elaborated in this paper. we studied the effect of distillation data size, for both stages of distillation. for the ecomm dataset, we used an in-house retriever to find related documents. for orcas-i, we use the provided docids (aggregated at per-query level) for retrieval augmentation. unless specified, we use (upto) the top-10 results for retrieval augmentation 1 . sequence length for models are based on the training set and features (set to 99%-percentile of sequence lengths). ## future work while we studied the problem of query intent classification in this paper, the approach proposed in our paper is general and could be applied to any query understanding task. following our approach, could enable myriad query understanding tasks use retrieval augmentation in a practically realistic and efficient manner. we leave this to future work though. we should also note that our experiments reveal non-trivial distillation gaps in both stages of distillation, which we believe is another open opportunity for future improvements. ## conclusion this paper provides a practical recipe for combining retrieval augmentation and large language models. in particular, we proposed quill as an approach to tackle the problem of query intent classification. our empirical study demonstrates conclusively that retrieval augmentation can provide significant value over existing approaches. furthermore we show that via our two-stage distillation approach, that quill not only learns better performing, more robust teachers, but also leads to even bigger gains when distilled into fast, realworld capable production student models.