Speed Up Without Losing Accuracy: A Word-Granularity Chinese WoBERT

提速不掉点：基于词颗粒度的中文WoBERT 苏剑林 · 2020-09-18 · Machine Learning language model attention

Currently, most Chinese pretrained models use characters as their basic unit, meaning Chinese sentences are split into individual characters. There are also some multi-granularity Chinese language models, such as Sinovation Ventures' ZEN and ByteDance's AMBERT, but the basic unit in these models is still the character—they merely find ways to incorporate word-level information as well. Pretrained Chinese models that use words as the basic unit are quite rare; as far as I know, only Tencent's UER has open-sourced a word-granularity BERT model, but in practice it doesn't perform particularly well.

So how well does a purely word-based Chinese pretrained model actually perform? Does it have any value? Recently, we pretrained and open-sourced a word-based Chinese BERT model, which we call WoBERT (Word-based BERT, "my BERT!"). Experiments show that this word-based WoBERT has distinct advantages on quite a few tasks—for instance, a noticeable speed-up—while performance essentially doesn't drop, and in some cases even improves. Here we summarize our work.

Open-source repository: https://github.com/ZhuiyiTechnology/WoBERT

Characters or Words?

Is it "characters" or "words" that work better? This is a rather maddening question in Chinese NLP, and there has been some systematic work investigating it. A relatively recent one is ShannonAI's "Is Word Segmentation Necessary for Deep Learning of Chinese Representations?", published at ACL 2019, which concluded that characters almost always outperform words. As mentioned above, current Chinese pretrained models are indeed mostly character-based. So does this mean the question is already settled—characters are simply better?

The reality is far from that simple. Take that ShannonAI paper as an example: its experimental results aren't wrong, but they're not representative. Why? Because it compares models whose embedding layers are all randomly initialized. In that setting, for the same task, a word-based model has more embedding parameters and is naturally more prone to overfitting, which hurts performance—you could guess this outcome without even running the experiment. The problem is that when we actually use word-based models, we typically don't initialize them randomly; we usually use pretrained word vectors (and downstream tasks may or may not fine-tune these vectors depending on the situation). That's the typical scenario for word-segmented NLP models, but the paper doesn't compare this scenario, so its conclusions aren't very convincing.

In fact, the phenomenon of "overfitting" is two-sided. We do need to guard against it, but overfitting also indicates that a model has relatively strong fitting capacity—and if we find ways to suppress overfitting, we can obtain a stronger model at the same complexity, or a lower-complexity model with the same performance. One important way to alleviate overfitting is more thorough pretraining. So comparisons that don't involve pretraining are unfair to word-based models, and our WoBERT is precisely what confirms the viability of word-based pretrained models.

The Benefits of Words

It's generally believed that the advantages of using characters as the basic unit are:

1. Fewer parameters, less prone to overfitting;
2. No dependence on a segmentation algorithm, avoiding word-boundary segmentation errors;
3. Less severe sparsity issues—out-of-vocabulary words basically don't occur.

As for the case for using words as the basic unit:

1. Shorter sequences, faster processing;
2. In text generation tasks, it can alleviate the exposure bias problem;
3. Word meanings have less ambiguity, reducing modeling complexity.

You might have some doubts about these benefits of using words. Take point 2, for instance: words can alleviate exposure bias because, in theory, the shorter the sequence, the less severe the exposure bias problem (a word-based model predicting a single $n$-character word in one step is equivalent to a character-based model taking $n$ steps, and all those $n$ steps are recursively dependent on each other, so the exposure bias problem is more severe for character-based models). As for point 3, although polysemous words do exist, the meaning of most words is fairly fixed—at least more so than the meaning of individual characters. This means that a single embedding layer might be enough to properly model word meaning, rather than needing, as with character-based models, several layers of general-purpose modeling just to combine characters into words.

At first glance, the two approaches seem evenly matched, but in fact the advantages of using characters are not necessarily the disadvantages of using words—with a few extra tricks, word-based models can avoid these issues to a considerable degree. For example:

