An Analysis of Simple Data Augmentation for Named Entity Recognition

In plain words

Named entity recognition finds mentions of things in text and labels each one with its type, and it needs labelled examples. In specialist fields, labelling is slow because only experts can do it, so the labelled sets are small. A common fix is data augmentation: make extra training examples by changing existing ones in ways that keep their meaning.

For sentence-level tasks, such tricks are well known. For entity recognition they are harder, because every word carries a label, so changing a word can break the label. This paper adapts a handful of simple tricks that do not need any other trained model. It swaps a word for another word with the same label, swaps a word for a WordNet synonym, swaps a whole entity for another entity of the same type, or shuffles the words inside a span that shares one label.

The authors test these tricks on two specialist data sets, one on patient reports and one on materials-science text, using two kinds of tagger. Their main finding is that the tricks help most when the training set is small, and that combining them helps more than any one of them alone on average.

TL;DR

• Four token-level augmentations for NER that do not rely on machine translation or parsing models §3.
• Label-wise token replacement and shuffle within segments keep the label sequence unchanged; synonym and mention replacement can change it tab. 1. All four can be combined §3.
• Author-reported: on the small training sets (50 sentences), every method gives a significant gain on both data sets for the LSTM tagger tab. 3.
• The gains shrink on the full training set, and some single methods lower the score there tab. 3 §4.4.
• Results come from one author-reported setup and five random seeds; some gains are smaller than one standard deviation tab. 3.

Problem

Modern neural taggers need many labelled sentences. In biomedical and materials-science text, expert annotation is slow and costly, so the labelled sets are small §1. Data augmentation adds transformed copies of training examples without changing their labels §1. Text-classification and translation methods use augmentation, but NER predicts a label for every token, so a change to one token can change that token's label §1. The paper asks which simple augmentations still work for NER and under which training-set sizes §1 §4.4.

Method

The paper surveys word replacement, mention replacement, word swapping and generative augmentation for sentence-level tasks §2. It then adapts four methods to token-level labels, with no external models §3. Each method takes a sentence with BIO labels and produces a new sentence with BIO labels. A replacement is decided per unit (token, mention or segment) with a binomial draw with probability p §3. The paper's Table 1 shows one sentence transformed by each method tab. 1.

Labels. Each token carries a BIO tag combined with an entity type §4.2. Table 1 shows one example sentence and one augmented variant per method, with changed tokens highlighted and the label row shown for each variant tab. 1.

Label-wise token replacement (LwTR) §3. For each token, draw a binomial variable with probability p. If it is 1, replace the token with another token drawn from a distribution over tokens that share the same label, built from the original training set. The label sequence stays unchanged.

Synonym replacement (SR) §3. The same token-level draw as LwTR, but the replacement is a WordNet synonym of the token. A synonym can have several words. The first word keeps the token's label. If the token starts a mention (B), the other words get I of the same type. If the token is inside a mention (I), all words get that label.

Mention replacement (MR) §3. For each mention, draw a binomial variable with probability p. If it is 1, replace the whole mention with another mention of the same entity type from the training set. The BIO labels are rewritten to match the new mention's length.

Shuffle within segments (SiS) §3. Split the sentence into segments. Each mention is one segment and each run of out-of-mention tokens is one segment. Each segment is shuffled with probability p. Only the token order inside a segment changes; the label order stays the same.

All §3. Create one augmented instance per method for each training instance. The paper does not describe any other way of combining the methods.

# Label-wise token replacement (LwTR)  [§3]  [code: augment.py#L156-L171]
def label_wise_token_replace(sentence, token_dist, p):
    # Keep stopwords and unselected tokens as they are  [code: augment.py#L163]
    augmented = []
    for token, label in sentence:
        if is_stopword(token) or not bernoulli(p):
            augmented.append((token, label))
        else:
            # Draw a token that already carries this full BIO label  [§3]
            new_token = sample(token_dist[label])
            augmented.append((new_token, label))
    return augmented

