That Leaderboard-Topping T5 Model Can Now Be Used for Chinese

Training

T5’s pre-training consists of both unsupervised and supervised parts. The unsupervised part uses nearly 800GB of corpora constructed by Google (referred to as C4 in the paper). The training objective is similar to BERT but adapted into a Seq2Seq version, which we can view as an advanced version of a cloze test:

Input: Bright moon, when did you appear? [M0] asking the blue sky, I do not know [M1], what year is it tonight? I want to [M2] return, but fear the jade towers and crystal palaces, high places [M3]; dancing to [M4] clear shadows, how can it compare to the human world.
Output: [M0] raising my wine cup [M1] the heavenly palace [M2] ride the wind [M3] are too cold [M4] play with

For the supervised part, common NLP supervised task data was collected and uniformly converted into Seq2Seq tasks for training. For example, sentiment classification can be converted like this:

Input: Identify the sentiment of this sentence: I felt very good about this trip to Beijing.
Output: Positive

Topic classification can be converted like this:

Input: What kind of news is the following? After eight months, we can finally see the women’s volleyball girls on the court again.
Output: Sports

Reading comprehension can be converted like this:

Input: Reading Comprehension: Trump and Biden are running together for the next President of the United States. Answer the question based on the above information: What nationality is Trump?
Output: American

As we can see, this transformation is consistent with the ideas of GPT-2, GPT-3, and PET. They all aim to express the task we want to perform in text and then convert it into text prediction. Readers can also refer to the previous post "Is GPT-3 Necessary? No, BERT’s MLM Model Can Also Do Few-Shot Learning" for related content. Overall, in our internal experiments, the large model size, massive data, and supervised pre-training were all key factors in T5’s success. "Everything is Seq2Seq" provided an effective scheme to integrate these key factors.

Results

The main achievements of T5 are summarized in the table below:

In addition to dominating multiple leaderboards, T5 also tuned many adjustable hyperparameters throughout the training process, such as whether the standard Encoder-Decoder architecture is better than the UniLM-style structure, whether the BERT-style unsupervised pre-training task is better than others, whether a 15% random mask ratio is optimal, and so on. Finally, they provided the following table and expressed regret, saying, "Actually, we feel that T5’s experiments are not yet fully exhaustive," which gives a sense of "taking the road so others have no road to walk." Regardless, these "alchemy" results are worth a careful look for anyone working on language models, as they might help us avoid some detours.

T5.1.1

Overall, mT5 follows the same lineage as T5 and is basically the same. However, in terms of model structure, mT5 uses the T5.1.1 scheme. Here is a basic introduction to it.

What many people don’t know is that since its release in October last year, T5 underwent a low-key minor upgrade this year. Specific details can be found at the GitHub link. The official team calls the T5 before the upgrade T5.1.0, and the one after the upgrade T5.1.1. Its main changes come from the paper "GLU Variants Improve Transformer", which primarily borrows the GLU (Gated Linear Unit) from "Language Modeling with Gated Convolutional Networks" to enhance the performance of the FFN part. Specifically, the original T5 FFN was (T5 has no Bias): \text{FFN}(x)=\text{relu}(xW_1)W_2 It has now been changed to: \text{FFN}_{\text{GEGLU}}(x)=\big(\text{gelu}(xW_1)\otimes xW_2\big)W_3 That is, the first transformation layer with relu activation was changed to a gated linear unit with gelu activation. This increases the FFN layer parameters by 50%, but according to the paper, the performance increases significantly. Additionally, T5.1.1 made changes to the Embedding layer. In T5.1.0, the Encoder and Decoder Embedding layers and the Softmax layer for final probability distribution prediction in the Decoder all shared the same Embedding matrix. Now, T5.1.1 only shares the Embedding layer between the Encoder and Decoder, while the Softmax layer in the Decoder uses an independent Embedding matrix. Of course, this significantly increases the number of parameters, but Google’s conclusion is that this yields better results, as summarized in the recent paper "Rethinking embedding coupling in pre-trained language models". One final change: T5.1.1 removed Dropout during the pre-training phase and only uses Dropout during the downstream fine-tuning phase.

After these adjustments, Google retrained and released the full series of T5.1.1 models. The download addresses can be found at the GitHub link mentioned earlier. Note that T5.1.1 only underwent unsupervised pre-training, but its performance is still quite outstanding. Since T5.1.1 offers significant improvements, mT5 continued to use the T5.1.1 structure.

Results

mT5 essentially reconstructed a multilingual dataset mC4 and then trained it using the T5.1.1 scheme. There are no major innovations in the technical route. Regarding training details, readers can refer to the original paper; the paper is not long, as T5 had already paved the way.

As for mT5’s performance, it is mainly concentrated in the following table:

Readers might wonder how to evaluate such a multilingual model. Simply put, we could directly fine-tune a cross-lingual machine translation task on this basis to see the performance improvement. However, for multilingual models, researchers are more concerned with their Zero-Shot performance on cross-lingual tasks. In short, for the same task, if fine-tuned on one language, can the model be directly applied to other languages? This is the meaning of "Cross-lingual zero-shot transfer (models fine-tuned on English data only)" in the figure above. As can be seen, mT5’s performance is quite impressive.