1. Word-based models have more parameters, but this can be alleviated through pretraining, so the problem isn't too serious;
2. Dependence on a segmentation algorithm is indeed an issue, but if we only keep the most common subset of words, then the results from different segmentation tools become quite similar, with little variation;
3. As for boundary segmentation errors, these are hard to avoid entirely, but tasks that require accurate boundaries are really just sequence-labeling tasks—text classification and text generation don't actually need precise boundaries, so this alone isn't grounds for dismissing word-based models;
4. If we add most individual characters to the vocabulary as well, out-of-vocabulary words won't occur either.

So in fact, there are quite a lot of benefits to using words. Aside from sequence-labeling tasks that require very precise boundaries, most NLP tasks won't have any real problem being word-based. That's why we set out to build a word-based BERT model.

Tokenizer

To add Chinese words into BERT, we first need the tokenizer to be able to split out words. Is it enough to just add words to the vocab.txt dictionary? Not quite. BERT's built-in tokenizer forcibly separates Chinese characters with spaces, so even if you add words to the vocabulary, it still won't split out Chinese words. Furthermore, when BERT does word-piece tokenization for English, it uses maximum-matching, which isn't precise enough for Chinese word segmentation either.

To make word segmentation work, we modified BERT's tokenizer by adding a "pre-tokenize" step, which lets us split out Chinese words. Specifically:

1. Add Chinese words to vocab.txt;
2. For an input sentence $s$, first apply pre_tokenize to perform a preliminary segmentation, obtaining $[w_1,w_2,\dots,w_l]$;
3. Iterate over each $w_i$: if $w_i$ is in the vocabulary, keep it as is; otherwise, run BERT's built-in tokenize function on $w_i$ again;
4. Concatenate the tokenize results for each $w_i$ in order to form the final tokenization result.

In bert4keras>=0.8.8, implementing the above change only requires passing one extra argument when constructing the Tokenizer, for example:

tokenizer = Tokenizer(
    dict_path,
    do_lower_case=True,
    pre_tokenize=lambda s: jieba.cut(s, HMM=False)
)

Here pre_tokenize is an externally supplied segmentation function; if none is passed, it defaults to None. For simplicity, WoBERT uses Jieba for word segmentation. We removed the redundant parts of BERT's original vocabulary (e.g., Chinese "words" prefixed with ##) and then added 20,000 additional Chinese words (the 20,000 highest-frequency words from Jieba's own built-in vocabulary). The final vocab.txt for WoBERT has 33,586 entries.

Model Details

The currently open-sourced WoBERT is the Base version, continuing pretraining on top of HIT's open-sourced RoBERTa-wwm-ext, with MLM as the pretraining task. During initialization, each word is split into characters using BERT's built-in tokenizer, and the average of the character embeddings is used to initialize the word embedding.

At this point, the technical essentials of WoBERT have basically all been covered—what remains is the actual training. We trained for 1 million steps on a single 24GB RTX GPU (taking roughly 10 days), with a sequence length of 512, a learning rate of 5e-6, a batch size of 16, and gradient accumulation over 16 steps—equivalent to a batch size of 256 trained for about 60,000 steps. The training corpus was roughly 30-plus GB of general-domain text. The training code has already been open-sourced at the link given at the beginning of this article.

In addition, we're also providing WoNEZHA, which continues pretraining based on Huawei's open-sourced NEZHA; the training details are essentially the same as for WoBERT. NEZHA's model architecture is similar to BERT's, except that it uses relative position embeddings rather than the absolute position embeddings BERT uses, so in theory NEZHA has no upper limit on the length of text it can handle. We provide this word-based WoNEZHA simply to give everyone an additional option.

Model Performance

Finally, let's talk about WoBERT's performance. In short, in our evaluations, compared with BERT, WoBERT essentially doesn't underperform on NLP tasks that don't require precise boundaries—some even show a certain improvement—while showing a clear speed advantage. In one sentence: "faster without losing accuracy."

For example, on two classification tasks from Chinese benchmark leaderboards:

