PMI-Masking: Principled masking of correlated spans

Masking tokens uniformly at random constitutes a common flaw in the pretraining of Masked Language Models ( MLMs ) such as BERT . We show that such uniform masking allows an MLM to minimize its training objective by latching onto shallow local signals , leading to pretraining inefficiency and suboptimal downstream performance . To address this flaw , we propose PMI-Masking , a principled masking strategy based on the concept of Pointwise Mutual Information ( PMI ) , which jointly masks a token n-gram if it exhibits high collocation over the corpus . PMIMasking motivates , unifies , and improves upon prior more heuristic approaches that attempt to address the drawback of random uniform token masking , such as whole-word masking , entity/phrase masking , and random-span masking . Specifically , we show experimentally that PMI-Masking reaches the performance of prior masking approaches in half the training time , and consistently improves performance at the end of training . 1 INTRODUCTION . In the couple of years since BERT was introduced in a seminal paper by Devlin et al . ( 2019a ) , Masked Language Models ( MLMs ) have rapidly advanced the NLP frontier ( Sun et al. , 2019 ; Liu et al. , 2019 ; Joshi et al. , 2020 ; Raffel et al. , 2019 ) . At the heart of the MLM approach is the task of predicting a masked subset of the text given the remaining , unmasked text . The text itself is broken up into tokens , each token consisting of a word or part of a word ; thus “ chair ” constitutes a single token , but out-of-vocabulary words like “ e-igen-val-ue ” are broken up into several sub-word tokens . In BERT , 15 % of tokens are chosen to be masked uniformly at random . It is the random choice of single tokens that we address in this paper : we show that this approach is suboptimal and offer a principled alternative . To see why Random-Token Masking is suboptimal , consider the special case of sub-word tokens . Given the masked sentence “ To approximate the matrix , we use the eigenvector corresponding to its largest e- [ mask ] -val-ue ” , an MLM will quickly learn to predict “ igen ” based only on the context “ e [ mask ] -val-ue ” , rendering the rest of the sentence redundant . The question is whether the network will also learn to relate the broader context to the tokens comprising “ eigenvalue ” . When they are masked together , the network is forced to do so , but such masking occurs with vanishingly small probability . One might hypothesize that the network would nonetheless be able to piece such meaning together from local cues ; however , we show that it often struggles to do so . We establish this via a controlled experiment , in which we reduced the size of the vocabulary , thereby breaking more words into sub-word tokens . We compared the extent to which such vocabulary reduction degraded regular BERT relative to so-called Whole-Word Masking BERT ( WWBERT ) ( Devlin et al. , 2019b ) , a version of BERT that jointly masks all sub-word tokens comprising an out-of-vocabulary word during training . We show that vanilla BERT ’ s performance degrades much more rapidly than that of WWBERT as the vocabulary size shrinks . The intuitive explanation is that Random-Token Masking is wasteful ; it overtrains on easy sub-word tasks ( such as predicting “ igen ” ) and undertrains on harder whole-word tasks ( predicting “ eigenvalue ” ) . The advantage of Whole-Word Masking over Random-Token Masking is relatively modest for standard vocabularies , because out-of-vocabulary words are rare . However , the tokenization of words is a very special case of a much broader statistical linguistic phenomenon of collocation : the cooccurrence of series of tokens at levels much greater than would be predicted simply by their individual frequencies in the corpus . There are millions of collocated word n-grams — multi-word expressions , phrases , and other common word combinations — whereas there are only tens of thousands of words in frequent use . So it is reasonable to hypothesize that Random-Token Masking generates many wastefully easy problems and too few usefully harder problems because of multiword collocations , and that this affects performance even more than the rarer case of tokenized words ; we show that this indeed is the case . Several prior works have considered the idea of masking across spans longer than a single word . Sun et al . ( 2019 ) and Guu et al . ( 2020 ) proposed Knowledge Masking and Salient Span Masking , respectively , in which tokens comprising entities or phrases , as identified by external parsers , are jointly masked . While extending the scope of Whole-Word Masking , the restriction to specific types of correlated n-grams , along with the reliance on imperfect tools for their identification , has limited the gains achievable by these approaches . With a similar motivation in mind , SpanBERT of Joshi et al . ( 2020 ) introduced Random-Span Masking , which masks word spans of lengths sampled from a geometric distribution at random positions in the text . Random-Span Masking was shown to consistently outperform Knowledge Masking , is simple to implement , and inspired prominent MLMs ( Raffel et al. , 2019 ) . However , while Random-Span Masking increases the chances of masking collocations , with high probability the selected spans break up correlated n-grams , such that the prediction task can often be performed by relying on local cues . In this paper we offer a principled approach to masking spans that consistently provide high signal , unifying the intuitions behind the above approaches while also outperforming them . Our approach , dubbed PMI-Masking , uses Pointwise Mutual Information ( PMI ) to identify collocations , which we then mask jointly . At a high level , PMI-Masking consists of two stages . First , given any pretraining corpus , we identify a set of contiguous n-grams that exhibit high cooccurrence probability relative to the individual occurrence probabilities of their components . We for- malize this notion by proposing an extended definition of Pointwise Mutual Information from bigrams to longer n-grams . Second , we treat these collocated n-grams as single units ; the masking strategy selects at random both from these units and from standard tokens that do not participate in such units . Figure 1 , detailed and reinforced by further experiments in section 5 , shows that ( 1 ) PMI-Masking dramatically accelerates training , matching the end-of-pretraining performance of existing approaches in roughly half of the training time ; and ( 2 ) PMI-Masking improves upon previous masking approaches at the end of pretraining . 2 MOTIVATION : MLMS ARE SENSITIVE TO TOKENIZATION . In this section we describe a simple experiment that motivates our PMI-Masking approach . We examined BERT ’ s ability to learn effective representations for words consisting of multiple subword tokens , treating this setting as an easily controlled analogue for the multi-word collocation problem that truly interests us . Our experiment sought to assess the performance gain obtained from always masking whole words as opposed to masking each individual token uniformly at random . We compared performance across a range of vocabulary sizes , using the same WordPiece Tokenizer1 that produced the original vocabulary of ∼ 30K tokens . As we decreased a 30K-token vocabulary to 10K and 2K tokens , the average length of a word over the pretraining corpus increased from 1.08 tokens to 1.22 and 2.06 tokens , respectively . Thus , by reducing the vocabulary size , we increased the frequency of multi-token words by a large factor . Table 1 presents the performance of BERT models trained with these vocabularies , measured as score on the SQuAD2.0 development set ( the experimental setup is described in section 4 ) . The downstream performance of Random-Token Masking substantially degraded as vocabulary size decreased and the number of spans of sub-word tokens increased . One reason for such degradation might be the model seeing less text as context ( 512 input tokens cover less text when more words are broken into multiple tokens ) . This possibly plays a role ; however , for models with the same vocabularies trained via Whole-Word Masking , this degradation was significantly attenuated . We therefore conjecture that this degradation occurred primarily because of the random masking strategy , which allows the model to use “ shortcuts ” for minimizing its loss , thus hindering its ability to learn the distribution of the entire multi-token word . If our conjecture is correct , such shortcuts are just as problematic in the case of inter-word collocations . In fact , for the regular 30K-token vocabulary , divided words are rare , so inter-word collocations would pose a larger problem than intra-word collocations in the common setting . One possible mitigation might be to expand the vocabulary to include multi-word collocations . However , there are millions of these , and such vocabulary sizes are currently infeasible . Even if we could get around the practical issue of size , this approach may suffer from generalization problems : the frequency of each multi-word collocation can be lower than the sample complexity for learning a meaningful representation . An alternative , more practical approach is to leave the vocabulary as is , but jointly mask co-located words , with the intention of cutting off local statistical “ shortcuts ” and allowing the model to improve further by learning from broader context . This is the approach we take in this paper . In what follows we detail such a masking approach and show its advantages experimentally . 3 MASKING CORRELATED n-GRAMS 3.1 EXISTING MASKING APPROACHES . We now more formally present the MLM setup as well as existing masking approaches , which we implement as baselines . Given text tokenized into a sequence of tokens , Masked Language Models are trained to predict a set fraction of “ masked ” tokens , where this fraction is called the masking budget and is traditionally set to 15 % . The modified input is inserted into the Transformer-based architecture ( Vaswani et al. , 2017 ) of BERT , and the pretraining task is to predict the original identity of each chosen token . Several alternatives have been proposed for choosing the set of tokens to mask . Random-Token Masking ( Devlin et al. , 2019a ) The original BERT implementation selects tokens for masking independently at random , where 80 % of the 15 % chosen tokens are replaced with [ MASK ] , 10 % are replaced with a random token , and 10 % are kept unchanged . Whole-Word Masking ( Devlin et al. , 2019b ) The sequence of input tokens is segmented into units corresponding to whole words . Tokens for masking are then chosen by sampling entire units at random until the masking budget is met . Following Devlin et al . ( 2019a ) , for 80 % /10 % /10 % of the units , all tokens are replaced with [ MASK ] tokens/ random tokens/ the original tokens , respectively . 1https : //github.com/huggingface/tokenizers Random-Span Masking ( Joshi et al. , 2020 ) Contiguous random spans are selected iteratively until the 15 % masking budget is spent . At each iteration , a span length ( in words ) is sampled from a geometric distribution ` ∼ Geo ( 0.2 ) , and capped at 10 words . Then , the starting point for the span to be masked is randomly selected . Replacement with [ MASK ] , random , or original tokens is done as above , where spans constitute the units . 3.2 PMI : FROM BIGRAMS TO n-GRAMS Our aim is to define a masking strategy that targets correlated sequences of tokens in a principled way . Of course , modeling such correlations in large corpora was widely studied in computational linguistics ( Zuidema ( 2006 ) ; Ramisch et al . ( 2012 ) ; inter alia ) . Particularly relevant to our work is the notion of Pointwise Mutual Information ( Fano , 1961 ) , which quantifies how often two events occur , compared with what we would expect if they were independent . Define the probability of any n-gram as the number of its occurrences in the corpus divided by the number of all the n-grams in the corpus . PMI leverages these probabilities to give a natural measure of collocation of bigrams : how surprising the bigram w1w2 is , given the unigram probabilities of w1 and w2 . Formally , given two tokens w1 and w2 , the PMI of the bigram “ w1w2 ” is PMI ( w1w2 ) = log p ( w1w2 ) p ( w1 ) p ( w2 ) . ( 1 ) Importantly , PMI is qualitatively different from pure frequency : a relatively frequent bigram may not have a very high PMI score , and vice versa . For example , the bigram “ book is ” appears 34772 times in the WIKIPEDIA+BOOKCORPUS dataset but is ranked around position 760K in the PMI ranking for bi-grams over this corpus , while the bigram “ boolean algebra ” appears 849 times in the corpus but is ranked around position 16K in the PMI ranking . What about contiguous spans of more than two tokens ? For a given n-gram , we would again like to measure how strongly its components indicate one another . We thus require a measure that captures correlations among more than two variables . A standard and direct extension of the PMI measure to more than two variables , referred to as ‘ specific correlation ’ in Van de Cruys ( 2011 ) , and as ‘ NaivePMIn ’ in this paper , is based on the ratio between the n-gram ’ s probability and the probabilities of its component unigrams : Naive-PMIn ( w1 . . . wn ) = log p ( w1 . . . wn ) ∏n j=1 p ( wj ) ( 2 ) As in the bivariate case , this measure compares the actual empirical probability of the n-gram in the corpus with the probability it would have if its components occurred independently . However , the above definition suffers from an inherent flaw : an n-gram ’ s Naive-PMIn will be high if it contains a segment with high PMI , even if that segment is not particularly correlated with the rest of the n-gram . Consider for example the case of trigrams : Naive-PMI3 ( w1w2w3 ) = log { p ( w1w2w3 ) p ( w1 ) p ( w2 ) p ( w3 ) · p ( w1w2 ) p ( w1w2 ) } = PMI ( w1w2 ) + log p ( w1w2w3 ) p ( w1w2 ) p ( w3 ) Where PMI ( w1w2 ) is defined in eq . 1 . When PMI ( w1w2 ) is high , the Naive-PMI3 measure of the trigram “ w1w2w3 ” will start at this high baseline . The added term of log p ( w1w2w3 ) p ( w1w2 ) p ( w3 ) quantifies the actual added information of “ w3 ” to this correlated bigram , i.e. , it quantifies how far p ( w1w2w3 ) is from being separable w.r.t . the segmentation into “ w1w2 ” and “ w3 ” . For example , since the PMI of the bigram “ Kuala Lumpur ” is very high , the Naive-PMIn of the trigram “ Kuala Lumpur is ” is misleadingly high , placing it at position 43K out of all trigrams in the WIKIPEDIA+BOOKCORPUS dataset . It is in fact placed much higher than obvious collocations such as the trigram “ editor in chief ” , which is ranked at position 210K out of all trigrams . In order to favor n-grams that can not be easily subdivided into shorter unrelated spans , we propose a measure of distance from separability with respect to all of an n-gram ’ s possible segmentations rather than with respect only to the segmentation into single tokens : PMIn ( w1 . . . wn ) = min σ∈seg ( w1 ... wn ) log p ( w1 . . . wn ) ∏ s∈σ p ( s ) ( 3 ) Here , seg ( w1 . . . wn ) is the set of all contiguous segmentations of the n-gram “ w1 . . . wn ” ( excluding the identity segmentation ) , where any segmentation σ ∈ seg ( w1 . . . wn ) is composed of sub-spans which together give “ w1 . . . wn ” . Intuitively , this measure effectively discards the contribution of high PMI segments ; the minimum in Eq . 3 implies that an n-gram ’ s collocation score is given by its weakest link , i.e. , by the segmentation that is closest to separability . When ranked by the above PMIn measure , the trigram “ Kuala Lumpur is ” is demoted to position 1.6M , since the segmentation into “ Kuala Lumpur ” and “ is ” yields unrelated segments , while the trigram “ editor in chief ” is upgraded to position 33K since its segmentations yield correlated components . As we will see , this definition is not only conceptually cleaner , but also leads to improved performance .

Summary: This paper proposes an improvement to how tokens are selected for masking in pre-training large masked language models (BERT and family). Specifically, it stipulates that purely random choice of words (or word pieces) makes the MLM task insufficiently hard. It then goes on to propose a data-driven approach for selecting n-grams to mask together. The approach, based on an extension of pointwise mutual information for n-grams, is shown to outperform random token and random spans masking strategies on performance of downstream tasks.