Simple and Effective Multi-Token Completion from Masked Language Models (2023)

Abstract: Pre-trained neural masked language models are often used for predicting a replacement token for a given sequence position, in a cloze-like task. However, this usage is restricted to predicting a single token, from a relatively small pre-trained vocabulary. Recent Sequence2Sequence pre-trained LMs like T5 do allow predicting multi-token completions, but are more expensive to train and run. We show that pre-trained masked language models can be adapted to produce multi-token completions, with only a modest addition to their parameter count. We propose two simple adaptation approaches, trading parameter counts for accuracy. The first method generates multi-token completions from a conditioned RNN. It has a very low parameter count and achieves competitive results. The second method is even simpler: it adds items corresponding to multi-token units to the output prediction matrix. While being higher in parameter count than the RNN method, it also surpasses current state-of-the-art multi-token completion models, including T5-3B, while being significantly more parameter efficient. We demonstrate that our approach is flexible to different vocabularies and domains and can effectively leverage existing pre-trained models available in different domains. Finally, a human evaluation further validates our results and shows that our solution regularly provides valid completions, as well as reasonable correctness for factual-sentence completions.

## introduction multi-token-completion (mtc) is the task of filling masked sentences with a sequence of tokens such that the completed sentence is probable and coherent. e.g., for the masked sentence "the 46th president of the us, @xcite , was elected in 2020" * authors contributed equally. while masked language models (mlms) such as bert @xcite and roberta @xcite successfully deal with a simpler variation of this task (single-token completion) these models were pre-trained on a limited vocabulary, not containing multi-word phrases. it is technically possible to complete numerous tokens simultaneously with mlms by introducing a sequence of @xcite tokens, but there are no clear methods of conditioning the tokens on each other, and, more importantly, the length of the completion needs to be pre-determined. expanding mlms' effective completion vocabulary -the actual vocabulary they are capable of completing well -will help a recent line of work using mlms to extract knowledge/information from corpora @xcite @xcite . most existing works using mlms to complete sentences either avoid the problem by limiting the completions to single tokens or use sub-optimal heuristics for mtc (e.g. presetting the length of the sequence and filling one token at a time). a trivial solution is to increase the number of tokens used in the pre-processing tokenization step, and also include multi-word phrases as tokens. however, apart from the longer tokenization time foot_0 , the main problem with this approach is that changing the input level of the model requires adapting the weights of the full mlm to the new input. 2 a recent family of seq2seq pre-trained lms, the prevalent of which is t5 @xcite , are trained directly on the mtc task, and indeed perform quite well on it, especially with larger model sizes. however, the seq2seq objective, coupled with very high parameter counts, make such models expensive to train and to run inference on, compared to mlms. in this work, we demonstrate how pre-trained mlms can be adapted to produce multi-token completions from (large) fixed vocabularies, using a self-supervised training objective and only a modest parameter count. specifically, we demonstrate an effective adaptation of pre-trained mlms to predict, on top of their pre-trained vocabularies, additional ∼ 100k noun-phrases (np-chunks) and entities ranging in length from 1 to 10 tokens. the only requirement is a textual corpus in which each of the desired phrases appears over k times (we used @xmath0. the adapted model surpasses the accuracy of t5-3b predictions, while using a fraction of the parameter count and being significantly more efficient to run. apart from evaluating with automatic measures, we also use human-annotators to explore different aspects of the proposed completions. after experimenting on the prediction of generalpurpose phrases, we also show our methods can be used to predict phrases in a specific domain. this is a significant benefit of the adaptation approach since the completion vocabularies can be tailored to a specific project's needs, using domain-specific pre-trained mlms, even in domains where huge t5-like models are not available 3 . we use any pre-trained mlm for mtc by extracting the informative representations the mlm uses for completing a single token, and feeding them into a small and simple model that chooses appropriate multi-token completions. we offer two different completion models. depending on the scenario, both are useful; the first solution, which expands the mlm's decoder matrix, achieves sota accuracy. the second solution, using a small iterative generation model, is well suited for large completion-vocabularies while achieving competitive results. the core reason for the success of our methods, despite being trained on less data and for less time than other solutions, is the fact that the mlm's pre-training is well suited for the mtc task. we provide two types of evidence to support this claim: in section 2 we present an experiment suggesting that mlms incorporate information regarding multi token phrases. in section 6, we show that pre-training on in-domain data helps our methods perform better. both of these signals indicate that, as expected, the mlm pre-training captures information required for mtc. therefore, what is left for our mtc adaptation models is only to extract this information from the mlm, an easier task that does not require a long training with a vast amount of data. our main contributions are mtc datasets (general purpose, pubmed-based, and a 3k dataset with human labeling), demonstrating that mlms learn the meaning of multi token phrases, and two methods for adapting any pre-trained mlms for mtc. we publish foot_3 our datasets, models and code foot_4 . ## mlms learn multi-token phrases our work relies on the assumption that mlms are capable of holding information about multi-token phrases. despite being a common belief and an essential property for the mlms to have in order to be as successful as they are, to the best of our knowledge this was not shown explicitly. how do we show that a token-completion mlm is encoding the semantics of multi-token phrases? looking at completions will fail, because we cannot directly probe for mtc. instead, we came up with the following experimental design, which measures the encoding of multi-token phrases indirectly by looking at masked positions influenced by them. we use multi-token phrases that have single token synonyms. we construct a dataset containing quadruples, each quadruple consisting of a multitoken phrase (e.g. "new york city"), a single-token synonym (e.g. "nyc"), a phrase similar in meaning to the synonyms (e.g. "chicago") and a random phrase (e.g. "dog"). we collect sentences containing the multi-token phrases, and in each sentence mask out an np-chunk different from the multitoken phrase (e.g. "[mask] is in new york city"). next, we replace each multi-token phrase with each of its other corresponding quadruple phrases, forming four similar masked sentences, each containing a different quadruple phrase (e.g. "[mask] is in nyc"). finally, we ask an mlm to complete the missing token in each masked sentence, and compare the suggested completions. if the mlm is aware that the multi-token phrase is semantically similar to the single-token synonym, we expect it to suggest similar completions regardless of which of the two it sees in the sentence. indeed, the model treats the multi-token phrase similar to the single-token synonym; we use a similarity measure comparing the mlm's suggested completions for the multitoken phrase sentences and each of the other types of sentences. we find that the average similarity between multi-token phrase sentences and singletoken synonym sentences is 0.76, while for similarphrase sentences it is 0.71 and for random-phrase sentences it is 0.63. our main conclusion from this experiment is that, as expected, mlms learn semantic meaning of multi-token phrases. full experiment details (including the similarity measure used) are in appendix b. ## masked language modeling data in absence of a known dataset for the mtc task, we curate one. we follow @xcite and use data from two sources: wikipedia articles and the books corpus @xcite . completion vocabulary. we start by building a vocabulary of phrases. as we aim for generalpurpose mtc, and following @xcite by considering phrases as np-chunks or entities, we simply focus on phrases that appear frequently in the corpus. we use spacy @xcite to extract 64m unique np-chunks and entities. we keep phrases appearing 500 times or more in the corpus, leaving us with ∼ 93k phrases. only 10.2% are single tokens using bert's cased tokenizer, indicating the importance of mtc. 52.9% of the phrases consist of 2 tokens, 36.9% consist of 3 or more. 53% of the phrases are single-words, 47% span two words or more. further statistics regarding the number of words and the number of tokens assembling the phrases selected are reported in masked sentences. we split the corpus using spacy's sentencer. we sample 50 unique sentences containing each vocabulary-phrase. we eliminate recurring sentences and mask out the vocabulary phrase to form masked sentences, using the masked span as the label. we randomly split the data into train (90%), validation (5%) and test (5%) sets. ## adapting mlms for mtc we propose two simple yet effective solutions, capable of integrating with any pretrained mlm and extending its completion vocabulary to perform mtc. both solutions utilize the mlm by using the contextual embedding vector of the masked place, which was originally pre-trained to be relevant for completion tasks. instead of using it to choose the appropriate single token replacement (see figure 1 (1)) like the mlm, we use it as an input to extension models tuned for completing out of an expanded vocabulary. we train only the relatively small extensions and use the pre-trained mlms as is, allowing short and effective training. ## results setup. measuring success in completion tasks is not trivial since a masked-sentence can have many suitable replacements. in many cases returning the expected phrase in a top-k place, and not necessarily first, is acceptable. hence, we measure accuracy@k (sklearn): the percentage of masked sentences where the label is among the top-k predictions. to better ground our performance expectation from mtc, we first evaluate the accuracy of an easier completion task; single token completion. we go on to evaluate naive mtc using mlms. as expected, results are not sufficient. finally, we evaluate our methods, adapting mlms for mtc. ## domain specific mtc we investigate mtc in a specific domain. datasets. we chose the biology domain, due to the availability of data and models. using pubmed (pubmed) abstracts as a corpus, we extract two mtc datasets. similarly to section 3, we first construct the completion vocabularies. key-phrase vocabulary: we use ∼ 20k mesh vocabulary foot_6 phrases appearing frequently as keyphrases in pubmed papers. this assures us phrases from the biology domain. we discard phrases appearing less than 50 times as np-chunks or entities in the corpus, leaving us with ∼ 12.5k phrases. frequent-phrases vocabulary: similarly to section 3, we extract a vocabulary of phrases appearing more than 500 times as np-chunks or entities in the corpus (∼ 144k phrases). to make sure we are considering mostly phrases from the biology domain, we discard phrases appearing in our general purpose vocabulary (described in section 3), leaving us with ∼ 118k phrases. finally, for both vocabularies we extract 50 sentences containing each of the vocabulary phrases, mask the phrase out in each sentence and split the data into train, development and test sets. base models. to evaluate the impact that the domain-specific pre-training has on the performance of our methods, we test our adaptation methods on top of models pretrained on different datasets. scibert @xcite is a bert-like model trained on the semantic scholar data -data from the scientific domain, closer to the biology domain than the general purpose bert. biobert @xcite is a bert-like model trained on pubmed abstracts. results are reported in conclusions. first, our method outperforms the strong t5-3b baseline in the domain-specific scenario as well, showing mlms can be effectively adapted for domain-specific mtc. second, the success of biobert and scibert suggests that the pretraining of the mlm is in fact utilized by our methods (i.e. there is no catastrophic forgetting). last, the key-phrases dataset is easier than the frequent np-chunks and entities one. this is likely due to the key-phrases dataset mostly containing phrases that are more frequent and significant. ## human evaluation completion tasks are difficult to evaluate since a sentence can have many different valid completion options including, but not limited to, the original masked span. we use accuracy@k to deal with this issue, but in some use-cases only the first completion is important. thus, we define a manual task to more accurately evaluate sentence completions. human-annotators are presented with the masked sentence, the original masked span and the first completion suggested by each of the methods we evaluate (full annotation task is in appendix d). annotators were not aware which model provided each completion, and the completions were presented in a random order for each sentence. the sentences are divided between three expert englishspeaking annotators. prior to performing the annotations, the annotators met for a calibration session, each annotating the same 50 sentences and discussing results until reaching high agreement. we sample 750 sentences from the test set described in section 3. each expert annotator is assigned with 250 sentences and the completions of each of the four methods, amounting to 1000 samples each. first, we ask annotators whether the completions are grammatically correct and make sense in the context of the sentence, regardless of whether they are factually correct (denoted as a "valid" completions). for example, any year or location is a valid completion for "jules verne was born in [mask]". in some cases, valid completions can be general and uninformative. for example, in the sentence "on january 1, 1998, [mask] was released publicly online as sgi freeware." the original span is "blender". while any sgi freeware name would make sense, the word "it" (suggested by t5-3b) is also a valid completion. this happens with general words like "he", "she", or "person" but can also occur with entities, e.g. in the sentence "jules verne was born in [mask]", the original masked span is "nantes" and a suggested completion might be "europe". we ask annotators to flag these cases, where the suggested completion is more general (less specific) than the original span. factual correctness of completions is also important, since some works use completion methods for knowledge extraction @xcite . to evaluate this, we first ask annotators to mark whether the masked span is a part of a specific fact. for example, "mlk, (born [mask])" should be marked as a fact, while "he died aged [mask]" shouldn't. we find that 50.93% of the masked sentences are part of a fact foot_7 . for the factual sentences, we ask the annotators to label whether the proposed completion is factually correct. if they do not know, annotators are instructed to use the web to verify the correctness. results are reported in correct is the percentage of sentences labeled as facts that are valid completions and are also factually correct. * marks that the difference between best and second best is statistically significant. ## related work transformer-based lms and mlms @xcite have revolutionized nlp in the past couple of years. while most of the impact has been achieved using these pretrained models as a source of meaningful contextual embeddings, recent works are using these models for the task they were pretrained for: masked language modeling @xcite @xcite @xcite . while most mlms are capable of completing multiple missing tokens simultaneously, they do so in an unconditional manner, yielding unsatisfying results. therefore, some works using mlms for masked language modeling simply restrict themselves to single token completions @xcite while others use heuristics in order to generate multi token completions @xcite . some attempts towards making mlms predict multi token units better by changing the masking technique during training include masking spans instead of tokens @xcite and masking whole words @xcite , but these methods still complete in an unconditioned manner during inference. xl-net @xcite avoids interdependence between masked tokens using a permutation language modeling pretraining objective. similarly to our solution, extending the decoder matrix with n-gram vocabulary items is done also by @xcite , but for a different purpose: they use them as an auxiliary signal during training, in order to improve bert's single-token pretraining. they mask n-grams, and predict a specific n-gram assigned id together with the single tokens assembling the n-gram, as to improve the singletoken embeddings. the bert model is pretrained end to end and the n-gram embeddings are discarded after pretraining and are not part of their final model used during fine-tuning and inference. while we share their extension of the decoder matrix, the mtc task adds additional requirements: to allow for easy mtc vocabulary support, we aim for a short mtc training and inference, building on the preexisting pretrained mlms and avoiding the need to pretrain from scratch. in addition, our solution does not distinguish between multi-token and single token-phrases during training since both are part of the completion vocabulary. in @xcite , the ilm framework is introduced. this is a framework for infilling, a task similar to mtc. ilm is a gpt-2 @xcite based solution, meaning the pretrained model used is a lm and not an mlm, requiring it to adapt to a different task. finally, t5 @xcite was pre-trained for the mtc task using a vast amount of data. it performs well but requires a lot of training data, time and memory, does not utilize existing mlms and is not available in many domains and languages. ## conclusions we show mlms can be adapted for multi token completion even though they were trained for single token completion. we presented two simple but effective solutions that leverage the pretrained mlms and offer quick adaptations to new vocabularies. the two solutions are trading-off performance and size: 1) an extended matrix decoder offering sota accuracy but size-dependent on the completion vocabulary; 2) an rnn decoder with slightly lower accuracy but size independent of the vocabulary size. we also demonstrate the flexibility of our approach to different vocabularies and domains by evaluating it on the pubmed dataset and showing that leveraging domain-pretrained mlms offers significant accuracy improvement. finally, we validate our results by conducting a human evaluation to account for valid completions that are not measured using our automatic metric. it shows that our extended matrix solution provides valid completions in 84% of the times, and can also correctly handle facts in 30% of the times, comparatively to the much larger t5-3b.