Basics

The basic code to load the mT5 model into Keras using bert4keras is:

# Model paths
config_path = '/root/kg/bert/mt5/mt5_small/t5_config.json'
checkpoint_path = '/root/kg/bert/mt5/mt5_small/model.ckpt-1000000'
spm_path = '/root/kg/bert/mt5/sentencepiece.model'

# Load tokenizer
tokenizer = SpTokenizer(spm_path, token_start=None, token_end='</s>')

# Load model
t5 = build_transformer_model(
    config_path=config_path,
    checkpoint_path=checkpoint_path,
    model='t5.1.1',
    return_keras_model=False,
    name='T5',
)

encoder = t5.encoder
decoder = t5.decoder
model = t5.model

As you can see, there isn’t much difference from loading BERT in bert4keras. The construction of t5_config.json and the download of model.ckpt-1000000 are detailed on GitHub, so please refer to that. The complete code (training and decoding details) can also be found on GitHub, so I won’t expand on it here.

It is worth mentioning that for Chinese, the tokenizer output contains words; that is, for Chinese, mT5 is word-based, although the word granularity is relatively small. This further confirms that our previous work "Speeding Up Without Dropping Points: Chinese WoBERT Based on Word Granularity" was moving in the right direction.

Chinese

I believe most readers of this blog only care about Chinese tasks, some might care about English tasks, and very few would care about tasks other than Chinese and English. However, mT5 covers 101 languages with a total vocabulary of 250,000, and its T5.1.1 structure’s Softmax does not share parameters. This results in the Embedding layer occupying a massive number of parameters. For example, the mT5 small model has 300 million parameters, of which 250 million are related to Embedding. The key is that we don’t use most of these parameters; it’s purely unnecessary waste. Therefore, for those of us primarily concerned with Chinese tasks, it is necessary to streamline this Embedding layer.

Streamlining the model is simple: just delete the unnecessary rows in the two Embedding matrices. The key lies in how to decide which tokens to keep and how to obtain a streamlined SentencePiece model. To decide which tokens to keep, a simple thought is to keep Chinese tokens, but not just Chinese—some English ones must also be kept. It seems like a regular expression problem, but it’s actually not that simple. Using English letters doesn’t necessarily mean it’s English, and using Chinese characters doesn’t necessarily mean it’s Chinese; it’s a confusing matter. So I thought of another way: use the 250,000-token tokenizer to segment tens of gigabytes of Chinese corpora I collected, count the segmentation results, and then select the top part based on word frequency (eventually keeping over 30,000 tokens). Although this is more time-consuming, it is more reliable and ensures that the tokens we need most are preserved. After deciding on the vocabulary, a new SentencePiece model must be modified and obtained, which is also a bit troublesome, but I eventually solved it after some searching. The processing methods are shared on GitHub.

After this processing, to build a new model, you only need to add three lines of code related to keep_tokens. The required video memory is greatly reduced, and the Chinese generation effect remains basically unchanged:

# Model paths
config_path = '/root/kg/bert/mt5/mt5_base/t5_config.json'
checkpoint_path = '/root/kg/bert/mt5/mt5_base/model.ckpt-1000000'
spm_path = '/root/kg/bert/mt5/sentencepiece_cn.model'
keep_tokens_path = '/root/kg/bert/mt5/sentencepiece_cn_keep_tokens.json'

# Load tokenizer
tokenizer = SpTokenizer(spm_path, token_start=None, token_end='</s>')
keep_tokens = json.load(open(keep_tokens_path))

# Load model
t5 = build_transformer_model(
    config_path=config_path,
    checkpoint_path=checkpoint_path,
    keep_tokens=keep_tokens,
    model='t5.1.1',
    return_keras_model=False,
    name='T5',
)

encoder = t5.encoder
decoder = t5.decoder
model = t5.model

Performance

Finally, everyone is likely concerned about whether the generation effect has improved after all this effort and whether it is worth using. Let’s put it this way: the CSL title generation model fine-tuned with the mT5 small version can match the BLEU index of the UniLM model based on WoBERT, and the decoding speed is 130% faster. The CSL title generation model fine-tuned with the mT5 base version can exceed the WoBERT-based UniLM model by more than 1% in metrics, and the decoding speed is also 60% faster.

Simply put, it is indeed both fast and good. As for equipment requirements, students who usually run BERT base should basically be able to run the mT5 small/base versions, and even the large version can be attempted. As for XL and XXL, those are quite difficult to handle, and I suggest giving up on them. More surprises are for everyone to discover themselves   By the way, I should remind you that when fine-tuning the T5 model, the learning rate needs to be more than 10 times larger than when fine-tuning BERT (i.e., at the 10^{-4} level, whereas BERT is generally at the 10^{-5} level). This is determined by the difference in model architectures.