PaperlayerView exact revision on arXiv

Source-linked AI summary

Studying Strategically: Learning to Mask for Closed-book QA

Qinyuan Ye, Belinda Z. Li, Sinong Wang, Benjamin Bolte, Hao Ma, Wen-tau Yih, Xiang Ren, Madian Khabsa

arXiv:2012.15856v2cs.CLcs.AI

TL;DR

Prior intermediate pre-training used a heuristic that masks named entities and dates, while this work asks how to learn a policy that targets knowledge likely to be queried. The paper trains a policy from context-answer examples to extract likely answer spans, then uses it to mask a generic corpus during intermediate pre-training before downstream fine-tuning. With BART-base, the TriviaQA-learned masking policy achieved 24.71% EM, outperforming the original BART-large model (24.28% EM) and saving 260M parameters.

  • Problem

    Prior intermediate pre-training used a heuristic that masks named entities and dates, while this work asks how to learn a policy that targets knowledge likely to be queried.

    Source: §1 · pp. 1–2

  • Method

    The paper trains a policy from context-answer examples to extract likely answer spans, then uses it to mask a generic corpus during intermediate pre-training before downstream fine-tuning.

    Source: §1 · p. 2, §2 · p. 2, §3 · p. 2

  • Results

    With BART-base, the TriviaQA-learned masking policy achieved 24.71% EM, outperforming the original BART-large model (24.28% EM) and saving 260M parameters.

    Source: §4.2 · p. 4

  • Takeaways & Limitations

    Learned masking policies transfer to some, but not all, cases across QA datasets.

    Source: §4.2 · p. 4

    The policy can merge separate answer spans because it predicts start and end positions independently, as in the example of two entities in one candidate span.

    Source: §4.3 · p. 5

Abstract

from arXiv · show

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.

1 Introduction

Closed-book QA tests whether language models can answer questions from knowledge stored in their parameters, and prior work improved it with an intermediate stage that masks entities and dates. This paper learns a task-supervised masking policy to target likely queried spans; on TriviaQA, it outperforms SSM and random masking with BART-base.

1 Introduction

  • Closed-book QA evaluates whether a language model can answer questions without an external knowledge base, drawing on knowledge implicitly stored in its parameters.Prior work found directly fine-tuned closed-book methods competitive with retrieval-based methods on open-domain QA benchmarks.Source: §1 · p. 1
  • Prior work used salient span masking to recover named entities or dates, with a reported 9%+ exact-match improvement on TriviaQA using T5-11B.The passage contrasts that gain with a 7% improvement from scaling T5 from 3B to 11B parameters.Source: §1 · p. 1
  • The paper learns a policy to extract answer spans likely to be queried, using context-answer supervision from QA tasks, then applies that policy during intermediate pre-training.The policy is intended to pack task-relevant knowledge into the language model's parameters.Source: §1 · pp. 1–2
  • 24.71%(±.21) exact match on TriviaQA with BART-base and a TriviaQA-learned policy surpasses 23.62%(±.29) with SSM and 22.93%(±.14) with random masking.The same trend holds with BART-large, and learned policies transfer to some—but not all—other QA datasets.Source: §1 · p. 2

2 Preliminaries

The paper frames its work as learning a masking policy for intermediate pre-training, between general pre-training and downstream fine-tuning. In masked language modeling, a policy selects tokens to mask and the language model learns to recover them.

2 Preliminaries

  • In masked language modeling, a policy g(x; φ) selects tokens for masking and a language model f(x; θ) learns to recover the masked tokens from corrupted input.Following SSM, the policy masks and reconstructs exactly one span, which can be represented by predicting its start and end positions.Source: §2 · p. 2
  • The proposed pipeline learns a masking policy from one or more anchor QA tasks, applies it to a generic corpus during intermediate pre-training, then fine-tunes on downstream closed-book QA.The stated evaluation uses a preset corpus, namely Wikipedia.Source: Figure 2 · p. 2, §2 · p. 2
  • The policy-learning setup assumes access to (context, question, answer) examples from at least one QA dataset, called an anchor task.The policy is evaluated by intermediate pre-training on the preset corpus followed by fine-tuning on downstream closed-book QA tasks.Source: §2 · p. 2

3 Approach

The approach trains a lightweight span-extraction policy from context-answer pairs, without using the question at policy-learning time. At intermediate pre-training, the policy selects likely answer spans for masking so the language model learns to unmask answer-like content.

