Passage-Mask: A Learnable Regularization Strategy for Retriever-Reader Models (2022)

Abstract: Retriever-reader models achieve competitive performance across many different NLP tasks such as open question answering and dialogue conversations. In this work, we notice these models easily overfit the top-rank retrieval passages and standard training fails to reason over the entire retrieval passages. We introduce a learnable passage mask mechanism which desensitizes the impact from the top-rank retrieval passages and prevents the model from overfitting. Controlling the gradient variance with fewer mask candidates and selecting the mask candidates with one-shot bi-level optimization, our learnable regularization strategy enforces the answer generation to focus on the entire retrieval passages. Experiments on different tasks across open question answering, dialogue conversation, and fact verification show that our method consistently outperforms its baselines. Extensive experiments and ablation studies demonstrate that our method can be general, effective, and beneficial for many NLP tasks.

## introduction retriever-reader based approaches are popularly considered in the knowledge-intensive tasks (e.g., open question answering (qa), fact verification). it is designed to retrieve a set of support documents and extract the answer from these documents. mostly adopted retrieve and read models (e.g., izacard and grave, 2020) are trained to generate the annotated gold answers using the reader model, based on passages obtained by the retrievers (e.g., @xcite . this training process of reader disregards the evidentiality of all retrieval passages and can easily overfit the top ranked passages @xcite . even if the top-rank passages in the test setting do not have the correct answers, these models still tend to find the answer in the top-rank passages and yield worse performance @xcite . it happens to the reader model due to the overfitting and the memorization of outdated information @xcite . to what extent does the reader model quality depend on the retrieval passages? we analyze the ranking impact of the retrieval passages from masking (e.g. mask out the top three passages), permuting, and removing. the overfitting, as well as the performance degradation, is observed. to desensitize the impact from the top-rank passages, we consider masking passages during training which serves as a desensitizer and can improve the reader model ability to reason over all retrieval passages. however, the standard masking and dropout strategies are not designed for our focused tasks and also bring an increased gradient variance during training due to their randomness. in the meantime, each neuron plays the same role and has the same mask. however, in the reader model, intuitively, top-rank passages often have a higher chance to overfit during the training. to this end, we introduce our passage mask (pm), which encourages to mask top-rank passages. reducing the gradient variance with fewer mask candidates and optimizing the mask candidates with bi-level optimization, the mask magnitude for each candidate can be learned. overall, the proposed mask parameters are jointly optimized with the entire network. we run extensive experiments across representative knowledge-intensive tasks: open-domain qa (natural questions open @xcite ; triviaqa unfiltered @xcite ), fact verification (faviq @xcite ), and knowledge grounded dialogue (wizard pf wikipedia @xcite ). our method shows large performance improvements across different tasks and datasets. furthermore, we provide extensive ablation studies on different design choices for the proposed method, including the designs of masking candidate space and efficiency. our analysis shows the passage mask contributes the performance improvement, helping the reader learn to focus on the retrieval passages without being distracted by high-ranked passages with more lexical overlaps. with little modification, our regularization can be easily applied to other nlp tasks for a better answer generation strategy. to the best of our knowledge, we present the first mask regularization in the open retriever-reader setting by preventing the rank-related overfitting in open qa, dialogue conversation, and fact verification. our contributions are summarized as follows: although fid does not use the unnormalized passage score as dpr, we still find out that fid has a preference over passages with higher retrieval passage scores. our analysis in ## experimental settings ## experiments we evaluate the performance of our mask and learning framework in this section. we bold the best result within each column block. the results of our method are obtained with three independent runs to determine the variance. see appendix a for full results with error bars. model nq triviaqa dev test dev test dpr (karpukhin et al., 2020) -41.5 -57.9 rag (lewis et al., 2020) -44.5 -56.1 colbert-qa (khattab et al., 2021) -48.2 -63.2 realm (guu et al., 2020) -40.4 --fid base (izacard and grave, 2021) 49.2 50.1 68.7 69.3 ours base 49.9 51.3 69.3 69.9 fid large (izacard and grave, 2021) 52.7 54.4 72.5 72.5 ours large 53.3 55.3 73.1 72.9 ## analysis what is the influence of the vanilla mask and dropout? here we verify whether pm is better model f1 fid base @xcite 17.1 eqa base @xcite 18.0 ours base 18.7 efficiency and running time. we provide the parameter sizes, gpu peak memory, and per step time position 1 st 2 nd 3 rd 4 th 5 th fid base permute top 3 ✓ ✓ ✓ 50.0 permute top 5 ✓ ✓ ✓ ✓ ✓ 50.0 remove 1 st ✓ 44.9 remove 2 nd ✓ 48.7 remove 3 rd ✓ 49.3 ablation studies on the components in pm. we conduct the ablation study to exam the role of bilevel optimization and reduced mask candidate set. for ablation, instead of searching the mask probability for different mask candidates, we randomly sample a candidate in the search space. through isolating performance of each components, our focus here is to identify the impact of the introduced mask parameter w and the reduced mask set. ❶ would we see improvements if finetuning the given checkpoint with baselines? as discussed in section 3.2, due to computation cost limitation, we use the provided checkpoint for the reader model and continue the finetuning with our method. however, if we continue finetuning the baseline checkpoint, would we still see the improvements? we conduct the experiments on open qa, dialogue and fact verification tasks. we adopt the best baseline models for each task such as fid base for nq and triviaqa, and eqa base for dialogue conversations and fact verification. in ## related work retrieval read architecture recent retriever models (e.g., @xcite @xcite to generate answers condition on the question and the retrieved passages @xcite @xcite . our work relies on this architecture and further fine-grain the reader model to introduce the passage-aware masking and promote the reasoning over the entire passage set. rank-related studies passage ranking has shown promising performance improvements. the most popular approach is combining the passage score and answer score together @xcite @xcite . other works (e.g., @xcite @xcite propose additional modules or operations to re-identify the passage rank. @xcite uses seq2seq model to identify the document's relevance to the query, @xcite introduces a passage re-ranking module, and @xcite proposes to use the calibrator as an answer reranker. there are some works that focus on the ranking efficiency. @xcite creates a simple neural model that combines the efficiency of dual encoders. similarly, we also find out that directly taking the rank makes the model overfitting. different from existing works, pm rethinks the impact of retrieval passage ranking from the regularization and generalization perspective. we focus on preventing the overfitting and improving the reasoning generalization during training. in the meantime, pm is also compatible with other previous ranking works with the potential to jointly improve the performance. ## conclusion our work demonstrates the benefits of introducing a passage mask mechanism. the proposed mask can desensitize the impact from the top-rank retrieval passages and prevent the model from overfitting. the proposed strategy shows noticeable gains in performance across open question answering, dialogue conversation, and fact verification. we further conduct the detailed study with the proposed masking strategy in different settings, e.g., comparing with vanilla masking, providing more evidence for rank-related overfitting, and verifying the impact of different components. to summarize, the proposed pm is effective and general, with the potential to be incorporated into existing models for various nlp tasks.