MELM: Data Augmentation with Masked Entity
Language Modeling for Low-Resource NER

To minimize the amount of generated tokens incompatible with the original labels, we design a labeled sequence linearization strategy to explicitly take label information into consideration during masked language modeling. Specifically, as shown in Figure 2(c), we add the label token before and after each entity token and treat them as normal context tokens. The yielded linearized sequence is utilized to further finetune our MELM so that its prediction is additionally conditioned on the inserted label tokens. Note that, we initialize the embeddings of label tokens with those of tokens semantically related to the label names (e.g., “organization” for $\langle$ B-ORG $\rangle$). By doing so, the linearized sequence is semantically closer to a natural sentence and the difficulty of finetuning on linearized sequence could be reduced Kumar et al. (2020).

Unlike MLM, only entity tokens are masked during MELM fine-tuning. At the beginning of each fine-tuning epoch, we randomly mask entity tokens in the linearized sentence $X$ with masking ratio $\eta$.

Then, given the corrupted sentence $\tilde{X}$ as input, our MELM is trained to maximize the probabilities of the masked entity tokens and reconstruct the linearized sequence $X$:

where $\theta$ represents the parameters of MELM, $n$ is the number of tokens in $\tilde{X}$, $x_{i}$ is the original token in $X$, $m_{i}=1$ if $x_{i}$ is masked and otherwise $m_{i}=0$. Through the above fine-tuning process, the proposed MELM learns to make use of both contexts and label information to predict the masked entity tokens. As we will demonstrate in Section 4.1, the predictions generated by the fine-tuned MELM are significantly more coherent with the original entity label, compared to those from other methods.

To generate augmented training data for NER, we apply the fine-tuned MELM to replace entities in the original training samples. Specifically, given a corrupted sequence, MELM outputs the probability of each token in the vocabulary being the masked entity token. However, as the MELM is fine-tuned on the same training set, directly picking the most probable token as the replacement is likely to return the masked entity token in the original training sample, and might fail to produce a novel augmented sentence. Therefore, we propose to randomly sample the replacement from the top k most probable components of the probability distribution. Formally, given the probability distribution $P(x_{i}|\tilde{X})$ for a masked token, we first select a set $V^{k}_{i}\subseteq V$ of the $k$ most likely candidates. Then, we fetch the replacement $\hat{x}_{i}$ via random sampling from $V^{k}_{i}$. After obtaining the generated sequence, we remove the label tokens and use the remaining parts as the augmented training data. For each sentence in the original training set, we repeat the above generation procedure $R$ rounds to produce $R$ augmented examples.

To increase the diversity of augmented data, we adopt a different masking strategy from train time. For each entity mention comprising of $n$ tokens, we randomly sample a dynamic masking rate $\epsilon$ from Gaussian distribution $\mathcal{N}(\mathrm{\mu,\,\sigma^{2}})$, where the Gaussian variance $\sigma^{2}$ is set as $1/n^{2}$. Thus, the same sentence will have different masking results in each of the $R$ augmentation rounds, resulting in more varied augmented data.

To remove noisy and less informative samples from the augmented data, the generated augmented data undergoes post-processing. Specifically, we train a NER model with the available gold training samples and use it to automatically assign NER tags to each augmented sentence. Only augmented sentences whose predicted labels are consistent with the their original labels are kept. The post-processed augmented training set $\mathbb{D}_{\text{aug}}$ is combined with the gold training set $\mathbb{D}_{\text{train}}$ to train the final NER tagger.

When extending low-resource NER to multilingual scenarios, it is straightforward to separately apply the proposed MELM on language-specific data for performance improvement. Nevertheless, it offers higher potential to enable MELM on top of code-mixing techniques, which proved to be effective in enhancing multilingual learning (Singh et al., 2019; Qin et al., 2020; Zhang et al., 2021). In this paper, with the aim of bridging MELM augmentation and code-mixing, we propose an entity similarity search algorithm to perform MELM-friendly code-mixing.

Specifically, given the gold training sets $\{\mathbb{D}^{\ell}_{\text{train}}~{}|~{}\ell\in\mathbb{L}\}$ over a set $\mathbb{L}$ of languages, we first collect label-wise entity sets $\mathbb{E}^{\ell,y}$, which consists of the entities appearing in $\mathbb{D}^{\ell}_{\text{train}}$ and belonging to class $y$. To apply code-mixing on a source language sentence $X^{\ell_{\text{src}}}$, we aim to substitute a mentioned entity $\mathbf{e}$ of label $y$ with a target language entity $\mathbf{e}_{sub}\in\mathbb{E}^{\ell_{\text{tgt}},y}$, where the target language is sampled as $\ell_{\text{tgt}}\sim\mathcal{U}(\mathbb{L}\setminus\{\ell_{\text{src}}\})$. Instead of randomly selecting $\mathbf{e}_{sub}$ from $\mathbb{E}^{\ell_{\text{tgt}},y}$, we choose to retrieve the entity with the highest semantic similarity to $\mathbf{e}$ as $\mathbf{e}_{sub}$. Practically, we introduce MUSE bilingual embeddings (Conneau et al., 2017) and calculate the entity’s embedding $\texttt{Emb}(\mathbf{e})$ by averaging the embeddings of the entity tokens:

