RoFormerV2: Pushing the Limits of Natural Language Understanding

About a year ago, we proposed Rotary Position Embedding (RoPE) and released the corresponding pretrained model RoFormer. As time has passed, RoFormer has fortunately received increasing attention and recognition — for example, EleutherAI's newly released 6-billion and 20-billion parameter GPT models both adopt RoPE positional encoding, and Google's newly proposed FLASH model paper explicitly points out that RoPE brings a clear improvement to Transformer performance.

Meanwhile, we have kept working on strengthening the RoFormer model, trying to push its performance to "the next level." After nearly half a year of effort, we believe we've achieved fairly good results, and are officially releasing them as "RoFormerV2":

Github: https://github.com/ZhuiyiTechnology/roformer-v2

Pushing the Limits

Since the rise of pretrained models, many researchers have been quite interested in one question: where lies the limit of pretrained models? Of course, "limit" is a rich term — the line of work represented by GPT-3 tries to explore the limits of parameter count and data volume, while Microsoft's recently proposed DeepNet explores the limit of depth. For us, we are more interested in the performance limit under a fixed parameter budget, trying to squeeze as much performance as possible out of a pretrained model. RoFormerV2 is exactly the product of this philosophy.

Put simply, RoFormerV2 starts from RoFormer and makes appropriate simplifications to the model architecture, thereby gaining a certain speedup. On the training side, in addition to the usual unsupervised MLM pretraining, we also collected over 20GB of labeled data and carried out supervised multi-task pretraining. With this supervised training, the model's performance improved substantially, essentially achieving an optimal trade-off between speed and performance for a given parameter count.

Notably, the 300-million-parameter RoFormer large surpassed several models with 1-billion+ parameters on the CLUE leaderboard, ranking 5th overall — and it is also the model with the fewest parameters among the top 5:

RoFormerV2 large's RoFormerV2 large's "report card" on CLUE

Model Overview

Compared with RoFormer, the main changes in RoFormerV2 are: simplifying the model architecture, expanding the training data, and adding supervised training. Together, these changes allow RoFormerV2 to achieve a "win-win" in both speed and performance.

Architectural Simplification

Architecturally, RoFormerV2 mainly removes all bias terms from the model, replaces Layer Norm with a simpler RMS Norm, and removes the gamma parameter from RMS Norm. These changes were mainly inspired by Google's T5 model.

Intuitively, one might think that bias terms and the beta/gamma parameters of Layer Norm involve very little computation, and thus have negligible impact on speed. But the actual result surprised us: after removing these seemingly "negligible" parameters, RoFormerV2's training speed improved noticeably!

Some reference figures are as follows (RoFormer and RoBERTa have similar speeds, so we omit them here; the base version was tested on a 3090, and the large version on an A100):

$$\begin{array}{c|cc|cc} \hline & \text{sequence length} & \text{training speed} & \text{sequence length} & \text{training speed} \\ \hline \text{RoBERTa base} & 128 & 1.0\text{x} & 512 & 1.0\text{x} \\ \text{RoFormerV2 base} & 128 & 1.3\text{x} & 512 & 1.2\text{x}\\ \hline \text{RoBERTa large} & 128 & 1.0\text{x} & 512 & 1.0\text{x} \\ \text{RoFormerV2 large} & 128 & 1.3\text{x} & 512 & 1.2\text{x} \\ \hline \end{array}$$

Unsupervised Training

Like RoFormer, RoFormerV2 is first pretrained in an unsupervised manner via the MLM task. There are two main differences:

1. RoFormer was trained starting from RoBERTa's weights, while RoFormerV2 is trained from scratch;
2. RoFormer's unsupervised training used only about 30GB of data, while RoFormerV2 uses 280GB.

Training from scratch is harder than continuing training from existing weights, mainly because the Post-Norm architecture is more difficult to converge. To address this, we propose a new training technique: designing the residual as

\begin{equation}\boldsymbol{x}_{t+1} = \text{Norm}(\boldsymbol{x}_t + \alpha F(\boldsymbol{x}_t)) \end{equation}

where $\alpha$ is initialized at 0 and increases linearly and slowly to 1. For further discussion, see A Brief Discussion on Transformer Initialization, Parameterization, and Normalization. This scheme is similar to ReZero, except that in ReZero $\alpha$ is a trainable parameter and the $\text{Norm}$ operation is removed; experiments show that our modification achieves better final performance than ReZero, coming close to the best solution found before DeepNet.

