Studying Strategically: Learning to Mask for Closed-book QA
This work offers an incremental improvement for researchers and practitioners working on closed-book QA by optimizing the intermediate pre-training process.
This paper addresses closed-book question answering by learning an optimal masking strategy for an intermediate pre-training stage. The learned policy identifies and masks spans likely to be tested, outperforming strong heuristic-based masking strategies on the TriviaQA dataset when pre-training BART.
Closed-book question-answering (QA) is a challenging task that requires a model to directly answer questions without access to external knowledge. It has been shown that directly fine-tuning pre-trained language models with (question, answer) examples yields surprisingly competitive performance, which is further improved upon through adding an intermediate pre-training stage between general pre-training and fine-tuning. Prior work used a heuristic during this intermediate stage, whereby named entities and dates are masked, and the model is trained to recover these tokens. In this paper, we aim to learn the optimal masking strategy for the intermediate pre-training stage. We first train our masking policy to extract spans that are likely to be tested, using supervision from the downstream task itself, then deploy the learned policy during intermediate pre-training. Thus, our policy packs task-relevant knowledge into the parameters of a language model. Our approach is particularly effective on TriviaQA, outperforming strong heuristics when used to pre-train BART.