Lattice-BERT: Leveraging Multi-Granularity Representations in Chinese Pre-trained Language Models

BERT (Devlin et al., 2019, Bidirectional Encoder Representations from Transformers) is a pre-trained language model comprising a stack of multi-head self-attention layers and fully connected layers. For each head in the $l_{th}$ multi-head self-attention layer, the output matrix $\mathbf{H}^{\mathrm{out},l}=\left\{\mathbf{h}_{1}^{\mathrm{out},l},\mathbf{h}_{2}^{\mathrm{out},l},\dots,\mathbf{h}_{n}^{\mathrm{out},l}\right\}\in\mathbb{R}^{n\times d_{k}}$ satisfies:

$$\mathbf{h}_{i}^{\mathrm{out},l}=\sum_{j=1}^{n}{\left(\frac{\exp{\alpha_{ij}^{l}}}{\sum_{j^{\prime}}{\exp{\alpha_{ij^{\prime}}^{l}}}}\mathbf{h}_{j}^{\mathrm{in},l}\mathbf{W}^{v,l}\right)}$$

$$\alpha_{ij}^{l}=\frac{1}{\sqrt{2d_{k}}}\left(\mathbf{h}_{i}^{\mathrm{in},l}\mathbf{W}^{q,l}\right){\left(\mathbf{h}_{j}^{\mathrm{in},l}\mathbf{W}^{k,l}\right)}^{T}$$

(1)

where $\mathbf{H}^{\mathrm{in},l}=\left\{\mathbf{h}_{1}^{\mathrm{in},l},\mathbf{h}_{2}^{\mathrm{in},l},\dots,\mathbf{h}_{n}^{\mathrm{in},l}\right\}\in\mathbb{R}^{n\times d_{h}}$ is the input matrix, and $\mathbf{W}^{q,l},\mathbf{W}^{k,l},\mathbf{W}^{v,l}\in\mathbb{R}^{d_{h}\times d_{k}}$ are learnable parameters. $n$ and $d_{h}$ are sequence length and hidden size, and the attention size $d_{k}$ = $d_{h}/n_{h}$, where $n_{h}$ is the number of attention heads.

To capture the sequential features in languages, previous PLMs adopt position embedding in either input representations Devlin et al. (2019); Lan et al. (2020) or attention weights Yang et al. (2019); Wei et al. (2019); Ke et al. (2020). For the input-level position embedding, the inputs of the first layer are $\mathbf{\widetilde{h}}_{i}^{\mathrm{in},0}=\mathbf{{h}}_{i}^{\mathrm{in},0}+\mathbf{P}_{i}$, where $\mathbf{P}_{i}$ is the embedding of the $i_{th}$ position. The other works incorporate position information in attention weights, i.e., $\widetilde{\alpha}_{ij}^{l}={\alpha}_{ij}^{l}+f\left(i,j\right)$, where $f$ is a function of the position pair $\left(i,j\right)$.

The BERT model is pre-trained on an unlabeled corpus with reconstruction losses, i.e., Masked Language Modeling (MLM) and Next Sentence Prediction (NSP), and then fine-tuned on downstream tasks to solve specific NLU tasks. Readers could refer to Devlin et al. (2019) for details.

We adopt a word lattice to consume all possible segmentation results of a sentence in one PLM. Each segmentation can be a mixture of characters and words. As shown in Figure 1, a word lattice is a directed acyclic graph, where the nodes are positions in the original sentences, and each directed edge represents a character or a plausible word. Word lattices incorporate all words and characters so that models could explicitly exploit the inputs of both granularities, despite some of the words are redundant. In the rest of this work, we use lattice tokens to refer to text units, including the characters and words, contained in lattice graphs.

As shown in Figure 2, we list the lattice tokens in a line and consume these tokens to transformers straightforwardly. However, the challenges of learning PLMs as BERT over the lattice-like inputs include: 1) encoding the lattice tokens while preserving lattice structures; 2) avoiding potential leakage brought by redundant information.

Since the original BERT is designed for sequence modeling, it is not straightforward for BERT to consume a lattice graph. The word lattices encode not only the character sequences but also nested and overlapping words from different segmentations. To accurately incorporate positional information from lattice graphs into the interactions between tokens, we extend the attention-level position embedding and propose lattice position attention.

The lattice position attention aggregates the attention score of token representations, $\alpha_{ij}$ in Eq. 1, with three position related attention terms, encoding the absolute positions, the distance, and the positional relationship, which can be formulated as:

$$\widetilde{\alpha}_{ij}={\alpha}_{ij}+\mathrm{att}_{ij}+b_{ij}+r_{ij}$$

(2)

The $\mathrm{att}_{ij}$ in Eq. 2 is the attention weight between the absolute positions:

$$\mathrm{att}_{ij}=\frac{1}{\sqrt{2d_{k}}}\left(\left[\mathbf{P}^{S}_{s_{i}};\mathbf{P}^{E}_{e_{i}}\right]\mathbf{W}^{q}\right)\left(\left[\mathbf{P}^{S}_{s_{j}};\mathbf{P}^{E}_{e_{j}}\right]\mathbf{W}^{k}\right)^{T}$$

