Speeding Up Without Performance Loss: Word-based Chinese WoBERT

Characters or Words?

Is "character" better or "word" better? This is a frustrating question in Chinese NLP, and several studies have investigated it systematically. A relatively recent one is "Is Word Segmentation Necessary for Deep Learning of Chinese Representations?" published by ShannonAI at ACL 2019, which concluded that characters are almost always superior to words. As mentioned earlier, current Chinese pre-trained models are indeed mostly character-based. So, does it seem like the problem is solved? Is character-based simply better?

The matter is far from that simple. Taking the ShannonAI paper as an example, its experimental results are not wrong, but they are not representative. Why? Because it compares the performance when the Embedding layers of both models are randomly initialized. In this case, for the same task, a word-based model has more parameters in the Embedding layer, making it naturally more prone to overfitting and resulting in worse performance—one could guess this without even running experiments. The issue is that when we use word-based models, we usually do not initialize them randomly; we often use pre-trained word vectors (choosing whether to fine-tune them based on the downstream task). This is the typical scenario for word-segmented NLP models, but the paper did not compare this scenario, so its results are not very convincing.

In fact, "overfitting" is a double-edged sword. We want to prevent overfitting, but overfitting also indicates that the model has a strong fitting capacity. If we find ways to suppress overfitting, we can obtain a stronger model with the same complexity, or a lower-complexity model with the same performance. One important means of alleviating overfitting is more thorough pre-training. Therefore, comparisons that do not introduce pre-training are unfair to word-based models. Our WoBERT confirms the viability of word-based pre-trained models.

Benefits of Words

Generally, the benefits of character-based models are:

1. Fewer parameters, less prone to overfitting;

2. No dependency on word segmentation algorithms, avoiding boundary segmentation errors;

3. Less severe sparsity, with virtually no Out-of-Vocabulary (OOV) words.

The reasons for using word-based models are:

1. Shorter sequences, leading to faster processing;

2. In text generation tasks, it can alleviate the Exposure Bias problem;

3. Lower semantic uncertainty of words, reducing modeling complexity.

Regarding the benefits of words, some might have doubts. For example, point 2: words can alleviate Exposure Bias because, theoretically, shorter sequences make the Exposure Bias problem less pronounced (a word model predicting an n-character word in one step is equivalent to a character model predicting n steps, where each step is recursively dependent, making the character model’s Exposure Bias more severe). As for point 3, although polysemous words exist, the meaning of most words is relatively fixed—at least more explicit than character meanings. Consequently, it might only take one Embedding layer to model word meaning well, unlike character models that require multiple layers to combine characters into words.

It seems like a toss-up, but in fact, the advantages of character-based models are not necessarily the disadvantages of word-based models. With a few techniques, word-based models can also avoid these issues to some extent:

1. While word-based models have more parameters, overfitting can be mitigated through pre-training, so this problem is not severe;

2. Dependency on segmentation algorithms is an issue, but if we only keep the most common words, the results from different segmentation tools are similar with little variation;

3. Boundary segmentation errors are hard to avoid, but accurate boundaries are only necessary for sequence labeling tasks. Text classification and text generation do not actually require precise boundaries, so word models should not be dismissed on this basis;

4. If we include most characters in the vocabulary, OOV words will not occur.

Therefore, the benefits of using words are quite numerous. Except for sequence labeling tasks that require very precise boundaries, most NLP tasks will not have issues with word-based units. Thus, we proceeded to create a word-based BERT model.

Tokenizer

To incorporate Chinese words into BERT, the Tokenizer must first be able to segment words. Is it enough to just add words to vocab.txt? Not quite. BERT’s built-in Tokenizer forcibly separates Chinese characters with spaces, so even if you add words to the dictionary, it won’t segment Chinese words. Furthermore, when BERT performs English word-piece tokenization, it uses the maximum matching method, which is not accurate enough for Chinese word segmentation.

To segment words, we modified BERT’s Tokenizer by adding a pre-tokenize operation. This allows us to segment Chinese words as follows:

1. Add Chinese words to vocab.txt;

2. For an input sentence s, use pre_tokenize to segment it first, obtaining [w_1, w_2, \dots, w_l];

3. Iterate through each w_i: if w_i is in the vocabulary, keep it; otherwise, segment w_i again using BERT’s built-in tokenize function;

4. Concatenate the tokenization results of each w_i in order as the final tokenization result.

In bert4keras >= 0.8.8, implementing this change only requires passing a parameter when constructing the Tokenizer:

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

Here, pre_tokenize is an external segmentation function; if not provided, it defaults to None. For simplicity, WoBERT uses Jieba segmentation. We removed the redundant parts of BERT’s original vocabulary (such as Chinese words with ##) and then added 20,000 additional Chinese words (the top 20,000 most frequent words from Jieba’s default dictionary). The final vocab.txt size for WoBERT is 33,586.

Model Details

The currently open-sourced WoBERT is the Base version, which underwent continued pre-training on the basis of the RoBERTa-wwm-ext open-sourced by HIT. The pre-training task was MLM. During the initialization phase, each word was split into characters using BERT’s built-in Tokenizer, and the average of the character embeddings was used to initialize the word embedding.

At this point, the technical essentials of WoBERT have been explained. The rest was the training process. We used a single 24G RTX card to train for 1 million steps (roughly 10 days), with a sequence length of 512, a learning rate of 5 \times 10^{-6}, a batch size of 16, and gradient accumulation for 16 steps, which is equivalent to training for about 60,000 steps with a batch size of 256. The training corpus was approximately 30GB of general-purpose data. The training code has been open-sourced at the link provided at the beginning of the article.

Additionally, we provide WoNEZHA, which is based on the NEZHA open-sourced by Huawei. The training details are basically the same as WoBERT. The model structure of NEZHA is similar to BERT, but it uses relative position encoding, whereas BERT uses absolute position encoding. Therefore, theoretically, the text length NEZHA can handle is unlimited. Providing the word-based WoNEZHA gives everyone an additional choice.

Model Performance

Finally, let’s discuss the performance of WoBERT. Simply put, in our evaluations, compared to BERT, WoBERT did not perform worse on NLP tasks that do not require precise boundaries; some even showed improvements. Meanwhile, the speed improved significantly. In short: "Speeding up without performance loss."

For example, comparing two classification tasks on Chinese benchmarks:

We also tested many internal tasks, and the results were similar, indicating that WoBERT and BERT are basically comparable on these NLU tasks. However, in terms of speed, WoBERT has a clear advantage. The table below compares the speed of the two models when processing text of different lengths:

We also tested Seq2Seq tasks (CSL/LCSTS headline generation) using WoBERT + UniLM. The results showed a significant improvement over character-based models:

This indicates that using words as units for text generation is actually more advantageous. If generating longer texts, this advantage would be further amplified.

Of course, we do not deny that when using WoBERT for sequence labeling tasks like NER, there might be a significant performance drop. For example, on People’s Daily NER, it dropped by about 3%. Perhaps surprisingly, after a bad case analysis, we found that the reason for the drop was not segmentation errors, but rather sparsity (on average, there are fewer samples per word, so training is less sufficient).

Regardless, we have open-sourced our work to provide everyone with an additional option to try when using pre-trained 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 units. Finally, through experiments, we demonstrated that word-based pre-trained models have unique value in many NLP tasks (especially text generation). On one hand, it offers a speed advantage; on the other hand, its performance is comparable to character-based BERT. We welcome everyone to test it.

Original Address: https://kexue.fm/archives/7758