3 Approach

  • The policy learns from (context, answer) pairs to extract answers from evidence documents, selecting likely answer spans for masking during intermediate pre-training.The intended effect is to focus the language model on learning to unmask spans resembling answers.Source: §3 · p. 2
  • A two-layer bidirectional LSTM encodes each context, and a linear layer predicts the start and end positions of a potential answer span.The policy uses lightweight architectures because it will be deployed on a large-scale corpus.Source: §3 · pp. 2–3
  • At inference, candidate spans are ranked by the sum of their start and end logits; the policy either masks the top span or samples one from the top five.Source: §3 · p. 3
  • The policy is trained using supervision from Natural Questions, TriviaQA, or a mixture of Natural Questions, TriviaQA, SQuAD, and ReCoRD.For each supervision source, it is trained for 30 epochs and the best checkpoint is selected by validation loss.Source: §3 · p. 3

4 Experiments

The experiments compare learned masking policies with random masking, salient span masking, and BART checkpoints across three closed-book QA datasets. On TriviaQA, the TriviaQA-trained policy outperforms SSM for both BART sizes, while transfer to other datasets is selective.

4.1 Experiment Setup

  • The setup evaluates Natural Questions, WebQuestions, and TriviaQA with BART and intermediate-pretraining checkpoints, reporting exact match averaged across three random seeds.Intermediate pre-training uses 100,000 updates on the December 20, 2018 Wikipedia snapshot.Source: §4.1 · p. 3

4.2 Results and Discussion

  • Random-mask intermediate pre-training also improves BART-base; the authors suggest continued pre-training or Wikipedia’s relation to the QA tasks may explain the gain.SSM is also a strong heuristic; its reported gains are 0.8% for BART-base and 2.0% for BART-large, which the authors hypothesize relates to model size.Source: §4.2 · pp. 3–4
  • EM=24.71% vs EM=24.28%: TriviaQA-trained masking lets BART-base surpass original BART-large, and the learned policy beats SSM with both BART sizes.Table 1 reports the TriviaQA comparison.Source: §4.2 · p. 4, Table 1 · p. 4
  • Learned masking transfers in some, but not all, settings; Learned(NQ, Top5) reaches 23.73%(±.21) EM on TriviaQA, on par with SSM.The authors suggest selective transfers may reflect meta-knowledge about information important to closed-book QA.Source: §4.2 · p. 4, Table 3 · p. 5
  • NQ is a difficult testbed for ranking these methods: BART-base results are close, intermediate-pretraining effects may fall within seed variance, and exact match may underestimate capability.The authors also note that random-initialized BART-base reaches 20% dev EM on NQ, while general pre-training adds only 4–5% EM.Source: §4.2 · p. 4

4.3 Case Study of the Learned Policy

  • The TriviaQA-trained policy predicts salient spans in development contexts, but independently selected start and end positions can combine into an overly broad candidate span.The example joins “Chadwick School” with “University of California” into one candidate.Source: §4.3 · pp. 4–5

5 Related Work

Related work connects pre-trained language models’ implicit knowledge to closed-book QA and situates this paper among heuristic and task-specific pre-training approaches. Unlike objectives developed without a specific task in mind, this work targets knowledge encoding for given tasks.

Implicit Knowledge in Pre-trained Language Models

  • Petroni et al. (2019) found that pre-trained language models store relational knowledge accessible through cloze queries, while Roberts et al. (2020) introduced closed-book QA.Closed-book QA asks models to generate answers using implicit knowledge rather than a retriever-reader strategy.Source: §5 · p. 5

Self-supervised Pre-training

  • Existing self-supervised objectives are mostly heuristic, including masked language modeling, span boundary representation learning, sentence reconstruction, and gap sentence generation.Source: §5 · p. 5
  • Unlike objectives designed for efficiency or pre-training from scratch, this work focuses on encoding implicit knowledge during intermediate pre-training for specified tasks.The passage contrasts this focus with ELECTRA, gradient-variance reduction, and PMI-based n-gram masking.Source: §5 · p. 5

Domain/Task-specific Pre-training

  • Related task-specific approaches include in-domain pre-training, reinforcement-learned mask generation, and task-guided masking based on predicted token importance.Source: §5 · p. 5

6 Conclusion

The paper proposes learning which spans to mask during intermediate pre-training to encode task-relevant knowledge. Its learned policy is particularly effective on TriviaQA, while results are less clear on other datasets and transfer is observed in some cases.

6 Conclusion

  • The method aims to encode task-relevant knowledge during intermediate pre-training, likened to identifying likely test content while studying.Source: §6 · pp. 5–6
  • The learned policy is particularly effective on TriviaQA, while the trend is less clear on other datasets and transfer to other QA datasets occurs in some cases.The authors discuss reasons for the less clear results on other datasets.Source: §6 · p. 6
  • Future work may explore meta-learning for the masking policy and unsupervised-QA-based self-training in a closed loop.The proposed loop would learn to memorize knowledge, ask questions, and provide answers.Source: §6 · p. 6

A Hyperparameters

Fine-tuning hyperparameters are selected in sequence: first learning rate, then batch size.

A Hyperparameters

Loading 2012.15856v2…