AMBERT: A Pre-trained Language Model with Multi-Grained Tokenization

Figure 1 gives an overview of AMBERT. AMBERT takes a text as input. Tokenization is conducted on the input text to obtain a sequence of fine-grained tokens and a sequence of coarse-grained tokens. AMBERT has two encoders, one for processing the fine-grained token sequence and the other for processing the coarse-grained token sequence. Each of the encoders has exactly the same architecture as that of BERT (Devlin et al., 2018). The two encoders share the same parameters at each corresponding layer, except that each has its own token embedding parameters. The fine-grained encoder generates contextualized representations from the sequence of fine-grained tokens through its layers. In parallel, the coarse-grained encoder generates contextualized representations from the sequence of coarse-grained tokens through its layers. AMBERT outputs a sequence of contextualized representations for the fine-grained tokens and a sequence of contextualized representations for the coarse-grained tokens.

AMBERT is expressive in that it learns and utilizes contextualized representations of the input text at both fine-grained and coarse-grained levels. The model retains all possibilities of tokenizations and learns the attention weights (importance) of representations of multi-grained tokens. AMBERT is also efficient through sharing of parameters between the two encoders. The parameters represent the same ways of combining representations, no matter whether representations are those of fine-grained tokens or coarse-grained tokens.

Pre-training of AMBERT is mainly conducted on the basis of masked language modeling (MLM), at both fine-grained and coarse-grained levels. Next sentence prediction (NSP) is not essential as indicated in many studies after BERT (Lan et al., 2019; Liu et al., 2019). We only use NSP in our experiments for comparison purposes. Let $\hat{\bf{x}}$ denote the sequence of fine-grained tokens with some of them being masked, and $\bar{\bf{x}}$ denote the masked fine-grained tokens. Let $\hat{\bf{z}}$ denote the sequence of coarse-grained tokens with some of them being masked, and $\bar{\bf{z}}$ denote the masked coarse-grained tokens. Pre-training is defined as optimization of the following function,

		$\displaystyle\min_{\theta}~{}-\log p_{\theta}(\bar{\mathbf{x}},\bar{\mathbf{z}}|\hat{\mathbf{x}},\hat{\mathbf{z}})\approx$
		$\displaystyle\min_{\theta}-\sum_{i=1}^{M}m_{i}\log p_{\theta}(x_{i}|\hat{\mathbf{x}})-\sum_{j=1}^{n}n_{j}\log p_{\theta}(z_{j}|\hat{\mathbf{z}}),$

where $m_{i}$ takes 1 or 0 as values and $m_{i}=1$ indicates that fine-grained token $x_{i}$ is masked, $m$ denotes the total number of fine-grained tokens; $n_{j}$ takes 1 or 0 as values and $n_{j}=1$ indicates that coarse-grained token $z_{j}$ is masked, $n$ denotes the total number of coarse-grained tokens; and $\theta$ denotes parameters.

In fine-tuning of AMBERT for classification, the fine-grained encoder and coarse-grained encoder create special [CLS] representations, and both representations are used for classification. Fine-tuning is defined as optimization of the following function, which is a regularized loss of multi-task learning, starting from the pre-trained model,

		$\displaystyle\min_{\theta}~{}-\log p_{\theta}({\bm{y}}|\mathbf{x})$
		$\displaystyle=\min_{\theta}-\log p_{\theta}({\bm{y}}|{\bm{r}}_{x0})-\log p_{\theta}({\bm{y}}|{\bm{r}}_{z0})$
		$\displaystyle-\log p_{\theta}({\bm{y}}|[{\bm{r}}_{x0},{\bm{r}}_{z0}])+\lambda\|\tilde{\bm{y}}_{x}-\tilde{\bm{y}}_{z}\|_{2},$

where $\mathbf{x}$ is the input text, ${\bm{y}}$ is the classification label, ${\bm{r}}_{x0}$ and ${\bm{r}}_{z0}$ are the [CLS] representations of fine-grained encoder and coarse-grained encoder, $[{\bm{a}},{\bm{b}}]$ denotes concatenation of vectors ${\bm{a}}$ and ${\bm{b}}$, $\lambda$ is a regularization coefficient, and $\|\|_{2}$ denotes L2 norm. The last term is based on agreement regularization (Brantley et al., 2019), which forces agreement between the predictions ($\tilde{\bm{y}}_{x}$ and $\tilde{\bm{y}}_{z}$).

Similarly, fine-tuning of AMBERT for span detection can be carried out, in which the representations of fine-grained tokens are concatenated with the representations of corresponding coarse-grained tokens. The concatenated representations are then utilized in the task.

We propose two ways of using AMBERT in inference. One is to utilize the AMBERT itself and the other to utilize only one encoder of AMBERT. The former performs better but needs more computation and the latter performs slightly worse but only needs computation comparable to BERT. One can choose either drop the fine-grained encoder or the coarse-grained encoder in AMBERT through evaluation using a development dataset, which makes the computational cost close to that of BERT.

We can consider two alternatives to AMBERT, which also rely on multi-grained tokenization. We refer to them as AMBERT-Combo and AMBERT-Hybrid and make comparisons of them with AMBERT in our experiments.

AMBERT-Combo has two individual encoders, an encoder (BERT) working on the fine-grained token sequence and the other encoder (BERT) working on the coarse-grained token sequence, without parameter sharing between them. In learning and inference AMBERT-Combo simply combines the output layers of the two encoders. Its fine-tuning is similar to that of AMBERT.