$$\begin{array}{c} \text{text classification effect comparison}\\ {\begin{array}{c|cc} \hline & \text{IFLYTEK} & \text{TNEWS} \\ \hline \text{BERT} & 60.31\% & 56.94\% \\ \text{WoBERT} & \textbf{61.15%} & \textbf{57.05%} \\ \hline \end{array}$$}

\end{array}

We also internally tested quite a few other tasks, and the results were similar, showing that on these NLU tasks WoBERT and BERT are essentially comparable. But in terms of speed, WoBERT has a clear advantage over BERT. The table below compares the speed of the two models when processing texts of different lengths:

$$\begin{array}{c} \text{speed comparison}\\ {\begin{array}{c|ccc} \hline & \text{128} & \text{256} & \text{512} \\ \hline \text{BERT} & \text{1.0x} & \text{1.0x} & \text{1.0x} \\ \text{WoBERT} & \textbf{1.16x} & \textbf{1.22x} & \textbf{1.28x} \\ \hline \end{array}$$}

\end{array}

We also tested the WoBERT+UniLM approach on Seq2Seq tasks (CSL/LCSTS title/summary generation), and the results showed a clear improvement over character-based models:

$$\begin{array}{c} \text{CSL summarization results}\\ {\begin{array}{c|c|cccc} \hline & \text{beam size} & \text{Rouge-L} & \text{Rouge-1} & \text{Rouge-2} & \text{BLEU} \\ \hline \text{BERT} & 1 & 63.81 & 65.45 & 54.91 & 45.52 \\ \text{WoBERT} & 1 & \textbf{66.38} & \textbf{68.22} & \textbf{57.83} & \textbf{47.76} \\ \hline \text{BERT} & 2 & 64.44 & 66.09 & 55.75 & 46.39 \\ \text{WoBERT} & 2 & \textbf{66.65} & \textbf{68.68} & \textbf{58.5} & \textbf{48.4} \\ \hline \text{BERT} & 3 & 64.75 & 66.34 & 56.06 & 46.7 \\ \text{WoBERT} & 3 & \textbf{66.83} & \textbf{68.81} & \textbf{58.67} & \textbf{48.6} \\ \hline \end{array}$$}\\

\\

\text{LCSTS Summary Generation Experimental Results}\\

{$$\begin{array}{c|c|cccc} \hline & \text{beam size} & \text{Rouge-L} & \text{Rouge-1} & \text{Rouge-2} & \text{BLEU} \\ \hline \text{BERT} & 1 & 27.99 & 29.57 & 18.04 & 11.72 \\ \text{WoBERT} & 1 & \textbf{31.51} & \textbf{32.9} & \textbf{21.13} & \textbf{13.74} \\ \hline \text{BERT} & 2 & 29.2 & 30.7 & 19.17 & 12.64 \\ \text{WoBERT} & 2 & \textbf{31.91} & \textbf{33.35} & \textbf{21.55} & \textbf{14.13} \\ \hline \text{BERT} & 3 & 29.45 & 30.95 & 19.5 & 12.93 \\ \text{WoBERT} & 3 & \textbf{32.19} & \textbf{33.72} & \textbf{21.81} & \textbf{14.29} \\ \hline \end{array}$$}

\end{array}

This shows that using words as the basic unit is actually advantageous for text generation. And for generating longer texts, this advantage would likely be further amplified.

Of course, we don't deny that using WoBERT for sequence-labeling tasks like NER can result in a noticeable drop in performance—for instance, on the People's Daily NER dataset, it dropped by about 3%. What might be surprising is that, after analyzing the bad cases, we found the drop wasn't caused by segmentation errors, but rather by sparsity (on average, each word has fewer training samples, so training is less thorough).

In any case, we've open-sourced our work to give everyone one more option to try when using pretrained models.

Summary

In this article, we open-sourced a word-based Chinese BERT model (WoBERT), and discussed the pros and cons of using words as the basic unit. Finally, through experiments, we showed that word-based pretrained models have distinct value on quite a few NLP tasks (especially text generation)—on the one hand offering a speed advantage, and on the other hand matching the performance of character-based BERT. We welcome everyone to give it a try.