SlowGuess/NAACL_2022_OCR huggingface dataset:
Yu Lin*, Zhecheng An*, Peihao Wu, Zejun Ma
AI-Lab Speech & Audio Team, ByteDance Inc., Beijing, China
{linyu.linyu,anzhecheng,wupeihao,mazejun}@bytedance.com
Though achieving impressive results on many NLP tasks, the BERT-like masked language models (MLM) encounter the discrepancy between pre-training and inference. In light of this gap, we investigate the contextual representation of pre-training and inference from the perspective of word probability distribution. We discover that BERT risks neglecting the contextual word similarity in pre-training. To tackle this issue, we propose an auxiliary gloss regularizer module to BERT pre-training (GR-BERT), to enhance word semantic similarity. By predicting masked words and aligning contextual embeddings to corresponding glosses simultaneously, the word similarity can be explicitly modeled. We design two architectures for GR-BERT and evaluate our model in downstream tasks. Experimental results show that the gloss regularizer benefits BERT in word-level and sentence-level semantic representation. The GR-BERT achieves new state-of-the-art in lexical substitution task and greatly promotes BERT sentence representation in both unsupervised and supervised STS tasks.
Pre-trained language models like BERT (Devlin et al., 2019) and its variants (Liu et al., 2019b; Lan et al., 2019; Zhang et al., 2019; Joshi et al., 2020) have achieved remarkable success in a wide range of natural language processing (NLP) benchmarks. By pre-training on large scale unlabeled corpora, BERT-like models learn contextual representations with both syntactic and semantic properties. Researches show the contextual representations generated by BERT capture various linguistic knowledge, including part-of-speech (PoS), named entities, semantic roles (Tenney et al., 2019; Liu et al., 2019a; Ettinger, 2020), word senses (Wiedemann et al., 2019), etc. Furthermore, with the fine-tuning procedure, the contextual representations show excel
(c)
Figure 1: Conditional token probability distribution of tokens given masked context (a) and full context (b) and (c). The ideal token distribution given full context is illustrated in (b), while (c) shows the full contextual token distribution generated by actual BERT.
lent transferability in downstream language understanding tasks, and lead to state-of-the-art (SOTA) performance.
The masked language model (MLM) plays a significant role in the pre-training stage of many BERT-like models (Liu et al., 2019b). In an MLM, a token  $w$  is sampled from a text sequence  $\mathbf{s}$ , and replaced with a [MASK] token. Let  $\mathbf{c}$  be the rest of tokens in  $\mathbf{s}$  except for  $w$ . We name  $\mathbf{c}$  as the masked context or surrounding context, and  $\mathbf{s}$  as the full context. During pre-training, BERT encodes the masked context  $\mathbf{c}$  into a contextual embedding vector  $h_c$ , and use it to generate a contextual token probability distribution  $p(x|\mathbf{c})$ , where  $x \in V$  and  $V$  denotes the token vocabulary. The training objective is to predict the masked token  $w$  by maximizing likelihood function  $\log p(w|\mathbf{c})$ . In the fine-tuning or inference stage, BERT takes the full context  $\mathbf{s}$  without masks as input, and encodes every token into its contextual representation for downstream tasks. We denote the contextual representation corresponds to token  $w$  as  $h_s$ .
We analyze the corresponding contextual token probability distribution  $p(x|\mathbf{c})$  and  $p(x|\mathbf{s})$  generated from  $h_c$  and  $h_s$ , as a proxy to study the representations (Li et al., 2020). Figure 1(a) shows an example when masked context  $\mathbf{c} = \text{"Tom is a [MASK] guy" }$ , the predicted tokens with high prob
*The first two authors contributed equally to this work.
Findings of the Association for Computational Linguistics: NAACL 2022, pages 907 - 920
abilities  $p(x|\mathbf{c})$  includes good, nice, great, tough, which are all reasonable answers to the Cloze task. Ideally, we want the context encoder to behave the same way when full context  $\mathbf{s}$  is given, as in Figure 1(b), the model should only propose contextual synonyms of bad such as dangerous, nasty and mean with  $p(x|\mathbf{s})$ . However, the actual BERT generates  $\hat{p}(x|\mathbf{s})$  as shown in Figure 1(c), which contains inappropriate token proposals such as good, rough and big.
The discrepancy between Figure 1(b) and 1(c) is because only the masked token distribution  $p(x|\mathbf{c})$  is explicitly modeled in BERT with the MLM, while the full contextual token distribution  $p(x|\mathbf{s})$  works in an agnostic way through model generalization. This leads to a gap between  $p(x|\mathbf{c})$  in pretraining and  $p(x|\mathbf{s})$  in fine-tuning and inference. It is shown in unsupervised semantic textual similarity (STS) tasks, BERT generates contextual embeddings that even underperforms static embeddings for sentence representation (Reimers and Gurevych, 2019). Although in BERT pre-training, random token replacement strategy is used to mitigate the mismatch that [MASK] token is never seen during fine-tuning, to the best of the authors' knowledge, there is no analysis on the gap of representation between masked context  $h_c$  and full context  $h_s$  in different phases when using BERT.
To address this issue, we perform an investigation on the inner structure of  $p(x|\mathbf{s})$ . Through theoretical derivation, we discover  $p(x|\mathbf{s})$  can be decomposed into the combination of masked contextual token distribution  $p(x|\mathbf{c})$  and a point-wise mutual information (PMI) term that describes contextual token similarity. Further analysis shows both the MLM and token replacement in BERT pre-training have potential shortcomings in modeling the contextual token similarity. Inspired by the decomposition of  $p(x|\mathbf{s})$ , we propose to add an auxiliary gloss regularizer (GR) module to the MLM task, where mask prediction and gloss matching are trained simultaneously in the BERT pre-training. We also design two model architectures to integrate the gloss regularizer into the original MLM task.
We examine our proposed model in downstream tasks including unsupervised lexical substitution (LS) (McCarthy and Navigli, 2007; Kremer et al., 2014), unsupervised STS and supervised STS Benchmark (Cer et al., 2017). By invoking gloss regularized pre-training, our model improves lexical substitution task from 14.5 to 15.2 points in the
LS14 dataset, leading to new SOTA performance. In unsupervised STS tasks, gloss regularizer improves the performance from 56.57 to 67.47 in terms of average Spearman correlation by a large margin. Such performance gain is also observed in supervised STS task. Empirical experiments prove our model effectively generates better contextual token distribution and representations, which contributes to word-level and sentence-level language understanding tasks.
Masked Language Models. Liu et al. (2019b) extend BERT into RoBERTa achieving substantial improvements. They claim the MLM task as the key contributor to contextual representation modeling, compared with next sentence prediction task. Many BERT variants focus on better masking strategies (Cui et al., 2019; Zhang et al., 2019; Joshi et al., 2020) to enhance the robustness and transferability of contextual representative learning. However, MLM suffers from the discrepancy between pre-training and fine-tuning since the [MASK] tokens are only introduced during pre-training. To tackle this issue, permutation language model from XLNet (Yang et al., 2019) and token replacement detection from ELECTRA (Clark et al., 2020) are proposed as alternative approaches to the MLM. Instead of avoiding MLM, we analyze how the mask modeling affects the full contextual representation in a probability perspective, and introduce gloss regularizer to mitigate the gap brought by MLM.
Contextual Representation Analysis. One way to analyze the contextual representation learned by pre-trained language model is through the probing tasks (Liu et al., 2019a; Miaschi and Dell'Orletta, 2020; Vulic et al., 2020), which are regarded as empirical proofs that pre-trained MLMs like BERT succeed in capturing linguistic knowledge. Many other researches focus on studying the geometry of contextual representations. Ethayarajh (2019) discovers anisotropy among the contextual embeddings of words when studying contextuality of BERT. Li et al. (2020) propose a method using normalizing flow to transform the contextual embedding distribution of BERT into an isotropic distribution, and achieve performance gains in sentence-level tasks.
Utilizing Word Senses. Because the BERT conveys contextualized semantic knowledge of polyse
mous, many researches use BERT as a backbone to build word sense disambiguation (WSD) models (Huang et al., 2019; Blevins and Zettlemoyer, 2020; Bevilacqua and Navigli, 2020). In these models, BERT is used as word senses and contexts encoders to perform the downstream matching task. One work that directly incorporates word sense knowledge into pre-training is SenseBERT (Levine et al., 2020) that introduces a weakly-supervised supersense prediction task, which leads to improvement on performance of WSD and word-in-context task. In SenseBERT, word prediction is enhanced with supersense category labels that act like an external knowledge source. However, the gloss regularizer in our model provides fine-grained semantic information, which aimed to align word representation space with the semantic space, and leads to better contextual representations.
3 Contextual Token Probability
3.1 Masked Language Model
Without loss of generality, the token probability distribution given full context  $p(x|\mathbf{s})$  can be decomposed into two parts,
$$
\begin{array}{l} \log p (x | \mathbf {s}) = \log p (x | \mathbf {c}) + \log \frac {p (x | w , \mathbf {c})}{p (x | \mathbf {c})} \\ = \log p (x | \mathbf {c}) + \operatorname {P M I} (x; w | \mathbf {c}) \quad (1) \\ \end{array}
$$
where  $\mathrm{PMI}(x;w|\mathbf{c})$  is the pointwise mutual information between  $x$  and  $w$  given  $\mathbf{c}$ . PMI describes how frequently two tokens co-occur than their independent occurrences, which is used as a measurement of the semantic similarity between tokens (Ethayarajh, 2019; Li et al., 2020). In Eqn. (1),  $\log p(x|\mathbf{c})$  only depends on masked context, which directly corresponds to the MLM training objective. However, the PMI term is not explicitly modeled.
In BERT,  $p(x|\mathbf{c})$  is generated from the encoded mask context  $h_c$  with a softmax operation as
$$
p (x | \mathbf {c}) = \operatorname {s o f t m a x} \left(\boldsymbol {h} _ {c} ^ {\top} \boldsymbol {v} _ {x}\right), \tag {2}
$$
where  $\pmb{v}_x$  stands for the embedding vector of token  $x$  in vocabulary  $V$ . During fine-tuning or inference stage, full context  $\mathbf{s}$  without masks is encoded into  $\pmb{h}_s$  as the contextual representation of token  $w$ . We can use the  $\pmb{h}_s$  to estimate  $p(x|\mathbf{s})$  in the same way as Eqn. (2), denoted by  $\hat{p} (x|\mathbf{s})$ ,
$$
\hat {p} (x | \mathbf {s}) = \hat {p} (x | w, \mathbf {c}) = \operatorname {s o f t m a x} \left(\boldsymbol {h} _ {s} ^ {\top} \boldsymbol {v} _ {x}\right). \tag {3}
$$
Under such approximation setup,  $\mathrm{PMI}(x;w|\mathbf{c})$  can be transformed into
$$
\begin{array}{l} \operatorname {P M I} (x; w | \mathbf {c}) \approx \log \frac {\hat {p} (x | w , \mathbf {c})}{p (x | \mathbf {c})} \\ = \left(\boldsymbol {h} _ {s} - \boldsymbol {h} _ {c}\right) ^ {\top} \boldsymbol {v} _ {x} + \varphi (\boldsymbol {w}, \mathbf {c}), \quad (4) \\ \end{array}
$$
where  $\varphi (w,\mathbf{c})$  is constant w.r.t.  $x$  (detailed in Appendix A). In a deep neural network parameterized model like BERT,  $h_s$  is encoded in an agnostic way. Thus, it is difficulty to further derive the PMI in Eqn. (4).
For a simpler case, if we consider a one-layer continuous bag-of-words (CBOW) model (Mikolov et al., 2013)  $^1$ , we have  $\pmb{h}_s - \pmb{h}_c = \pmb{h}_w$ , where  $\pmb{h}_{\mathcal{W}}$  is a context vector only related to the center token  $w$ . Now the PMI is formulated as
$$
\operatorname {P M I} _ {\mathrm {C B O W}} (x; w | \mathbf {c}) = \log p (x | w) + \psi (w, \mathbf {c}), \tag {5}
$$
where  $\psi (w,\mathbf{c})$  is another constant w.r.t.  $x$  (also detailed in Appendix A). In this case, only the similarity information between  $x$  and  $w$  plays a role when comparing  $\mathrm{PMICBOW}(x;w|\mathbf{c})$  among different candidate tokens  $x\in V$  , while the context information is irrelevant.
Although  $h_s - h_c = h_w$  is not satisfied in a deep model like BERT, the input sequences for  $h_s$  and  $h_c$  share the most identical tokens  $\mathbf{c}$ , and their only difference is whether to mask  $w$ . Therefore, there is a potential risk that  $\mathrm{PMI}(x; w|\mathbf{c})$  in MLM loses information related to the condition  $\mathbf{c}$ , and degrades to the marginal  $\mathrm{PMI}(x; w)$ , especially when the MLM lacks modeling  $p(x|\mathbf{s})$  in its training objective.
3.2 Replaced Language Model
In the BERT training process, a portion of tokens are replaced with random real tokens other than [MASK], and the model is trained to predict the original tokens. We name this task as the replaced language model (RLM). Different from MLM, an RLM takes full context without masked tokens as input, and directly generates token distribution  $p(x|\mathbf{s})$ , which seems to be a better way for full contextual representation modeling.
We take a closer look at the RLM training process. Let  $p(x|\mathbf{s}) = p(x|w,\mathbf{c})$  be the probability that token  $w$  is replaced with token  $x$  in context  $\mathbf{c}$ . According to the Bayes' theorem, we have
$$
p (x \mid w, \mathbf {c}) = \frac {p (x \mid \mathbf {c}) p (w \mid x , \mathbf {c})}{\sum_ {x ^ {\prime} \in V} p \left(x ^ {\prime} \mid \mathbf {c}\right) p \left(w \mid x ^ {\prime} , \mathbf {c}\right)}. \tag {6}
$$
1The CBOw model can be considered as a kind of masked language model.
In a well-trained model,  $p(w|x,\mathbf{c})$  should be the replacing probability during training. Since the process of randomly replacing words is irrelevant to the context,  $p(w|x,\mathbf{c}) = p(w|x)$ . Let  $\alpha$  be the probability when a token remains unchanged, and  $1 - \alpha$  be the replacing probability. Therefore,
$$
p (x | \mathbf {s}) = \frac {(1 - \alpha) p (x | \mathbf {c})}{\alpha | V | p (w | \mathbf {c}) + (1 - \alpha) \sum_ {x ^ {\prime} \neq w} p \left(x ^ {\prime} | \mathbf {c}\right)}, \tag {7}
$$
where  $|V|$  denotes the vocabulary size. Eqn. (7) shows in RLM  $p(x|\mathbf{s})$  is proportional to  $p(x|\mathbf{c})$ . Therefore,  $\mathrm{PMI}(x;w|\mathbf{c})$  is constant w.r.t.  $x$ , i.e.
$$
\begin{array}{l} \operatorname {P M I} (x; w | \mathbf {c}) = \log \frac {p (x | w , \mathbf {c})}{p (x | \mathbf {c})} = \log \frac {p (x | \mathbf {s})}{p (x | \mathbf {c})} \\ = \log \frac {1 - \alpha}{\alpha | V | p (w | \mathbf {c}) + (1 - \alpha) \sum_ {x ^ {\prime} \neq w} p \left(x ^ {\prime} \mid \mathbf {c}\right)}. \tag {8} \\ \end{array}
$$
Combining Eqn. (8) with Eqn. (1), we conclude that the distribution of  $x$  only relies on surrounding context  $\mathbf{c}$ , but pays no attention to the center token  $w$ . This infers the RLM actually models the token distribution conditioning on almost only the surrounding context, even if it takes full context as input. As a result, the RLM fails to contribute better full contextual representation performance to the MLM, since the PMI term, as a component part of the full contextual token distribution  $p(x|\mathbf{s})$ , is completely ignored in the RLM. Solely using RLM would lead to worse contextual representation than using only MLM for pre-training.
4 Gloss Regularizer
4.1 Invoking Gloss Matching
As shown in Eqn. (1),  $p(x|\mathbf{s})$  consists of  $p(x|\mathbf{c})$  and  $\mathrm{PMI}(x;w|\mathbf{c})$ . Both MLM and RLM succeed in modeling  $p(x|\mathbf{c})$ . However, the analysis in Section 3 shows RLM completely ignores  $\mathrm{PMI}(x;w|\mathbf{c})$ , and MLM may suffer from potential risks that the contextual information in  $\mathrm{PMI}(x;w|\mathbf{c})$  would be lost, in either way the model generates poor estimation of  $p(x|\mathbf{s})$ .
$\mathrm{PMI}(x;w|\mathbf{c})$  describes co-occurrence probability of  $x$  and  $w$  normalized by their marginal probabilities under context  $\mathbf{c}$  as condition. Ideally, it should be learned by training with labeled dataset  $\{(\mathbf{s}_1,\mathbf{s}_2)\}$ , where  $\mathbf{s}_1 = \{x_1,\mathbf{c}\}$  and  $\mathbf{s}_2 = \{x_2,\mathbf{c}\}$  are semantically similar text samples with shared context  $\mathbf{c}$  and exchangeable token pair  $(x_{1},x_{2})$ .
However, such labeled data is expansive to build and not suitable for large-scale pre-training setup.
Intuitively,  $\mathrm{PMI}(x;w|\mathbf{c})$  can be regarded as semantic similarity between tokens under context. Although the contexts of similar tokens are hard to obtain, we can use the glosses of tokens as an alternative. Since the semantic of a word can be defined by its gloss, contextual token similarity can be determined by detecting whether tokens are matching to similar glosses under context. Therefore, in order to better model the contextual token similarity defined by  $\mathrm{PMI}(x;w|\mathbf{c})$ , we introduce gloss matching an auxiliary task named the gloss regularizer. Two architectures to integrate gloss regularizer into MLM are detailed in Section 4.2 and 4.3.
4.2 Multitask Model
A straight-forward method is to perform mask prediction and gloss matching as joint multitasks (denoted as MT). In this architecture, the masked context  $\mathbf{c}$  and the full context  $\mathbf{s}$  are encoded by a context encoder into the contextual vector  $h_c$  and  $h_s$ . The loss function of the MLM task is
$$
\mathcal {L} _ {\mathrm {M L M}} = - \boldsymbol {h} _ {c} ^ {\top} \boldsymbol {v} _ {w} + \log \sum_ {w ^ {\prime} \in V} \exp \left(\boldsymbol {h} _ {c} ^ {\top} \boldsymbol {v} _ {w ^ {\prime}}\right). \tag {9}
$$
For the gloss matching task, as illustrated in Figure 2(a), let  $\mathbf{g}_t$  be the gloss text of token  $w$  under context  $\mathbf{c}$ . Another gloss encoder is used to encode  $\mathbf{g}_t$  into a gloss vector  $\boldsymbol{e}_t$ . Gloss matching is performed by calculating the similarity between the contextual token representation  $\boldsymbol{h}_s$  and the gloss vector  $\boldsymbol{e}_t$ . The gloss regularizing loss is
$$
\mathcal {L} _ {\mathrm {G R}} = - \operatorname {s i m} [ \boldsymbol {h} _ {s}, \boldsymbol {e} _ {t} ] + \log \sum_ {t ^ {\prime} \in T} \exp \sin [ \boldsymbol {h} _ {s}, \boldsymbol {e} _ {t ^ {\prime}} ], \tag {10}
$$
where  $\mathrm{sim}[\cdot ,\cdot ]$  is a similarity measurement function, and  $T$  is a set of negative glosses. The final loss function is the combination of the two losses,
$$
\mathcal {L} _ {\mathrm {M T}} = \mathcal {L} _ {\mathrm {M L M}} + \lambda \mathcal {L} _ {\mathrm {G R}}, \tag {11}
$$
where  $\lambda$  denotes the regularizing weight.
This setting resembles the bi-encoders model (BEM) for WSD proposed by (Blevins and Zettlemoyer, 2020). However, in our model, the context encoder is trained on mask prediction task simultaneously with the gloss matching task, while the BEM takes gloss matching as a fine-tuning task. We train the two tasks together for better contextual and semantic representation modeling. As a
(a) The framework of the gloss regularized BERT
Type I: Multi-Task (MT) Model
(b) Two types of the context encoder structures
Type II: Separate Context Encoder (SC) Model
Figure 2: (a) shows the gloss regularizer aligns contextual representation space with the gloss space. (b) Two GR architectures: the MT trains MLM and GR as multitask, while the SC utilizes two independent context encoders (the loss  $\mathcal{L}_{\mathrm{MLM}}^{(2)}$  of SC is not shown).
result, the model learns token distribution not only conditioning on the masked context, but also influenced by semantic similarity with center token, which gives a better estimation of  $p(x|\mathbf{s})$ .
4.3 Separate Context Encoder Model
Another method is directly inspired by the decomposition from Eqn. (1). Different from the multi-task model, we use two context encoders instead of one (denoted as SC). The first context encoder, denoted by  $\mathsf{enc}_1$ , encodes the masked context as  $h_c^{(1)} = \mathsf{enc}_1(\mathbf{c})$ , and learns purely from MLM task with loss  $\mathcal{L}_{\mathrm{MLM}}^{(1)}$  derived similar as Eqn. (9).
The full context  $\mathbf{s}$  is encoded into  $h_s^{(2)} = \mathrm{enc}_2(\mathbf{s})$  by the second context encoder. Eqn. (4) shows  $\mathrm{PMI}(x;w|\mathbf{c})$  is entailed in the linear difference between the encoding of full and masked context. Therefore, we use  $(h_s^{(2)} - h_c^{(1)})$  for gloss matching, where the loss function is formulated as
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {G R}} ^ {\mathrm {s}} = - \operatorname {s i m} \left[ e _ {t}, h _ {s} ^ {(2)} - h _ {c} ^ {(1)} \right] \\ + \log \sum_ {t ^ {\prime} \in T} \exp \sin \left[ e _ {t ^ {\prime}}, h _ {s} ^ {(2)} - h _ {c} ^ {(1)} \right]. \tag {12} \\ \end{array}
$$
In order to make the gloss matching learned by  $\mathsf{enc}_2$  aligned with the word embedding space, another MLM task is added to the training of  $\mathsf{enc}_2$ , with loss  $\mathcal{L}_{\mathrm{MLM}}^{(2)}$ . Thus, the complete loss function
of the SC model is
$$
\mathcal {L} _ {\mathrm {S C}} = \mathcal {L} _ {\mathrm {M L M}} ^ {(1)} + \mathcal {L} _ {\mathrm {M L M}} ^ {(2)} + \lambda \mathcal {L} _ {\mathrm {G R}} ^ {\mathrm {s}}. \tag {13}
$$
Although one gloss encoder and two contextual encoders are involved during training, only  $\mathsf{enc}_2$  is used at the inference stage. The estimation of contextual token distribution is given by  $\hat{p} (x|\mathbf{s}) = \mathrm{softmax}(\boldsymbol{v}_w^\top \boldsymbol {h}_s^{(2)})$ . By using two separate contextual encoders, the MLM task and the gloss matching task can be trained individually, which leads to better performance for each task. Besides, the combination of the two tasks corresponds to the theoretical derivation of  $p(x|\mathbf{s})$ , and integrates the gloss regularizer in a more natural and explainable way.
4.4 Gloss Regularized Pre-training
To pre-train the GR-BERT, we employ the gloss dataset from the online Oxford dictionary released by Chang et al. (2018); Chang and Chen (2019), formatted in triplets: word, sentence and definition  $(word_{i},\text{sent}_{ij},\text{def}_{ik})$ , where the human understandable gloss text  $def_{ik}$  describes the sense of target  $word_{i}$  in the sentence context  $sent_{ij}$ . The data consists of 677,191 pieces in total, including 31,889 words and 78,105 glosses.
We train the GR-BERT model with the gloss matching loss. For each triplet sample, context
sentij with word is encoded by the context encoder, and the target gloss defik is encoded by the gloss encoder. The cosine function is used as the similarity function  $\mathrm{sim}[\cdot ,\cdot ]$  in Eqn. (10). To build the negative gloss set in training, we use the in-batch negative sampling strategy (Chen et al., 2017). For each triplet sample (wordi, sentij, defik) in a batch, the positive glosses of other samples in the batch make up the negative gloss set  $T$  for the target word wordi. Since it would be relatively easy to distinguish the positive gloss defik from the randomly build-up in-batch negative set, we add the hard negative gloss to the negative gloss set. For each triplet (wordi, sentij, defik), we randomly pick another gloss defil of the target word wordi as the hard negative.
We utilize the BERT and RoBERTa models to initialize the context and gloss encoders, both of which are pre-trained on the gloss dataset for about 10 epochs. The gloss-matching accuracy is used as the metric to evaluate the model performance in pre-training. Detailed pre-training settings and hyper-parameters are provided in Appendix B.