AMBERT-Hybrid has only one encoder (BERT) working on both the fine-grained token sequence and the coarse-grained token sequence. It creates representations on the concatenation of two sequences and lets the representations of the two sequences interact with each other at each layer. Its pre-training is formalized in the following function,

		$\displaystyle\min_{\theta}~{}-\log p_{\theta}(\bar{\mathbf{x}},\bar{\mathbf{z}}|\hat{\mathbf{x}},\hat{\mathbf{z}})$
		$\displaystyle\approx\min_{\theta}-\sum_{i=1}^{m}m_{i}\log p_{\theta}(x_{i}|\hat{\mathbf{x}},\hat{\mathbf{z}})$
		$\displaystyle-\sum_{j=1}^{n}n_{j}\log p_{\theta}(z_{j}|\hat{\mathbf{x}},\hat{\mathbf{z}}),$

where the notations are the same as in (LABEL:equ:mlm-ambert). Its fine-tuning is the same as that of BERT.

For Chinese, we use a corpus consisting of 25 million documents (57G uncompressed text) from Jinri Toutiao¹¹1Jinri Toutiao is a popular news app. in China.. Note that there is no common corpus for training of Chinese BERT. For English, we use a corpus of 13.9 million documents (47G uncompressed text) from Wikipedia and OpenWebText (Gokaslan and Cohen, 2019)²²2Unfortunately, BookCorpus, one of the two corpora in the original paper for English BERT, is no longer publicly available..

The characters in the Chinese texts are naturally taken as fine-grained tokens. We conduct word segmentation on the texts and treat the words as coarse-grained tokens. We employ a word segmentation tool based on a n-gram model. Both tokenizations exploit WordPiece embeddings (Wu et al., 2016). There are 21,128 characters and 72,635 words in the vocabulary of Chinese.

The words in the English texts are naturally taken as fine-grained tokens. We perform coarse-grained tokenization on the English texts in the following way. First, we calculate the n-grams in the Wikipedia documents using KenLM (Heafield, 2011). We next build a phrase-level dictionary consisting of phrases whose frequencies are sufficiently high and whose last words highly depend on their previous words. We then employ a left-to-right search algorithm to perform phrase-level tokenization on the texts. There are 30,522 words and 77,645 phrases in the vocabulary of English.

We make use of the same parameter settings for the AMBERT and BERT models. All models in this paper are ‘base-models’ having 12 layers of encoder. It is too computationally expensive for us to train the models as ‘large models’ having 24 layers. To retain consistency, the masked spans in the coarse-grained encoder are also masked in the fine-grained encoder. The details of pre-training and fine-tuning are the same as those in the original BERT paper (Devlin et al., 2018), which are given in Appendix C.

We also conduct experiments on the efficient inference method of AMBERT on CLUE/GLUE/SQuAD/RACE. We choose the fine-grained encoder for the span detection tasks (CMRC2018 and SQuAD) because it performs much better in the tasks. We choose the coarse-grained encoder for the other Chinese tasks and the fine-grained encoder for the other English tasks because they perform better on average. All the decisions are made based on the results from the Dev datasets. The detailed results are shown in Table 7. We conclude that, a) for the English tasks, AMBERT with one chosen encoder achieves similar results as AMBERT with two encoders and outperforms the single-grained “Our BERT” models with a large margin; b) for the Chinese tasks, AMBERT with one chosen encoder performs slightly worse than AMBERT but performs much better than the single-grained “Our BERT” models. Therefore, in practice, one can train an AMBERT with two encoders and use only one of the encoders in inference.

We further investigate the reason that AMBERT is superior to AMBERT-Combo. Figure 3 shows the distances between the [CLS] representations of the fine-grained encoder and coarse-grained encoder in AMBERT-Combo and AMBERT after pre-training, in terms of cosine distance (one minus cosine similarity) and normalized Euclidean distance. One can see that the distances in AMBERT-Combo are larger than the distances in AMBERT. We perform the assessment using the data in different tasks and find similar trends. The results indicate that the representations of fine-grained encoder and coarse-grained encoder are closer in AMBERT than in AMBERT-Combo. These are natural consequences of using AMBERT and AMBERT-Combo, whose parameters are respectively shared and unshared across encoders. It implies that the higher performances by AMBERT is due to its parameter sharing, which can learn and represent similar ways of combining tokens no matter whether they are fine-grained or coarse-grained. An intuitive explanation is that the ways of combining representations of fine-grained tokens and the ways of combining representations of coarse-grained tokens “in the same contexts” are exactly the same.

We also examine the reasons that AMBERT works better than AMBERT-Hybrid, while both of them exploit multi-grained tokenization. Figure 2 shows the attention weights of first layers in AMBERT and AMBERT-Hybrid, as well as the single-grained BERT models, after pre-training. In AMBERT-Hybrid, the fine-grained tokens attend more to the corresponding coarse-grained tokens and as a result the attention weights among fine-grained tokens are weakened. In contrast, in AMBERT the attention weights among fine-grained tokens and those among coarse-grained tokens are intact. It appears that attentions among single-grained tokens (fine-grained ones or coarse-grained ones) play important roles in downstream tasks.

To answer the question why the improvements by AMBERT on Chinese are larger than on English in the same pre-training settings, we further make an analysis. We respectively tokenize 10,000 randomly selected Chinese sentences from five tasks in CLUE with our Chinese word tokenizer. The average proportion of words is 51.5%, which indicates that about half of the tokens are fine-grained and half are coarse-grained in Chinese. Similarly, we tokenize 10,000 randomly selected English sentences from five different tasks in GLUE with our English phrase tokenizer. The average proportion of phrases is only 13.1%, which means that there are much less coarse-grained tokens than fine-grained tokens in English. (Please refer to Table 10 in the Appendix for more details of the experiments.) Therefore, we postulate that for Chinese it is necessary for a model to process the language at both fine-grained and coarse-grained levels. AMBERT indeed has the capability.