[Paper Review] Position Masking for Language Models

Andy Wagner, Tiyasa Mitra+3 more|arXiv (Cornell University)|Jun 2, 2020

Topic Modeling5 references4 citations

TL;DR

This paper proposes position masking—a novel extension of masked language modeling (MLM) in BERT-like models—by additionally masking and predicting the positions of masked tokens using a dedicated classifier head. The method improves training convergence and performance, achieving a 0.3% F1 gain on SQuAD and reducing convergence time to 50% of the original BERT training tokens on IPU hardware, while also enabling faster convergence via enhanced gradient flow through dropout layers.

ABSTRACT

Masked language modeling (MLM) pre-training models such as BERT corrupt the input by replacing some tokens with [MASK] and then train a model to reconstruct the original tokens. This is an effective technique which has led to good results on all NLP benchmarks. We propose to expand upon this idea by masking the positions of some tokens along with the masked input token ids. We follow the same standard approach as BERT masking a percentage of the tokens positions and then predicting their original values using an additional fully connected classifier stage. This approach has shown good performance gains (.3\% improvement) for the SQUAD additional improvement in convergence times. For the Graphcore IPU the convergence of BERT Base with position masking requires only 50\% of the tokens from the original BERT paper.

Motivation & Objective

• To improve training efficiency and performance of BERT-style models by extending masked language modeling to include position prediction.
• To investigate whether adding position masking leads to faster convergence and better generalization in transformer-based language models.
• To evaluate the impact of position masking on both MLM pre-training and downstream SQuAD fine-tuning tasks.
• To explore the effect of modified gradient flow through dropout layers on model performance.
• To compare training efficiency and accuracy across GPU and IPU hardware platforms under the same training regime.

Proposed method

• A new fully connected classifier head is added to predict the original positions of masked tokens, alongside the standard token reconstruction head in MLM.
• The model uses the same masking strategy as BERT: 90% of tokens are masked, with 5% using correct positions and 5% using random positions.
• Position masking is applied in addition to token masking, providing extra supervision for position encoding learning.
• The architecture is trained using standard BERT pre-training objectives, with sequence lengths of 128 and 384, and a pipelined IPU implementation.
• An enhanced dropout gradient technique is introduced, where dropout masks are not applied during backpropagation, improving fine-tuning performance.
• Experiments are conducted on both GPU and IPU platforms, using BERT Base as the baseline and SQuAD as the primary evaluation metric.

Experimental results

Research questions

• RQ1 Does masking token positions in addition to token values improve convergence speed and final model performance in BERT-like models?
• RQ2 To what extent does position masking enhance the learning of positional embeddings in transformer models?
• RQ3 How does the enhanced dropout gradient technique affect fine-tuning performance compared to standard dropout?
• RQ4 Can position masking reduce the number of training tokens required to reach convergence, especially on specialized hardware like IPU?
• RQ5 What is the relative performance gain of position masking on IPU versus GPU platforms in terms of F1 score and convergence time?

Key findings

• Position masking improved SQuAD v1.1 F1 score by 0.3% and EM by 0.4% on the IPU, compared to standard BERT training.
• On the IPU, convergence time was reduced to 50% of the original BERT paper’s token count, with position masking enabling faster convergence.
• The GPU results showed a 0.3% F1 improvement on SQuAD, consistent with the IPU results, though peak performance was not reached due to sequence length constraints.
• A 10% position masking rate yielded the best performance, with 15% masking leading to equivalent results to no masking, and 5% showing similar gains to 10%.
• The enhanced dropout gradient technique improved SQuAD performance by 0.5% over standard dropout, suggesting a potential regularization effect.
• The IPU outperformed GPU in both convergence speed and steady-state performance, with a 10% faster convergence when using position masking.