# Synonym replacement (SR)  [§3]  [code: augment.py#L174-L206]
def synonym_replace(sentence, p):
    augmented = []
    for token, label in sentence:
        # Stopwords and unselected tokens are kept  [code: augment.py#L181]
        if is_stopword(token) or not bernoulli(p):
            augmented.append((token, label))
            continue
        # Pool WordNet lemma names of all senses  [code: augment.py#L186-L192]
        synonyms = wordnet_synonyms(token) - {token}
        if len(synonyms) == 0:
            augmented.append((token, label))
            continue
        # Pick one synonym at random  [§3]  [code: augment.py#L199]
        words = random_choice(list(synonyms)).split()
        for i, word in enumerate(words):
            # The first word keeps the label; B-x gives I-x to the rest  [§3]
            word_label = label
            if i > 0 and label.startswith("B-"):
                word_label = "I-" + label[2:]
            augmented.append((word, word_label))
    return augmented

# Mention replacement (MR)  [§3]  [code: augment.py#L43-L83]
def mention_replace(sentence, mentions_by_type, p):
    augmented = []
    position = 0
    while position < len(sentence):
        token, label = sentence[position]
        if label == "O":
            # Tokens outside mentions are copied  [code: augment.py#L49-L52]
            augmented.append((token, label))
            position += 1
            continue
        # Find the end of this mention  [code: augment.py#L70-L76]
        end = mention_end(sentence, position)
        if bernoulli(p):
            # Swap in a random training mention of the same type  [§3]
            new_words = random_choice(mentions_by_type[label[2:]])
            augmented += labelled_words(new_words, label[2:])
        else:
            augmented += sentence[position:end]
        position = end
    return augmented

# Shuffle within segments (SiS)  [§3]  [code: augment.py#L86-L131]
def shuffle_within_segments(sentence, p):
    # One segment per mention and per run of O tokens  [§3]
    segments = split_into_segments(sentence)
    augmented = []
    for segment in segments:
        # Shuffle a segment of more than one token with probability p  [§3]
        if len(segment) > 1 and bernoulli(p):
            tokens = random_permutation(segment.tokens)
        else:
            tokens = segment.tokens
        # Labels stay in position; only the tokens move  [§3]
        augmented += zip(tokens, segment.labels)
    return augmented

Hyperparameters

The paper gives the values below in its text and tables. The "code default" column lists what the released code uses when no flag is passed; it is not a paper value. "not stated" means the paper does not give the value.

Data set sizes tab. 2:

Results

Author-reported. Setup: MaSciP and i2b2-2010 §4.1; an LSTM tagger with GloVe and a SciBERT tagger, both with a CRF §4.2; micro-averaged span F1 on the test set, mean and standard deviation over five seeds tab. 3. The Δ column is the mean of the eight differences from the no-augmentation row, which the table bears out for the first two rows. "*" marks a cell the paper underlines, which the caption defines as significantly better than the baseline under a paired t-test at p = 0.05 tab. 3. Columns are S, M, L (50, 150, 500 training sentences) and F (full set) tab. 3.

Headline reading the paper gives §4.4:

• All methods beat the baseline on average, with no single winner across both models. Synonym replacement is the best single method on average for SciBERT; mention replacement is the best single method on average for the LSTM §4.4 tab. 3.
• Applying all methods together beats any single method on average, but on the full training set a single method is sometimes better (MaSciP LSTM, i2b2 SciBERT) §4.4 tab. 3.
• Gains are concentrated in the small sets, and all four single methods are significant there for both models and both data sets §4.4 tab. 3.
• The authors attribute the weaker effect of combining methods on large sets to a trade-off between diversity and label validity §4.4.
• The authors note significant improvements with pretrained transformers and argue that augmentation should also be studied on pretrained models such as BERT §4.4.

Paper vs code

Code paths are relative to the repository at code_commit (raw/code/2010.11683). The repository is the official release (boschresearch/data-augmentation-coling2020, Apache-2.0); its README names the paper and both authors README.md.