$\left[\cdot;\cdot\right]$ means the concatenation of vectors. $\mathbf{W}^{q},\mathbf{W}^{k}\in\mathbb{R}^{2d_{e}\times d_{k}}$ are learnable parameters, $d_{e}$ and $d_{k}$ are embedding size and attention size. $s_{i},e_{i}$ are positions of start and end characters of the $i_{th}$ token. Taking the word 研究/research in Figure 1 as an example, it starts at the first character and ends at the second one, thus, its $s_{i}$ and $e_{i}$ are 1 and 2, respectively. $\mathbf{P}^{S}$ and $\mathbf{P}^{E}$ are learnable position embedding matrices. $\mathbf{P}^{S}_{t},\mathbf{P}^{E}_{t}\in\mathbb{R}^{d_{e}}$ is the $t_{th}$ embedding vector of $\mathbf{P}^{S}$ or $\mathbf{P}^{E}$. The $\mathrm{att}_{ij}$ exploit the prior of attention weight between the start and end positions of the token pairs.

The $b_{ij}$ in Eq. 2 is the attention term for the distance between the $i_{th}$ and $j_{th}$ tokens, which consists of four scaling terms considering the combinations of the start and end positions:

$$b_{ij}={b}^{ss}_{{s_{j}}-{s_{i}}}+{b}^{se}_{{s_{j}}-{e_{i}}}+{b}^{es}_{{e_{j}}-{s_{i}}}+{b}^{ee}_{{e_{j}}-{e_{i}}}$$

$b^{ss}_{t}$ reflects the attention weight brought by the relative distance $t$ between the start positions of two tokens. The other terms, i.e., $b^{se}_{t}$, $b^{es}_{t}$, and $b^{ee}_{t}$, have similar meanings. In practice, the distance t is clipped into $\left[-128,128\right]$.

$r_{ij}$ in Eq. 2 is a scaling term represents the positional relation between the $i_{th}$ and $j_{th}$ tokens. We consider seven relations, including (1) self, (2) left and detached, (3) left and overlapped, (4) containing, (5) contained by, (6) right and overlapped, (7) right and detached. Figure 3 shows an illustration of these 7 relations. Formally, for the $i_{th}$ and $j_{th}$ tokens, they are overlapped means $s_{i}\leq s_{j}<e_{i}\leq e_{j}$ or $s_{j}\leq s_{i}<e_{j}\leq e_{i}$, and if $e_{i}<s_{j}$ or $e_{j}<s_{i}$, they are detached. If $s_{i}\leq s_{j}\leq e_{j}\leq e_{i}$ and $i\neq j$, the $i_{th}$ token contains the $j_{th}$ token and the $j_{th}$ token is contained by the $i_{th}$ token. Intuitively, only two detached tokens can be concurrent in one Chinese word segmentation result. Moreover, the containing relation reflects a sort of lexical hierarchy in the lattices. We think $r_{ij}$ can explicitly model the positional relations between tokens in lattice graphs.

We argue that the attention scores for distances and token relations capture different aspects of the multi-granularity structures in lattice graphs, thus, meeting the needs of various downstream tasks, such as distance for coreference resolution and positional relation for named entity recognition. With the information of absolute positions, distances, and positional relations, PLMs could accurately exploit the lattice structures in attention layers.

The lattice position attention weights are shared over all layers. $b_{ij}$, $r_{ij}$, $\mathbf{W}^{q}$, and $\mathbf{W}^{k}$ are diverse in different attention heads to capture diverse attention patterns. We follow Ke et al. (2020) to reset the positional attention scores related to [CLS] tokens, which is the special token as prefix of the input sequences to capture the overall semantics.

Vanilla BERT is trained to predict the randomly masked tokens in the sentences, i.e., the masked language modeling (MLM). For the case of consuming multi-granularity inputs, the input tokens are redundant which means a character can occur in its character forms and multiple words it belongs to. Directly adopting the randomly masking strategy may simplify the prediction problem in our case because the masked token can be easily guessed via peeking the unmasked tokens overlapping with the masked one. Taking the word 研究/research in Figure 2 as an example, supposing the masked input is [M]/究/研究, the model will consult 研究 rather than the context to predict the masked token, 研.

We investigate this problem and find that the tokens within a minimal segment of the lattice provide strong clue for the prediction of other tokens. A segment is a connected subgraph of a lattice where no token exists outside the subgraph that overlaps with any token inside the subgraph. To identify these minimal segments, we enumerate the character-level tokens in sentence order, checking if all the word-level tokens which contain this character end at this character. If so, all the tokens containing previous and current characters are considered as a segment, and the next segment starts from the next character, see the example in Figure 2. After the segment detection, we propose a masked segment prediction (MSP) task as a replacement of the MLM in the original BERT. In MSP, we mask out all the tokens in a segment and predict all these tokens (see Figure 2) to avoid the potential leakage.

In addition to MSP, we also pre-train our models with the sentence order prediction (SOP) task in Lan et al. (2020), where the model predicts whether two consecutive sentences are swapped in inputs.