where $\textsc{MUSE}_{\ell_{\text{src}},\ell_{\text{tgt}}}$ denotes the $\ell_{\text{src}}-\ell_{\text{tgt}}$ aligned embeddings and $\mathbf{e}_{i}$ is the $i$-th token of $\mathbf{e}$. Next, we obtain the target-language entity $\mathbf{e}_{sub}$ semantically closest to $\mathbf{e}$ as follows:

$f(\cdot,\cdot)$ here is the cosine similarity function. The output entity $\mathbf{e}_{sub}$ is then used to replace $\mathbf{e}$ to create a code-mixed sentence more suitable for MELM augmentation. To generate more augmented data with diverse entities, we further apply MELM on the gold and code-mixed data. Since the training data now contains entities from multiple languages, we also prepend a language marker to the entity token to help MELM differentiate different languages, as shown in Figure 3.

We conduct experiments on CoNLL NER dataset (Tjong Kim Sang, 2002; Tjong Kim Sang and De Meulder, 2003) of four languages where $\mathbb{L}$ = {English (En), German (De), Spanish (Es), Dutch (Nl)}. For each language $\ell\in\mathbb{L}$, we first sample $N$ sentences from the full training set as $\mathbb{D}^{\ell,N}_{\text{train}}$, where $N\in\{100,200,400,800\}$ to simulate different low-resource levels. For a realistic data split ratio, we also downscale the full development set to $N$ samples as $\mathbb{D}^{\ell,N}_{\text{dev}}$. The full test set for each language is adopted as $\mathbb{D}^{\ell}_{\text{test}}$ for evaluation.

For monolingual experiments on language $\ell$ with low-resource level $N\in\{100,200,400,800\}$, we use $\mathbb{D}^{\ell,N}_{\text{train}}$ as the gold training data, $\mathbb{D}^{\ell,N}_{\text{dev}}$ as the development set and $\mathbb{D}^{\ell}_{\text{test}}$ as the test set. For zero-shot cross-lingual experiments with low-resource level $N\in\{100,200,400,800\}$, we use $\mathbb{D}^{\text{En},N}_{\text{train}}$ as the source language gold training data, $\mathbb{D}^{\text{En},N}_{\text{dev}}$ as the development set and $\mathbb{D}^{\text{De}}_{\text{test}}$, $\mathbb{D}^{\text{Es}}_{\text{test}}$ and $\mathbb{D}^{\text{Nl}}_{\text{test}}$ as target language test sets. Under multilingual settings where $N$ training data from each language is available $(N\in\{100,200,400\})$, we use $\bigcup_{\ell\in\mathbb{L}}\mathbb{D}^{\ell,N}_{\text{train}}$ as the gold training data, $\bigcup_{\ell\in\mathbb{L}}\mathbb{D}^{\ell,N}_{\text{dev}}$ as the development set and evaluate on $\mathbb{D}^{\text{En}}_{\text{test}}$, $\mathbb{D}^{\text{De},}_{\text{test}}$, $\mathbb{D}^{\text{Es}}_{\text{test}}$ and $\mathbb{D}^{\text{Nl}}_{\text{test}}$, respectively.

To elaborate the effectiveness of the proposed MELM, we compare it with the following methods:

Apart from the quantitative results, we further analyze the augmented data to demonstrate the effectiveness of our MELM in maintaining the consistency between the original label and the augmented token. Table 3 presents examples of the top-5 predictions from pretrained MLM, MELM w/o linearize and MELM. As we can see, the pretrained MLM, which does not introduce any design or contraint on data augmentation, tends to generate high-frequency words such as “the”, “he” and “she”, and the majority of generated words do not belong to the original entity class. Being finetuned on NER data with entity-oriented masking, MELM w/o linearize is able to generate more entity-related tokens.

However, without the explicit guidance from entity labels, it is still too difficult for MELM w/o linearize to make valid predictions solely based on the ambiguous context (e.g., both “Pompeo” (PER) and “Reuters” (ORG) are compatible with the context of Example #2), which leads to token-label misalignment. Compared to the above methods, our MELM take both label information and context into consideration, and thus generates more entities that fit into the context and align with the original label as well. Moreover, it is noteworthy that MELM can leverage the knowledge from pretrained model to generate real-world entities that do not exist in the original NER dataset (e.g., “Greenpeace” and “Amnesty”), which essentially increases the entity diversity in training data.

As demonstrated in Lin et al. (2020) and our preliminary experiments in Figure 1, introducing unseen entities can effectively provide more entity regularity knowledge, and helps to improve NER performance. Therefore, we examine the amount of unique entities introduced by different methods. As there might be token-label misalignment in the augmented data, we firstly train an ‘oracle’ NER model on the full CoNLL dataset and then use it to tag training data of MELM and different baseline methods. For each method, we count the total number of unique entities whose labels match the labels assigned by the ‘oracle’ model. As shown in Figure 4, while many augmented entities from MLM-Entity, DAGA and MELM w/o linearize are filtered out due to token-label misalignment, we note that MELM introduces a significantly larger number of unseen entities in the augmented data. Therefore MELM is able to provide richer entity regularity knowledge, which explains its superiority over the baseline methods.