Optimizing Retrieval-augmented Reader Models via Token Elimination (2023)

Abstract: Fusion-in-Decoder (FiD) is an effective retrieval-augmented language model applied across a variety of open-domain tasks, such as question answering, fact checking, etc. In FiD, supporting passages are first retrieved and then processed using a generative model (Reader), which can cause a significant bottleneck in decoding time, particularly with long outputs. In this work, we analyze the contribution and necessity of all the retrieved passages to the performance of reader models, and propose eliminating some of the retrieved information, at the token level, that might not contribute essential information to the answer generation process. We demonstrate that our method can reduce run-time by up to 62.2%, with only a 2% reduction in performance, and in some cases, even improve the performance results.

## introduction the task of open-domain question answering (odqa) @xcite consists of answering questions using external knowledge, which is used as a source of relevant information that might be helpful for a model to extract or generate the right answer for a question. the expected answer can be short and concise @xcite , or long and detailed @xcite , in which it is called long-form question answering (lfqa). the retriever-reader architecture has been widely-used and adopted for odqa tasks @xcite . the retriever fetches the most relevant passages using the question as a query. then, the reader extracts or generates an answer, using the question and the relevant passages. the explicit structure of the system, consisting of these two sub-modules, allows for a decoupled optimization of either the retrieving or the reading process. in this work, we exclusively focus on the optimization 1 we provide the source code for our work at https:// github.com/mosheber/token_elimination . here, two passages are considered as the input to the decoder module. the token filtering operation is performed when generating token t, between decoder layer l and l + 1. using the representation of token t at layer l, the cross-attention is computed for all the input tokens. then, we filter out the lowest ranked tokens from the input (marked in yellow), with only the highest ranking input tokens being used from the next generated token onward. there has been rapid and remarkable progress in retriever-reader systems for solving odqa tasks using a generative approach @xcite . one such prominent approach is fusion-in-decoder (fid) @xcite ) that utilizes a generative text-to-text model to generate an answer. despite the significant performance improvements, there are several computational bottlenecks associated with fid that have a negative impact on efficiency. the most prominent ones are: (a) the need to retrieve a relatively large amount of documents to reach peak performance, and (b) an extensive cross-attention operation, caused by processing multiple concatenated retrieved passages, applied repeatedly for every generated token. these bottlenecks are negatively amplified in the case of long-form question answering. previous works have attempted to mitigate these bottlenecks, either by limiting the input to the reader or by directly optimizing it in a variety of methods. @xcite included a passage re-ranker inside the reader which aimed to filter out the least relevant passages during encoding. @xcite optimized the decoder module by pretraining a modified and optimized architecture, and @xcite modified the attention operations performed to be less computationally intensive. in this work, we tackle the heavy cross-attention computation in the decoder by introducing token filtering, a method that removes redundant tokens from input passages during the decoding stage, by dynamically computing their salience during generation. using token filtering eliminates uninformative tokens from the cross-attention matrix, and prevents them from being utilized during answer generation, directly contributing to the reduction of the overall generation time. to further boost efficiency and reduce latency, we combine our token filtering approach with dynamic decoder layer skipping introduced by @xcite , referred to as calm. by combining both approaches and by conducting experiments on three lfqa datasets, we find that this approach presents a better performance vs. efficiency trade-off than by using the methods separately, in most cases. overall, our contributions are as follows: ## preliminaries in a retriever-reader system, the reader, which is typically a language model, receives a query along with a collection of passages, where each passage often consists of a title and a context. additionally, we are provided with the ground truth, which can be an expected answer or a gold passage that is most relevant to the query. since our main focus is on generative models, we employ the widelyused fusion-in-decoder (fid) model @xcite , which is a cutting-edge encoderdecoder model based on the t5 model @xcite . the encoder module of the fid model processes the input passages in parallel, with each layer taking the output of the previous layer, and the final output of the encoder is the output of its last layer. similarly, each layer of the decoder processes its input by receiving the output of the preceding layer. the decoder module then cross-attends to the large number of concatenated input representations and assimilates the information from the different passages to generate an answer. at each decoding step, the decoder computes the attention scores based on the precomputed input tokens' representations which serve as the query for the multi-headed attention operation, concurrently taking into account the current decoded sequence. ## method following our analysis in section 3.2, we turn to implementing a method for filtering out the redundant information during the decoding stage. we aim to find a subset of the input tokens that is the most relevant for generating the correct answer. as pointed out previously, we can utilize the crossattention scores computed between the generated tokens and the passages as basic signal for filtering out irrelevant tokens, similarly to @xcite . thus, we suggest a token filtering approach, using the cross-attention scores computed at a predetermined layer and generated token index during inference. at that point, for each input token, we compute the average attention scores over all attention heads, similarly to @xcite ; @xcite . once these scores are computed, we keep the top k%-scored input tokens, which will be the only tokens to be considered towards the next tokens' predictions. formally, the cross-attention scores per input token are defined as follows: we perform a descending argsort operation on the scores above, and take the top p% from the sorted input token indices. hence, we denote t as the total number of input tokens from all the passages, and t ′ as the amount of tokens we keep after filtering, which is p% from t : since the filtering percentage, token index and layer can effect the quality of this mechanism, as inspected in section 3.2, we obtain the optimal values for them by performing a hyperparameterlike search over their possible values, which is described in section 5.4. we produce new past key and value representations for the input sequence (across all the decoder layers), containing only the selected tokens, resulting in a more compact tensor to attend to. we name the filtering mechanism token filtering, with an overview of the approach presented in figure 1 . note that we remove the irrelevant tokens from the keys and values of the encoder output sequence during inference time only once, hence reducing their dynamic dimension during computation for all the subsequent tokens. for additional details about the cross-attention scoring computation we refer the reader to appendix b. ## related work open-domain question answering many previous works utilized setups which are based on the retriever and the language model reader components @xcite @xcite . the goal of the retriever is to fetch the most relevant passages to a given question @xcite . the reader processes the question and the relevant passages to extract or generate the answer, with generative approaches 7 since the results in encoder-decoder efficiency due to computation bottlenecks in generative models, particularly in encoder-decoder models @xcite , previous works attempt to mitigate them. the encoder model can be used to filter out irrelevant passages during generation @xcite . nevertheless, the encoder's impact on the model's latency in longform scenarios is negligible, as depicted in figure 2 . consequently, our research centers on analyzing the computational aspects of the decoder instead. in particular, the decoder model utilizes many redundant and heavy cross-attention operations, which can be removed or replaced with simpler alternatives @xcite . since encoder-decoder models perform compute heavy operations in multiple layers, previous works have proposed stopping the layer propagation dynamically by assessing the model's confidence for prediction at a certain layer @xcite . other works have adapted this mechanism to decoder models as well @xcite . however, these studies fail to tackle the issue of input size during generation, thereby resulting in computations being performed on irrelevant input to some degree, which we address through complementary token filtering. ## conclusions we analyze the precise performance vs. efficiency trade-off of the fid's encoder and decoder in long-form settings, with an analysis of the crossattention operation in the decoder model. we show that the decoder has more impact on the latency, particularly for long outputs, and that the decoder attends to more salient information early on during generation. hence, our proposed approach for efficiency reduction, namely a combined approach of token filtering and calm, removes irrelevant layers and tokens during the generation process, for every token produced. our approach achieves a significant reduction in resources (up to 62.2%), while not sacrificing more than 2% of the performance, and is the current state-of-the-art on the eli5 kilt leaderboard. future work can further develop a more dynamic method for choosing the most relevant tokens from the input, instead of using predetermined hyperparameters, and train the cross-attention patterns to better attend to the salient information during generation.