Rethinking Positional Encoding in Language Pre-training

1 INTRODUCTION . The Transformer model ( Vaswani et al. , 2017 ) is the most widely used architecture in language representation learning ( Liu et al. , 2019 ; Devlin et al. , 2018 ; Radford et al. , 2019 ; Bao et al. , 2020 ) . In Transformer , positional encoding is an essential component since other main components of the model are entirely invariant to sequence order . The original Transformer uses the absolute positional encoding , which provides each position an embedding vector . The positional embedding is added to the word embedding , which is found significantly helpful at learning the contextual representations of words at different positions . Besides using the absolute positional encoding , Shaw et al . ( 2018 ) ; Raffel et al . ( 2019 ) further propose the relative positional encoding , which incorporates some carefully designed bias term inside the self-attention module to encode the distance between any two positions . In this work , we revisit and study the formulation of the widely used absolute/relative positional encoding . First , we question the rationality of adding the word embedding with the absolute positional embedding in the input layer . Since the two kinds of embeddings are apparently heterogeneous , this addition operation brings mixed correlations1 between the positional information and word semantics . For example , by expanding the dot-production function of keys and values in the self-attention module of the first layer , we find that there are explicit terms that use “ word ” to query “ positions ” and vice versa . However , words may only have weak correlations to where they appear in the sentence . Our empirical analysis also supports this by showing that in a well-trained model , such correlation is noisy . Second , we notice that the BERT model does not only handle natural language words . A special symbol [ CLS ] is usually attached to the sentence . It is widely acknowledged that this symbol ∗Correspondence to : { guolin.ke , dihe } @ microsoft.com 1The term “ correlation ” mainly refers to the dot product between Key and Query in the self-attention module . receives and summarizes useful information from all the positions , and the contextual representation of [ CLS ] will be used as the representation of the sentence in the downstream tasks . As the role of the [ CLS ] symbol is different from regular words that naturally contain semantics , we argue that it will be ineffective if we treat its position the same as word positions in the sentence . For example , if we apply the relative positional encoding to this symbol , the attention distribution of some heads will likely be biased to the first several words , which hurts the understanding of the whole sentence . Based on the investigation above , we propose several simple , yet effective modifications to the current methods , which lead to a new positional encoding called Transformer with Untied Positional Encoding ( TUPE ) for language pre-training , see Figure 1 . In TUPE , the Transformer only uses the word embedding as input . In the self-attention module , different types of correlations are separately computed to reflect different aspects of information , including word contextual correlation and absolute ( and relative ) positional correlation . Each kind of correlation has its own parameters and will be added together to generate the attention distribution . A specialized positional correlation is further set to the [ CLS ] symbol , aiming to capture the global representation of the sentence correctly . First , we can see that in TUPE , the positional correlation and word contextual correlation are de-coupled and computed using different parameters . This design successfully removes the randomness in word-to-position ( or position-to-word ) correlations and gives more expressiveness to characterize the relationship between a pair of words or positions . Second , TUPE uses a different function to compute the correlations between the [ CLS ] symbol and other positions . This flexibility can help the model learn an accurate representation of the whole sentence . We provide an efficient implementation of TUPE . To validate the method , we conduct extensive experiments and ablation studies on the GLUE benchmark dataset . Empirical results confirm that our proposed TUPE consistently improves the model performance on almost all tasks . In particular , we observe that by imposing this inductive bias to encode the positional information , the model can be trained more effectively , and the training time of the pre-training stage can be largely improved . 2 PRELIMINARY . 2.1 ATTENTION MODULE . The attention module ( Vaswani et al. , 2017 ) is formulated as querying a dictionary with key-value pairs , e.g. , Attention ( Q , K , V ) = softmax ( QK T √ d ) V , where d is the dimensionality of the hidden representations , and Q ( Query ) , K ( Key ) , V ( Value ) are specified as the hidden representations of the previous layer . The multi-head variant of the attention module is popularly used which allows the model to jointly attend to the information from different representation sub-spaces , and is defined as Multi-head ( Q , K , V ) = Concat ( head1 , · · · , headH ) WO headk = Attention ( QW Q k , KW K k , V W V k ) , ( 1 ) where WQk ∈ Rd×dK , WKk ∈ Rd×dK , WVk ∈ Rd×dV , and WO ∈ RHdV ×d are learnable project matrices , H is the number of heads . dK and dV are the dimensionalities of Key and Value . The self-attention module is one of the key components in Transformer and BERT encoder ( Devlin et al. , 2018 ) . For simplicity , we use the single-head self-attention module and set dK = dV = d for a demonstration . We denote xl = ( xl1 , x l 2 · · · , xln ) as the input to the self-attention module in the l-th layer , where n is the length of the sequence and each vector xli ∈ Rd is the contextual representation of the token at position i. zl = ( zl1 , z l 2 · · · , zln ) is the output of the attention module . Then the self-attention module can be written as zli = n∑ j=1 exp ( αij ) ∑n j′=1 exp ( αij′ ) ( xljW V , l ) , where αij = 1√ d ( xliW Q , l ) ( xljW K , l ) T . ( 2 ) As we can see , the self-attention module does not make use of the order of the sequence , i.e. , is permutation-invariant . However , natural language is well-structured and word order is important for language understanding ( Sutskever et al. , 2014 ) . In the next section , we show several previous works that proposed different ways of incorporating positional information into the Transformer model . 2.2 POSITIONAL ENCODING . Generally , there are two categories of methods that encode positional information in the Transformer model , absolute positional encoding and relative positional encoding . Absolute Positional Encoding . The original Transformer ( Vaswani et al. , 2017 ) proposes to use absolute positional encoding to represent positions . In particular , a ( learnable ) real-valued vector pi ∈ Rd is assigned to each position i . Given a sentence , pi will be added to the word embedding wi at position i , and wi + pi will be used as the input to the model , e.g , x1i = wi + pi . In such a way , the Transformer can differentiate the words coming from different positions and assign each token position-dependent attention . For example , in the self-attention module in the first layer , we have αAbsij = 1√ d ( ( wi + pi ) W Q,1 ) ( ( wj + pj ) W K,1 ) T . ( 3 ) Relative Positional Encoding . In absolute positional encoding , using different pi for different position i helps Transformer distinguish words at different positions . However , as pointed out in Shaw et al . ( 2018 ) , the absolute positional encoding is not effective for the model to capture the relative word orders . Therefore , besides using absolute positional encoding , Shaw et al . proposes a relative positional encoding as an inductive bias to help the learning of attention modules , αRelij = 1√ d ( xliW Q , l ) ( xljW K , l + alj−i ) T , ( 4 ) where alj−i ∈ Rd is learnable parameter and can be viewed as the embedding of the relative position j − i in layer l. In this way , embedding alj−i explicitly models the relative word orders . T5 ( Raffel et al. , 2019 ) further simplifies it by eliminating alj−1 in Query-Key product . αT5ij = 1√ d ( xliW Q , l ) ( xljW K , l ) T + bj−i . ( 5 ) For each j − i , bj−i is a learnable scalar2 and shared in all layers . 2Specifically , in Shaw et al . ( 2018 ) ; Raffel et al . ( 2019 ) , the relative position j − i will be first clipped to a pre-defined range , e.g. , clip ( j − i , −t , t ) , t = 128 . The embedding is defined over the possible values of the clipped range , i.e. , [ −t , t ] . Besides , Shaw et al . also tried to add vector aV , lj−i to the value V in the output of self-attention , but the experiment results indicate that it did not improve much . 3 TRANSFORMER WITH UNTIED POSITIONAL ENCODING . 3.1 UNTIE THE CORRELATIONS BETWEEN POSITIONS AND WORDS . In absolute positional encoding , the positional embedding is added together with the word embedding to serves as the input to the neural networks . However , these two kinds of information are heterogeneous . The word embedding encodes the semantic meanings of words and word analogy tasks can be solved using simple linear arithmetic on word embeddings ( Mikolov et al. , 2013 ; Pennington et al. , 2014 ; Joulin et al. , 2016 ) . On the other hand , the absolute positional embedding encodes the indices in a sequence , which is not semantic and far different from word meanings . We question the rationality of the linear arithmetic between the word embedding and the positional embedding . To check clearly , we take a look at the expansion of Eq . ( 3 ) . αAbsij = ( ( wi + pi ) W Q,1 ) ( ( wj + pj ) W K,1 ) T√ d = ( wiW Q,1 ) ( wjW K,1 ) T√ d + ( wiW Q,1 ) ( pjW K,1 ) T√ d + ( piW Q,1 ) ( wjW K,1 ) T√ d + ( piW Q,1 ) ( pjW K,1 ) T√ d ( 6 ) The above expansion shows how the word embedding and the positional embedding are projected and queried in the attention module . We can see that there are four terms after the expansion : word-to-word , word-to-position , position-to-word , and position-to-position correlations . We have several concerns regarding this formulation . First , it is easy to see that the first and the last term characterize the word-word and position-position relationships respectively . However , the projection matrices WQ , l and WK , l are shared in both terms . As the positional embedding and the word embedding encode significantly different concepts , it is not reasonable to apply the same projection to such different information . Furthermore , we also notice that the second and the third term use the position ( word ) as the query to get keys composed of words ( positions ) . As far as we know , there is little evidence suggesting that the word and its location in a sentence have a strong correlation . Furthermore , in BERT and recently developed advanced methods such as RoBERTa ( Liu et al. , 2010 ) , sentences are patched in a random way . For example , in BERT , each input contains multiple sentences and some of the sentences are negatively sampled from other documents to form the next sentence prediction task . Due to the random process of batching , it is possible that a word can even appear at any positions and the correlations between words and positions could be weak . To further investigate this , we visualize the four correlations in Eq . ( 6 ) on a pre-trained BERT model . We find that the second and the third term looks uniform across positions , as shown in Figure 2 . This phenomenon suggests that there are no strong correlations3 between the word and the absolute position and using such noisy correlation may be inefficient for model training . 3Some recent works ( Yang et al. , 2019 ; He et al. , 2020 ) show that correlations between relative positions and words can improve the performance . Our results have no contradiction with theirs as our study is on the correlations between word embeddings and absolute positional embeddings . Our modification . To overcome these problems , we propose to directly model the relationships between a pair of words or positions by using different projection matrices and remove the two terms in the middle . That is , we use αij = 1√ 2d ( xliW Q , l ) ( xljW K , l ) T + 1√ 2d ( piU Q ) ( pjU K ) T , ( 7 ) where UQ , UK ∈ Rd×d are the projection matrice for the positional embedding , and scaling term 1√ 2d is used to retain the magnitude of αij ( Vaswani et al. , 2017 ) . A visualization is put in Figure 3 . Our proposed method can be well combined with the relative positional encoding in Raffel et al . ( 2019 ) by simply changing Eq . ( 5 ) to αij = 1√ 2d ( xliW Q , l ) ( xljW K , l ) T + 1√ 2d ( piU Q ) ( pjU K ) T + bj−i . ( 8 ) 3.2 UNTIE THE [ CLS ] SYMBOL FROM POSITIONS Note that in language representation learning , the input sequence to the Transformer model is not always a natural sentence . In BERT , a special symbol [ CLS ] is attached to the beginning of the input sentence . This symbol is designed to capture the global information of the whole sentence . Its contextual representation will be used to make predictions in the sentence-level downstream tasks after pre-training ( Devlin et al. , 2018 ; Liu et al. , 2019 ) . We argue that there will be some disadvantages if we treat this token the same as other natural words in the attention module . For example , the regular words usually have strong local dependencies in the sentence . Many visualizations ( Clark et al. , 2019a ; Gong et al. , 2019 ) show that the attention distributions of some heads concentrate locally . If we process the position of [ CLS ] the same as the position of natural language words , according to the aforementioned local concentration phenomenon , [ CLS ] will be likely biased to focus on the first several words instead of the whole sentence . It will potentially hurt the performance of the downstream tasks . Our modification . We give a specific design in the attention module to untie the [ CLS ] symbol from other positions . In particular , we reset the positional correlations related to [ CLS ] . For better demonstration , we denote vij as the content-free ( position-only ) correlation between position i and j . For example , when using the absolute positional encoding in Eq . ( 7 ) , vij = 1√2d ( piU Q ) ( pjU K ) T ; when using relative positional encoding in Eq . ( 8 ) , vij = 1√2d ( piU Q ) ( pjU K ) T + bj−i . We reset the values of vij by the following equation : resetθ ( v , i , j ) = vij i 6= 1 , j 6= 1 , ( not related to [ CLS ] ) θ1 i = 1 , ( from [ CLS ] to others ) θ2 i 6= 1 , j = 1 , ( from others to [ CLS ] ) , ( 9 ) where θ = { θ1 , θ2 } is a learnable parameter . A visualization is put in Figure 4 .

Summary: The paper delves into the nature of positional encoding in Transformer and variants (especially BERT). The paper points out that, attention weight computation obtained from the addition of the word embedding and the position embedding in the first layer can be expanded into four terms, namely word-word, position-position, and two word-position (in both directions). The paper points out that word-position relationship is largely meaningless, which is also demonstrated via visualizing the values in pertained BERT. Then the paper proposes to remove these two terms when computing the attention. Secondly, the paper mentions that special tokens in BERT (e.g. [CLS]) should be treated in a different way than typical words. To do so, it proposes to replace the attention logic with a learnable parameter if the attention involves one or two special tokens. Both methods combined show a significant advantage on multiple GLUE tasks when applied to BERT.