Towards Multi-Sense Cross-Lingual Alignment of Contextual Embeddings

1 INTRODUCTION . Cross-lingual word embeddings ( CLWE ) provide a shared representation space for knowledge transfer between languages , yielding state-of-the-art performance in many cross-lingual natural language processing ( NLP ) tasks . Most of the previous works have focused on aligning static embeddings . To utilize the richer information captured by the pre-trained language model , more recent approaches attempt to extend previous methods to align contextual representations . Aligning the dynamic and complex contextual spaces poses significant challenges , so most of the existing approaches only perform coarse-grained alignment . Schuster et al . ( 2019 ) compute the average of contextual embeddings for each word as an anchor , and then learn to align the static anchors using a bilingual dictionary . In another work , Aldarmaki & Diab ( 2019 ) use parallel sentences in their approach , where they compute sentence representations by taking the average of contextual word embeddings , and then they learn a projection matrix to align sentence representations . They find that the learned projection matrix also works well for word-level NLP tasks . Besides , unsupervised multilingual language models ( Devlin et al. , 2018 ; Artetxe & Schwenk , 2019 ; Conneau et al. , 2019 ; Liu et al. , 2020 ) pretrained on multilingual corpora have also demonstrated strong cross-lingual transfer performance . Cao et al . ( 2020 ) and Wang et al . ( 2020 ) show that unsupervised multilingual language model can be further aligned with parallel sentences . Though contextual word embeddings are intended to provide different representations of the same word in distinct contexts , Schuster et al . ( 2019 ) find that the contextual embeddings of different senses of one word are much closer compared with that of different words . This contributes to the anisomorphic embedding distribution of different languages and causes problems for cross-lingual alignment . For example , it will be difficult to align the English word bank and its Japanese translations 銀行 and岸 that correspond to its two different senses , since the contextual embeddings of different senses of bank are close to each other while those of 銀行 and 岸 are far . Recently , Zhang et al . ( 2019 ) propose two solutions to handle multi-sense words : 1 ) remove multi-sense words and then align anchors in the same way as Schuster et al . ( 2019 ) ; 2 ) generate cluster level average anchor for contextual embeddings of multi-sense words and then learn a projection matrix in an unsupervised way with MUSE ( Conneau et al. , 2017 ) . They do not make good use of the bilingual dictionaries , which are usually easy to obtain , even in low-resource scenarios . Moreover , their projection-based approach still can not handle the anisomorphic embedding distribution problem . In this work , we propose a novel sense-aware cross entropy loss to model multiple word senses explicitly , and then leverage a sense level translation task on top of it for cross-lingual model pretraining . The proposed sense level translation task enables our models to provide more isomorphic and better aligned cross-lingual embeddings . We only use the cross-lingual signal from bilingual dictionaries for supervision . Our pretrained models demonstrate consistent performance improvements on zero-shot cross-lingual NER , sentiment classification and XNLI tasks . Though pretrained on less data , our model achieves the state-of-the-art result on zero-shot cross-lingual German NER task . To the best of our knowledge , we are the first to perform sense-level contextual embedding alignment with only bilingual dictionaries . 2 BACKGROUND : PREDICTION TASKS OF LANGUAGE MODELS . Next token prediction and masked token prediction are two common tasks in neural language model pretraining . We take two well-known language models , ELMo ( Peters et al. , 2018 ) and BERT ( Devlin et al. , 2018 ) , as examples to illustrate these two tasks ( architectures are shown in Appendix A ) . Next token prediction ELMo uses next token prediction tasks in a bidirectional language model . Given a sequence of N tokens ( t1 , t2 , . . . , tN ) , it first prepares a context independent representation for each token by using a convolutional neural network over the characters or by word embedding lookup ( a.k.a . input embeddings ) . These representations are then fed into L layers of LSTMs to generate the contextual representations : hi , j for token ti at layer j . The model assigns a learnable output embedding w for each token in the vocabulary , which has the same dimension as hi , L . Then , the forward language model predicts the token at position k with : p ( tk|t1 , t2 , . . . , tk−1 ) = softmax ( hTk−1 , Lwk′ ) = exp ( hTk−1 , Lwk′ ) ∑V i=1 exp ( h T k−1 , Lwi ) ( 1 ) where k′ is the index of token tk in the vocabulary , V is the size of the vocabulary , and ( w1 , . . . , wV ) are the output embeddings for the tokens in the vocabulary . The backward language model is similar to the forward one , except that tokens are predicted in the reverse order . Since the forward and backward language models are very similar , we will only describe our proposed approach in the context of the forward language model in the subsequent sections . Masked token prediction The Masked Language Model ( MLM ) in BERT is a typical example of masked token prediction . Given a sequence ( t1 , t2 , . . . , tN ) , this approach randomly masks a certain percentage ( 15 % ) of the tokens and generates a masked sequence ( m1 , m2 , . . . , mN ) , where mk = [ mask ] if the token at position k is masked , otherwise mk = tk . BERT first prepares the context independent representations ( x1 , x2 , . . . , xN ) of the masked sequence via token embeddings . It is then fed into L layers of transformer encoder ( Vaswani et al. , 2017 ) to generate “ bidirectional ” contextual token representations . The final layer representations are then used to predict the masked token at position k as follows : p ( mk = tk|m1 , . . . , mN ) = softmax ( hTk , Lwk′ ) = exp ( hTk , Lwk′ ) ∑V i=1 exp ( h T k , Lwi ) ( 2 ) where k′ , V , h and w are similarly defined as in Eq . 1 . Unlike ELMo , BERT ties the input and output embeddings . 3 PROPOSED FRAMEWORK . We first describe our proposed sense-aware cross entropy loss to model multiple word senses explicitly in language model pretraining . Then , we present our joint training approach with sense alignment objective for cross-lingual mapping of contextual word embeddings . The proposed framework can be applied to most of the recent neural language models , such as ELMo , BERT and their variants . See Table 1 for a summary of the main notations used in this paper . 3.1 SENSE-AWARE CROSS ENTROPY LOSS Table 1 : Summary of the main notations Notation Description tk k-th token in sentence tk , s s-th sense of tk k′ index of token tk in vocabulary L number of LSTM/Transformer layers V size of vocabulary S maximum number of senses per token hk , j contextual representation of token tk in layer j hk∗ , L contextual representation used in softmax function for predicting tk vi i-th word in vocabulary vi , s s-th sense of vi wi output embedding of vi wi , s context-dependent output embedding ( i.e . sense vector ) of vi , s ci , s sense cluster center of vi , s Ci sense cluster centers of vi d dimension of contextual representations P projection matrix for dimension reduction Limitations of original training objectives The training tasks with Eq . 1 and 2 maximize the normalized dot product of contextual representations ( hk−1 , L or hk , L ) with a weight vector wk′ . The only difference is that hk−1 , L in Eq . 1 encodes the information of previous tokens in the sequence , while hk , L in Eq . 2 encodes the information of the masked sequence . Therefore , without loss of generality , we use hk∗ , L to denote the contextual representation for predicting the next or masked token tk . Even though contextual language models like ELMo and BERT provide a different token representation for each distinct context , the learned representations are not guaranteed to be sense separated . For example , Schuster et al . ( 2019 ) computed the average of ELMo embeddings for each word as an anchor , and found that the average cosine distance between contextual embeddings of multi-sense words and their corresponding anchors are much smaller than the average distance between anchors , which mean that the embeddings of different senses of one word are relatively near to each other comparing to that of different words . We also observed the same with BERT embeddings . This finding suggests that sense clusters of a multi-sense word ’ s appearances are not well separated in the embedding space , and the current contextual language models still have room for improvement by considering finer-grained word sense disambiguation . Notice that there is only one weight vector wk′ for predicting the token tk in the original training tasks . Ideally , we should treat the appearances of a multi-sense word in different contexts as different tokens , and train the language models to predict different senses of the word . In the following , we propose a novel sense-aware cross entropy loss to explicitly model different senses of a word in different contexts . Sense-aware cross entropy loss Given a sequence ( t1 , t2 , . . . , tN ) , our proposed framework generates contextual representations ( hk , j for token tk in layer j ∈ { 1 , . . . , L } ) in the same way as the standard LMs . Different from existing methods , our approach maintains multiple context-dependent output embeddings ( henceforth , sense vectors ) for each token . Specifically , let S be the maximum number of senses per token . Each word vi in the vocabulary contains S separate sense vectors ( wi,1 , wi,2 , . . . , wi , S ) , where each wi , s corresponds to a different sense ( see Appendix for some interesting visualization examples ) . Following the notation in Section 2 , we use k′ to denote the index of the output token tk in the vocabulary . Therefore , the sense vectors of tk can be represented by ( wk′,1 , wk′,2 , . . . , wk′ , S ) , which are randomly initialized and of the same dimension as hk∗ , L . Note that we untie the input and output embeddings in our framework . We propose a word sense selection method shown in Algorithm 1 to select the most likely sense vector when training with sense-level cross entropy loss . Figure 1 shows the architecture of our proposed models . Assuming sense s′ is selected for token tk ( which means sense vector wk′ , s′ should be used ) , we have the following new prediction task : p ( tk , s′ |context ) = softmax ( hTk∗ , Lwk′ , s′ ) = exp ( hTk∗ , Lwk′ , s′ ) ∑V i=1 ∑S s=1 exp ( h T k∗ , Lwi , s ) ( 3 ) The sense-aware cross entropy loss for word sense prediction is defined as follows : LSENSE = − log ( p ( tk , s′ |context ) ) ( 4 ) Word sense selection algorithm Word sense selection when training the language model can be handled as a non-stationary data stream clustering problem ( Aggarwal et al. , 2004 ; Khalilian & Mustapha , 2010 ; Abdullatif et al. , 2018 ) . The most intuitive way to select the corresponding sense 1Since the backward language model is similar to the forward , we only show the forward one for simplicity . vector for hk∗ , L is to select the vector wk′ , s with the maximum dot product value hTk∗ , Lwk′ , s , or cosine similarity value cossim ( hk∗ , L , wk′ , s ) . However , our experiments show that these methods do not work well due to curse of dimensionality , suboptimal learning rate and noisy hk∗ , L . We apply an online k-means algorithm to cluster different senses of a word in Algorithm 1 . For each sense vector wi , s , we maintain a cluster center ci , s which is of the same dimension as wi , s . Therefore , each token vi in the vocabulary has S such cluster center vectors , denoted by Ci = ( ci,1 , ci,2 , . . . , ci , S ) . When predicting token tk in a given sequence , we apply Algorithm 1 to select the best sense vector based on hk , L ( see Figure 1 ) . Notice that hk , L is different from hk∗ , L for next token prediction ( Figure 1a ) for which hk∗ , L = hk−1 , L . The cluster centers Ci are not neural network parameters ; instead , they are randomly initialized using a normal distribution N ( 0 , σ2 ) and updated through Algorithm 1 . In addition , we also maintain a projection matrix P for dimension reduction to facilitate effective sense clustering . P ∈ Rd×d′ projects hk , L and ci , s from dimension d to d′ , and is shared by all tokens in vocabulary . Similar to C , P is also randomly initialized with normal distribution N ( 0 , 1 ) , and then updated through Algorithm 2 . Both Algorithm 1 and 2 run in parallel , and are interrupted when the language model stops training . Some rationales behind our algorithm design are the following : Algorithm 1 Word sense selection 1 : Hyper-parameters : number of senses S , sense learning rate α 2 : Initialize the set of all sense cluster centers C 3 : repeat 4 : input : hk , L , vocabulary index k′ of the token to predict 5 : Lookup sense cluster centers for k′ : Ck′ = { ck′,1 , ck′,2 , . . . , ck′ , S } 6 : P = updated projection matrix from Alg . 2 7 : if cosine similarity between ck′ , s′P and h′kP is the largest among the vectors in Ck′ then 8 : ck′ , s′ = ( 1− α ) ck′ , s′ + αhk , L 9 : output : s′ ( wk′ , s′ should be selected ) 10 : end if 11 : until interrupted Algorithm 2 Projection matrix P update 1 : Hyper-parameters : projection dimension d′ , up- date interval M , queue size Q 2 : Initialize P withN ( 0 , 1 ) , queue H = ∅ , m = 0 3 : repeat 4 : input : hk , L 5 : m = m+ 1 6 : Add hk , L to queue H 7 : if size ( H ) > Q then 8 : Pop the oldest element from queue H . 9 : end if 10 : if m > =M then 11 : P = the first d′ PCA components of H 12 : m = 0 13 : end if 14 : output : P 15 : until interrupted • Directly computing cosine similarity between ck′ , s and hk , L suffers from the curse of dimensionality . We maintain P for dimension reduction . Although many algorithms use random projection for dimension reduction , we find using PCA components can help improve clustering accuracy . • Since the neural model parameters keep being updated during training , the sense clusters become non-stationary , i.e. , their locations keep changing . Experiments shows that when using P for dimension reduction , a slightly larger projection dimension d′ will make the clustering algorithm less sensitive to cluster location change . We use d′ = 16 for ELMo , and d′ = 14 for BERT . We also notice that the sense clustering works well even if P is updated sporadically . We can set a relatively large update interval in Algorithm 2 to reduce computation cost . • A separate sense learning rate α should be set for the clustering algorithm . A large α makes the algorithm less robust to noise , while a small α leads to slow convergence . • It is essential to use the current token ’ s contextual representation hk , L for sense selection even though we use hk∗ , L = hk−1 , L in the next token prediction task . If we use hk−1 , L for sense selection , experiments show that most of the variance comes from input embedding xk−1 . This introduces too much noise for word sense clustering . Dynamic pruning of redundant word senses To make the training more efficient , we keep track of relative sense selection frequency for each token in the vocabulary . Assume token vi has initial senses ( vi,1 , vi,2 , . . . , vi , S ) , for which we compute the relative frequency ρ ( vi , s ) such that 0 ≤ ρ ( vi , s ) ≤ 1 and ∑ s ρ ( vi , s ) = 1 . A lower ρ ( vi , s ) means the sense is less frequently selected compared with others . We check the relative frequencies after every E training steps , and if ρ ( vi , s ) < β ( a threshold hyper-parameter ) , vi , s is removed from the list of senses of vi . Remark on model size and parameters The sense cluster centers C and the projection matrix P are only used to facilitate sense selection during model pretraining , which are not neural model parameters . The sense vectors wi , s will no longer be used after pretraining , which can also be discarded . Therefore , our models and the original models have exactly the same number of parameters when transferred to downstream tasks . Remark on model complexity The computational complexity of our algorithm is linear with respect to the size of data , so our method is scalable to train on very large datasets .

Summary: This paper proposes to introduce multiple senses into pre-trained models. The proposed method selects senses dynamically while pretraining the model and applies a sense-aware cross-entropy loss for pretraining. This paper further proposes to jointly pre-train a sense-aware cross-lingual model with sense-level translation. The proposed model yields better performance than the baseline models under both monolingual and cross-lingual setting.