We explore four kinds of downstream tasks, i.e., sentence/sentence-pair classification, multiple choices, sequence labeling, and span selection machine reading comprehension (MRC). For the sentence/sentence-pair classification, both vanilla and Lattice-BERT classify input instances base on logistic regressions over the representation of [CLS] tokens in the last layer. The circumstances are similar in multiple choice tasks, where softmax regressions are conducted over the representations of [CLS] tokens to choose the best options. However, for the span selection MRC, and the sequence labeling tasks like named entity recognition (NER), models need to perform token-wise classification. Vanilla BERT predicts labels for the input characters, but lattice-BERT has additional words. In Lattice-BERT, we extract the character chains (word pieces for numbers and English words) from lattices for training and prediction for a fair comparison with vanilla BERT. Pilot studies show that this strategy performs comparably with the more complex strategies, which supervise the labels over words and obtain a character’s label via ensembles of all tokens containing that character.

We test our models on 11 Chinese NLU tasks, including the text classification and Machine Reading Comprehension (MRC) tasks in the Chinese Language Understanding Evaluation benchmark (Xu et al., 2020, CLUE), and two additional tasks to probe the effectiveness in sequence labeling.

CLUE text classification: natural language inference CMNLI, long text classification IFLYTEK (IFLY.), short text classification TNEWS, semantic similarity AFQMC, coreference resolution (CoRE) CLUEWSC 2020 (WSC.), and key word recognition (KwRE) CSL.

CLUE MRC: Span selection based MRC CMRC 2018 (CMRC), multiple choice questions C³, and idiom cloze ChID.

Sequence Labeling: Chinese word segmentation (CWS) MSR dataset from SIGHAN2005 Emerson (2005), and named entity recognition (NER) MSRA-NER Levow (2006).

We probe our proposed Lattice-BERT model thoroughly with these various downstream tasks. The statistics and hyper-parameters of each task are elaborated in Appendix B. We tune learning rates on validation sets and report test results with the best developing performances for CLUE tasks.⁵⁵5We report accuracy / exact-match scores for CLUE tasks, and label-F1 / F1 for NER / CWS tasks. CLUE results see https://www.cluebenchmarks.com/rank.html. For MSR and MSRA-NER, we run the settings with the best learning rates five times and report the average scores to ensure the reliability of results.

RoBERTa Cui et al. (2020) is the Chinese version RoBERTa model Liu et al. (2019), which adopts the whole word masking trick and external pre-training corpus, known as RoBERTa-wwm-ext.⁶⁶6https://huggingface.co/hfl/chinese-roberta-wwm-ext
NEZHA Wei et al. (2019) is one of the best Chinese PLMs with a bag of tricks, which also explores attention-level position embedding.
AMBERT Zhang and Li (2020) is the state-of-the-art multi-granularity Chinese PLM, with two separated encoders for words and characters.
BERT-word is a Chinese PLM baseline, taking words as single-granularity inputs. We obtain the results from Zhang and Li (2020) directly.
BERT-our is our implemented BERT model, with the same pre-training corpus, model structures, hyper-parameters, and training procedure with Lattice-BERT, but taking characters as inputs. We also adopt the whole word masking trick.
LBERT is our proposed Lattice-BERT model, with word lattices as inputs, equipping with lattice position attentions and masked segment prediction.

In Table 1, we can see in text classification, MRC, and sequence labeling tasks, with both base and lite sizes, LBERT works better than our character-level baselines consistently. LBERT-base outperforms all previous base-size PLMs in average scores and obtain the best performances in 7 of the 11 tasks.

Comparing with the mono-granularity PLMs in base-size, LBERT takes benefits from word-level information and outperforms its character-level counterpart, BERT-our, by 1.5% averagely. Meanwhile, LBERT performs better than the word-level model, BERT-word, remarkably on CLUE tasks. We think the lattice inputs incorporate coarse-grained semantics while avoiding segmentation errors by combining multiple segmentation results. Therefore, with the multi-granularity treatments in word lattices, PLMs obtain better performances in downstream tasks than the mono-granularity settings.

Furthermore, LBERT outperforms the previous state-of-the-art (sota) multi-granularity PLM, AMBERT Zhang and Li (2020), by 0.9% in text classification and 1.3% in MRC, averagely. Different from modeling the characters and words separately, the graph representations of word lattices could enhance the interaction between multi-granularity tokens and utilize all possible segmentation results simultaneously. As a result, LBERT achieves a new sota among the base-size models on the CLUE leaderboard as well as the sub-leaderboards for text classification and MRC tasks.⁷⁷7https://www.cluebenchmarks.com/rank.html, the sota by the time of submission, Oct. 31st, 2020.

With lite-size settings, LBERT brings 2.0% improvement over BERT-our on average, which is larger than the case in base-size. In CWS, TNEWS, and CSL, the lite-size LBERT even outperforms the base-size BERT-our. With more coarse-grained inputs, the shallower architectures do not require complicated interactions to identify character combinations but utilizing word representations explicitly, thus, narrowing the gap with the deeper ones.