PaperlayerView exact revision on arXiv

Source-linked AI summary

Train No Evil: Selective Masking for Task-Guided Pre-Training

Yuxian Gu, Zhengyan Zhang, Xiaozhi Wang, Zhiyuan Liu, Maosong Sun

arXiv:2004.09733v2cs.CL

TL;DR

Task-agnostic pre-training and limited supervised data can leave domain- and task-specific patterns underlearned. The paper adds task-guided pre-training with selective masking between general pre-training and fine-tuning. It reports comparable or better performance than conventional training with less than 50% of overall computation cost.

  • Problem

    Conventional pre-training is task-agnostic, while limited supervised fine-tuning data can hinder models’ ability to capture domain- and task-specific patterns.

    Source: §1 · p. 1

  • Method

    The framework adds masked-language-model pre-training on in-domain unsupervised data and selectively masks important tokens identified through task-specific scoring.

    Source: §1 · p. 2

  • Results

    With less than 50% of overall computation cost, the method achieves comparable or better performance than conventional pre-train-then-fine-tune.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude that task-guided pre-training can help models learn domain- and task-specific patterns efficiently and effectively.

    Source: §4 · p. 5

    Applying task-guided pre-training when in-domain data is limited remains a direction for future work.

    Source: §4 · p. 5

Abstract

from arXiv · show

Recently, pre-trained language models mostly follow the pre-train-then-fine-tuning paradigm and have achieved great performance on various downstream tasks. However, since the pre-training stage is typically task-agnostic and the fine-tuning stage usually suffers from insufficient supervised data, the models cannot always well capture the domain-specific and task-specific patterns. In this paper, we propose a three-stage framework by adding a task-guided pre-training stage with selective masking between general pre-training and fine-tuning. In this stage, the model is trained by masked language modeling on in-domain unsupervised data to learn domain-specific patterns and we propose a novel selective masking strategy to learn task-specific patterns. Specifically, we design a method to measure the importance of each token in sequences and selectively mask the important tokens. Experimental results on two sentiment analysis tasks show that our method can achieve comparable or even better performance with less than 50% of computation cost, which indicates our method is both effective and efficient. The source code of this paper can be obtained from https://github.com/thunlp/SelectiveMasking.

1 Introduction

Conventional pre-training is task-agnostic, while limited supervised data can hinder learning domain- and task-specific patterns. The paper adds task-guided pre-training with selective masking between general pre-training and fine-tuning.

1 Introduction

  • Prior approaches either require large in-domain datasets and substantial computation or use intermediate supervised pre-training that can hurt transferability.The latter concern is attributed to Wang et al. (2019).Source: §1 · pp. 1–2
  • The three-stage framework adds task-guided pre-training on mid-scale in-domain unsupervised data between general pre-training and fine-tuning.It uses masked language modeling and selectively masks tokens considered important for downstream tasks.Source: §1 · p. 2
  • Task-relevant tokens vary by task, and selectively masking important tokens is proposed as a more efficient way to learn task-specific patterns than random masking.The passage gives sentiment words such as “like” and “hate” and relation-extraction predicates and verbs as examples.Source: §1 · p. 2
  • A token-scoring method identifies important tokens, and a neural network trained on labeled downstream examples learns to select them on unlabeled in-domain data.This lets the method use mid-scale in-domain data despite limited supervised downstream data.Source: §1 · p. 2
  • With less than 50% of overall computation cost, the method achieves comparable or better performance than conventional pre-train-then-fine-tune.Experiments cover two sentiment-analysis tasks.Source: §1 · p. 2

2 Methodology

The methodology inserts task-guided pre-training between general pre-training and fine-tuning, using selective masking on in-domain data. It identifies important tokens from labeled task data and learns to select them in unlabeled in-domain text.

2.1 Training Framework

  • The framework has three stages: general pre-training, task-guided pre-training on in-domain data, and downstream fine-tuning.General pre-training randomly masks 15% of general-domain tokens; TaskPT selectively masks important tokens, while fine-tuning follows the conventional approach.Source: §2.1 · p. 2