Multi-task Training

As mentioned above, RoFormerV2's architecture is simplified to gain speed, but since there's "no free lunch," under the same training setup RoFormerV2 shows a slight performance drop compared to RoBERTa and RoFormer. To make up for this drop and further exploit the model's potential, we added supervised multi-task pretraining.

Specifically, we collected 77 labeled datasets totaling 20GB, and constructed 92 tasks for multi-task training. These datasets cover common natural language understanding tasks such as text classification, text matching, reading comprehension, information extraction, and coreference resolution, aiming to equip the model with fairly comprehensive natural language understanding capabilities. To carry out this training, we further developed a multi-task training framework on top of bert4keras that flexibly supports mixed training across tasks with different formats, and incorporated techniques such as gradient normalization (see Notes on Multi-task Learning (II): The Business of Gradient Alignment) to ensure each task achieves as good a result as possible.

RoFormerV2 is not the first model to attempt multi-task pretraining — prior work such as MT-DNN, T5, and, more recently, ZeroPrompt have all already confirmed the value of multi-task pretraining. Our main contribution is thoroughly validating this approach on Chinese and being the first to open-source it in this setting.

Experimental Results

We mainly compare results on the CLUE leaderboard:

$$\small{\begin{array}{c|ccccccccccc} \hline & \text{iflytek} & \text{tnews} & \text{afqmc} & \text{cmnli} & \text{ocnli} & \text{wsc} & \text{csl} & \text{cmrc2018} & \text{c3} & \text{chid} & \text{cluener}\\ \hline \text{BERT base} & 61.19 & 56.29 & 73.37 & 79.37 & 71.73 & 73.85 & 84.03 & 72.10 & 61.33 & 85.13 & 78.68\\ \hline \text{RoBERTa base} & 61.12 & 58.35 & 73.61 & 80.81 & 74.27 & 82.28 & \textbf{85.33} & 75.40 & 67.11 & 86.04 & 79.38\\ \text{RoBERTa large} & 60.58 & 55.51 & 75.14 & \textbf{82.16} & 75.47 & 81.97 & 85.07 & 78.85 & 76.74 & \textbf{88.65} & \textbf{80.19}\\ \hline \text{RoFormer base} & 61.08 & 56.74 & 73.82 & 80.97 & 73.10 & 80.57 & 84.93 & 73.50 & 66.29 & 86.30 & 79.69\\ \hline \text{RoFormerV2 small} & 60.46 & 51.46 & 72.39 & 76.93 & 67.70 & 69.11 & 83.00 & 71.80 & 64.49 & 77.35 & 78.20\\ \text{RoFormerV2 base} & 62.50 & \textbf{58.74} & 75.63 & 80.62 & 74.23 & 82.71 & 84.17 & 77.00 & 75.57 & 85.95 & 79.87\\ \text{RoFormerV2 large} & \textbf{62.65} & 58.06 & \textbf{76.95} & 81.20 & \textbf{75.83} & \textbf{88.03} & 84.97 & \textbf{80.50} & \textbf{78.34} & 87.68 & \textbf{80.17}\\ \hline \end{array}}$$

As we can see, the improvement from multi-task training is quite substantial. On most tasks, RoFormerV2 not only "recovers" the performance gap caused by architectural simplification, but also gains further improvement, achieving on average what could be considered the best result among models of the same tier. On CMNLI and CHID, however, RoFormerV2 underperforms RoBERTa — this is because both tasks have very large amounts of training data (hundreds of thousands of examples), and when the training data is abundant enough, model performance is mainly determined by model capacity, so the benefit from multi-task training becomes relatively small.

So, in summary: if your task is fairly standard and the data size isn't especially large, RoFormerV2 is often a good choice; if you also want a bit of a training speedup, RoFormerV2 is a good option too. But if your task has an especially large amount of data, RoFormerV2 typically won't have much of an advantage.

Summary

This post has given a basic introduction to our newly released RoFormerV2 model. It improves speed through architectural simplification, and improves performance by combining unsupervised and supervised pretraining, thereby achieving a "win-win" in both speed and performance.

English translation of a post from 科学空间 | Scientific Spaces by 苏剑林. Original: https://kexue.fm/archives/8998
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.