• LSTM encoder is missing. The paper's LSTM tagger §4.2 has no counterpart in the code: the encoder asserts embedding_type == "bert" models.py, and the README runs only the BERT path README.md. Only the SciBERT results can be reproduced from this release.
• Optimiser and schedule. The paper gives none §4.3. The code uses AdamW (named "adam") at 3e-5, batch 16, up to 100 epochs, a reduce-on-plateau schedule and early stopping main.py train.py.
• One learning rate for everything. The encoder, the output layer and the CRF share one optimiser at one rate train.py. The paper says nothing about per-part rates.
• Gradient clipping per module. Norm 5.0 is applied to the encoder, the output layer and the CRF separately train.py, not to the whole model. The paper does not mention clipping.
• Output layer. The code puts a hidden linear layer of the BERT width and an output linear layer with no activation between them, then a linear-chain CRF main.py models.py. The paper describes the CRF but not this extra layer §4.2.
• CRF transition constraints. Viterbi decoding forbids invalid BIO transitions models.py models.py. Training uses the unconstrained transition matrix models.py. The paper mentions neither.
• Output label set. The tag dictionary is built from the train, dev and test sets data.py main.py. The paper does not say this.
• Augmentation timing. The paper does not say when augmented instances are made. The code regenerates them at every epoch from the original sentences, with fresh random draws train.py train.py.
• Stopwords. The code never replaces NLTK English stopwords in LwTR or SR augment.py augment.py augment.py augment.py. The paper does not mention stopwords.
• Smoothing exponent. The code exposes an exponent on the LwTR token counts, default 1.0, which leaves the distribution unchanged main.py augment.py. The paper does not describe it.
• Replace probability and instance count. The paper tunes them on dev §4.3. The code defaults to 0.3 and 1 main.py and has no search loop, so the grid must be run externally.
• Subset construction and dev split. The paper's training subsets and 15% dev split §4.1 are not in the code. The README says to copy prepared files into data/ README.md; data/ holds only sample files train.txt.
• Metric. The code matches spans on label, text and token indices train.py data.py. A span with the right text at a different position is a miss. The paper says "string match" without this detail §4.2.
• Shipped result files are not the paper's numbers. results/*.json record --data_folder data, the folder that holds only sample files all.json train.txt. The README table (F1 37.9 for no augmentation) has no stated setup and does not match tab. 3 README.md.
• Unused flag. --variational_dropout is parsed but not used in the files read main.py.
• Model selection. Best-dev checkpoint, then test: the code matches the paper train.py train.py.

Pitfalls

• Many settings are unstated. Optimiser, learning rate, batch size, epochs, dropout, clipping and seeds are missing from the paper §4.3. The code's defaults are the only concrete values, so use them as the reference and say so when you deviate.
• Keep the label set consistent. The paper says a replacement keeps the label. In code, LwTR draws from a distribution keyed by the full BIO tag, so B-type and I-type tokens have separate pools augment.py. Keep that split, or an I- token can be replaced by a B- token.
• Synonym choice depends on set order. SR picks from list(set) augment.py, so the same seed can pick different synonyms across processes unless you sort the candidates first. Sorting is a suggestion for reproducibility, not something the paper or code does.
• SR pools all senses. Lemma names from every WordNet sense of the word are pooled augment.py, and the paper does not say whether senses are filtered by part of speech.
• MR can redraw the same mention. Candidates come from the whole training set and may include the original mention augment.py.
• SiS moves words inside a mention. Mention segments are shuffled too, so "neuropathic pain syndrome" can become "pain neuropathic syndrome" with labels B I I §3 augment.py. The label row is kept, but the entity text changes.
• No truncation. The encoder does not truncate inputs models.py. Sentences longer than the BERT input limit are not handled in the code read.
• Word dropout without rescaling. Whole token vectors are zeroed with a Bernoulli mask and not rescaled models.py.
• Dev and test labels in the output space. The tag set includes labels that appear only in dev or test data.py. For strict hygiene, build it from train only.
• The text is more positive than the table. The text says all methods improve on the baseline §4.4. On the full set, several single-method cells are equal or lower: SciBERT i2b2 F is 87.5 (LwTR), 87.7 (SR), 87.8 (MR), 87.7 (SiS) and 87.2 (All) against 87.8 for no augmentation; LSTM MaSciP F is 76.3 (LwTR) and 76.0 (MR) against 76.4 tab. 3. Treat the F column as no reliable gain.
• Wide spreads on small sets. Some cells have large standard deviations. The SciBERT i2b2 L cell for synonym replacement is 70.5 ± 2.3 against 70.2 ± 0.3 with no augmentation and is not underlined tab. 3. With five seeds, compare methods only within one recipe.
• Data access. Neither data set ships with the code. The paper links portals for both §4.1. Access terms are not stated here.
• The compiled analysis is excluded. The source has an analysis section that is commented out (%\input{7-analysis}) and does not appear in the compiled paper. Its tables are not summarised here.
• Figures. The source's figure is inside a disabled block, so the card cites tables only.

Lineage

• Builds on: word replacement, mention replacement and token swaps from sentence-level NLP work, which the paper surveys §2. The paper adapts them to token labels and drops methods that need external models §3.
• Compared against: the no-augmentation baseline for each model and data set tab. 3.