2.2 Selective Masking

  • TaskPT uses selective masking on mid-scale in-domain data to learn domain-specific and task-specific language patterns.The selection rule for task tokens is first learned from the downstream supervised data.Source: §2.1 · p. 2, §2.2 · p. 2

2.1 Training Framework

  • TaskPT is intended to reduce the need for full general pre-training, lowering the combined pre-training time cost.The data scales are about 1000M words for general unsupervised data, 10M for in-domain unsupervised data, and 10K for downstream supervised data.Source: §2.1 · p. 2, §2 · p. 2

Finding important tokens

  • For labeled task examples, tokens are scored sequentially by their impact on classification confidence, and scores below threshold δ mark important tokens.The method removes previously marked important tokens from the buffer so they do not influence later scores; experiments set δ = 0.05.Source: Finding important tokens · p. 3

Masking on in-domain unsupervised data

  • Because in-domain data lacks classification labels, the method trains a token-importance classifier on labeled task examples and uses it to score in-domain tokens without supervision.The classifier is fine-tuned on examples annotated with important-token labels and used to select tokens for masking.Source: Masking on in-domain unsupervised data · p. 3

3 Experiments

Experiments on sentiment-analysis tasks compare task-guided pre-training and masking strategies across four task–domain combinations. The reported results cover accuracy and pre-training cost, the effect of domain similarity, and a token-selection case study.

3.1 Experimental Settings

  • The evaluation uses BERTBASE on MR and SemEval14 restaurant and laptop tasks, with Yelp and Amazon as in-domain unsupervised data.Fine-tuning runs for 10 epochs, selecting the version with highest dev-set accuracy.Source: §3.1 · p. 3

Efficiency

  • Across all four settings, the method achieves comparable or better performance with less than 50% of pre-training cost.Figure 2 plots test accuracy against overall pre-training steps for the four task–domain combinations.Source: Efficiency · p. 3, Figure 2 · p. 4
  • Sem14-Rest+Yelp exceeds fully pre-trained BERTBASE by 1.4% with half the training steps, while MR+Yelp is 1.94% lower.The authors attribute the difference to greater text-domain similarity between Sem14-Rest and Yelp than between MR and Yelp.Source: Efficiency · p. 4

Effectiveness

  • Selective masking consistently outperforms random masking, including when the in-domain data are less similar to the downstream task.Task-guided pre-training improves results over GenePT in three settings with either masking strategy; random masking harms performance in MR+Yelp.Source: Effectiveness · p. 4

3.3 Case Study

  • The case study reports that selective masking picks sentiment tokens on both supervised MR data and unsupervised Yelp data.Table 2 presents example sequences masked by the selective method and by the PLM scoring function, with selected tokens shown in bold.Source: §3.3 · pp. 4–5, Table 2 · p. 4

4 Conclusion

The conclusion presents task-guided pre-training with selective masking as a three-stage framework for learning domain- and task-specific patterns. Experiments are limited to two sentiment-classification tasks using BERT, and the authors identify scarce in-domain data and masking design as directions for future work.

4 Conclusion

  • The paper presents a three-stage framework that adds task-guided pre-training with selective masking to PLM training.The authors report that it helps models learn domain- and task-specific patterns and benefits downstream tasks.Source: §4 · p. 5
  • Experiments cover two sentiment-classification tasks using BERT, while the authors state the method can generalize to other masked-language-modeling approaches and text-classification tasks..Source: §4 · p. 5
  • Future work includes applying task-guided pre-training when in-domain data are limited and designing more effective selective-masking strategies.Source: §4 · p. 5

A Cost of Selective Masking

Selective masking adds a small share of pre-training time while the framework saves about half of conventional pre-training cost. Its token-selection overhead is described as insignificant relative to the time saved.

A Cost of Selective Masking

  • The pipeline fine-tunes BERT, labels important downstream tokens, trains a token-level classifier, then uses it to mask important tokens in in-domain data.The classifier transfers token-importance selection from supervised downstream data to in-domain unsupervised datasets.Source: A Cost of Selective Masking · p. 6
  • The four-step selective-masking overhead was considered insignificant compared with pre-training time saved and was excluded from pre-training cost comparisons.Table 3 compares the four token-selection steps and their summed cost with saved pre-training time.Source: A Cost of Selective Masking · p. 6, Table 3 · p. 6
  • About 36 hours—roughly half the conventional pre-training cost—was saved, while selective masking accounted for only a small part of total pre-training time.Figure 3 depicts the conventional cost as a whole pie, with colored sectors for GenePT, selective masking, and TaskPT and a white sector for time saved.Source: A Cost of Selective Masking · p. 6, Figure 3 · p. 7

B Detailed Experimental Setup

GenePT generally follows BERT's pre-training procedure and hyperparameters, with a 256-token sequence limit and FP16 precision. It ran for 1M steps on four NVIDIA V100 GPUs and saved checkpoints during training.

B.1 GenePT

  • GenePT generally followed BERTBASE's pre-training procedure and hyperparameters, with the maximum sequence length set to 256 tokens and FP16 precision used for efficiency.The passage cites Devlin et al. (2019) for BERTBASE.Source: §B.1 · p. 7
  • GenePT ran on four NVIDIA V100 GPUs, and the full 1M-step training took about three days.Source: §B.1 · p. 7
  • Checkpoints were saved at 100k, 200k, and 300k steps during training.Source: §B.1 · p. 7

B.2 Selective Masking.

Selective masking derives token-importance labels from downstream supervised data, trains a token-level binary classifier, and applies it to in-domain unsupervised data for masking.

B.2 Selective Masking

  • The authors fine-tuned GenePT checkpoints at 100k, 200k, 300k, and 1M steps on MR and SemEval14, selecting the highest-valid-accuracy model after at most 10 epochs.Fine-tuning used batch size 64, maximum sequence length 256, learning rate 2e-5, and random seed 42.Source: §B.2 · p. 7
  • Fine-tuned MR and SemEval14 models selected important tokens on downstream data, using δ = 0.05 and labeling selected tokens 1 and others 0.The sentences were tokenized with BERT’s sub-word tokenizer.Source: §B.2 · pp. 7–8
  • A token-level binary classification head was trained on the GenePT checkpoints using the annotated data, with label 1 weighted 1.5 to balance the labels.Training used maximum sequence length 128, batch size 64, and learning rate 1e-5.Source: §B.2 · p. 8
  • The trained classifier labeled tokens in mid-scale in-domain data, and tokens predicted as important were masked for subsequent pre-training.Source: §B.2 · p. 8

B.3 TaskPT

TaskPT continued pre-training GenePT checkpoints on selectively masked in-domain data for at most 200k steps. The resulting model was fine-tuned on MR and SemEval14, with performance averaged over 10 random seeds.

B.3 TaskPT

  • TaskPT continued pre-training GenePT checkpoints on selectively masked in-domain data for at most 200k steps.Its hyperparameters were almost the same as GenePT's.Source: §B.3 · p. 8

B.4 Fine-tuning

  • The post-TaskPT model was fine-tuned on MR and SemEval14, with performance averaged over 10 random seeds.The fine-tuning hyperparameters generally matched the earlier Fine-tune BERT stage.Source: §B.4 · p. 8

C Detailed Datasets Description

The section describes four sentiment-classification datasets and their preparation: MR receives a random split, SemEval14 is converted to conventional classification, and Yelp and Amazon review text supplies in-domain unsupervised data.

C Detailed Datasets Description

D Results on SemEval14-Laptop

On SemEval14-Laptop, task-guided pre-training is reported to improve performance at less than half the training cost; experiments compare early-stopped efficiency runs with full-pretraining effectiveness runs.

D Results on SemEval14-Laptop

  • Sem14-Lap is evaluated with Yelp and Amazon as in-domain datasets, using 100k, 200k, and 300k early stops for efficiency and full BERTBASE pre-training for effectiveness.Source: D Results on SemEval14-Laptop · p. 9
  • Table 5 compares test accuracies after 1M general-pretraining steps for models without task-guided pre-training and models using different masking strategies.Source: Table 5 · p. 9
  • Better performance is reached with less than 50% training cost on Sem14-Lap.Source: D Results on SemEval14-Laptop · p. 9
Loading 2004.09733v2…