Finding One Missing Puzzle of Contextual Word Embedding: Representing Contexts as Manifold

1 INTRODUCTION . In modern natural language processing , using vector representation of words in a low dimension space is a common practice . In early days , word embeddings were static , i.e. , the representation is solely determined by a word ’ s identity , such as the case of word2vec ( Mikolov et al. , 2013 ) and GloVe ( Pennington et al. , 2014 ) . In contrast , contextualized word embeddings , e.g. , ELMo ( Peters et al. , 2018 ) , BERT ( Devlin et al. , 2019 ) and GPT ( Brown et al. , 2020 ) , revolutionized this technique by introducing information from the contexts to the central word . When compared to static embedding , studies to understand the contextual embedding generally lag behind . This is particularly unfavoured as the superior performance of pre-trained contextualized representations is gaining more and more attention . To tackle this problem , previous studies have proposed several directions to explore . For example , researchers in Hewitt & Manning ( 2019 ) ; Coenen et al . ( 2019 ) ; Wu et al . ( 2020 ) develop several probing techniques to relate contextualized word embedding to syntactic features . On the other hand , several works try to uncover how contexts influence the geometry of embedding space . Ethayarajh ( 2019 ) measures the context-specificity of the word embeddings from different layers of several deep representation models and confirms the majority variance is provided by the context but not the token . They also point out that the embedding space is generally anisotropic . Then , results from Cai et al . ( 2021 ) indicate that isotropy actually exists in certain isolated clusters of words in the representation space . Mamou et al . ( 2020 ) proposes to analyze the manifolds defined by word identity and other linguistic features . Under both “ predictive ” ( i.e. , with [ MASK ] tokens ) and “ contextualized ” ( i.e. , the real sentence ) settings , the authors analyze how well the manifolds can be separated from each other . While these pioneering works offer some useful insight , their starting points are generally “ tokencentric ” , i.e. , the analysis focuses on how the context affects the representation of tokens . Therefore , very few studies touch upon another critical problem , i.e. , for contextualized word embedding methods , how exactly does the model represent the context itself . In this study , to fill in this gap , we introduce a novel perspective to analyze the representation of context . To give a rigorous discussion , we first provide a set of mathematical definitions by formulating the language representation prob- lem using category theory . Then , inspired by manifold learning , we provide two analysis scenarios and also methods to probe the representation of contexts . As a realization of these theoretical analyses , we further provide empirical results and cases studies , which give more insights into the current contextual representation models . A particularly interesting example can be the re-thinking of the functions of Multi-head attention and Feed Forward Network ( FFN ) in a Transform-based model . 2 MOTIVATION AND CONTRIBUTIONS . This work tackles an untouched fundamental problem in language representation learning , i.e. , how a representation model represent different contexts . Therefore , it provides a novel perspective to re-think and characterize these models . In a contextualized representation model , the representation at a specific sequence position is determined by both the token at that place and the surrounding context . However , most of the attention has been paid to one direction , i.e. , understanding how the vector represents the token , leaving the context part largely unexplored . Intuitively , we argue that the effect of different context conditions ( termed as observed condition below ) should be semantically related . This point , however , is not explicitly modeled in the current representation scheme , which motivates us to develop this study . Ideally , a representation scheme of context should have the following two features . Firstly , it should be compatible with the word representation mentioned above . Secondly , the effect of the central token and observed condition should be disentangled , which makes the latter focusing on how different they are in affecting the central word ’ s semantics . To achieve these goals , we design the following representational scheme . Firstly , for each observed condition , we fix the context , mutate the central token , and record how much is changed about the representation vector of the central token . This generates a collection of vectors whose each element corresponds to a contextualized representation . We hope this kind of collection reflects how each observed condition affects the semantics of the different central tokens , and comparing different collections reflect how different contexts affect the semantics . On the other hand , as mentioned above , to disentangle the effect of observed condition and central token , we study the nature of the collection , other than its elements . However , characterizing this kind of information is challenging . Obviously , taking the collection of mutated sequences ( and the associated vectors ) as a set would miss a lot of information among them . Therefore , we propose a systematic and rigorous mathematical formulation based on category theory which states clearly the representation of context should reflect both the set of representations of the substituted central tokens as well as how these sets of word representations can transit to each other . In particular , we introduce the concept “ morphism ” to characterize the mappings between different token substitution sets ( corresponding to the semantic effects of different observed conditions ) . Furthermore , with of help of manifold learning , we also develop a tool set to analyze the aforementioned information experimentally , from either a functional or a topological perspective . Besides the above contributions , our empirical results , enabled by these theoretical foundations , also generate several interesting findings . First , our analysis of context around [ CLS ] token suggests that high-confidence correct predictions tend to occur when contexts show robustness in modulating the semantics of the central token as well as consistency of the modulating effect across different samples . Last but not least , our dissection of the attention and FFN layers identifies a game between them when information gets through a deep model , and their final agreement actually supports our hypothesis that they both of them try to fit an identity morphism that reflects bona fide linguistic associations . 3 PRELIMINARY : CONTEXTUAL LANGUAGE REPRESENTATION . Following the distributional hypothesis ( Harris , 1954 ) , the output vector of a representational model ( denoted by G ) represents the probability of the appearance of certain tokens ( i.e. , context ) around a given token ( i.e. , central token x ) , i.e. , p ( context|x ) . Note that here context is a random variable , whose formal definition will be given in the next section . In contextualized language modeling ( Peters et al. , 2018 ) , the distribution of context is conditioned not only on the central token but also on what really appears around the central token in each sample ( e.g. , the unmasked tokens in a given sentence ) . In this sense , we introduce observed condition C to denote the observed context . That is to say , observed condition C corresponds to an actual value of random variable context given a specific sample , which makes the modeling target of contextualized language modeling be p ( context|C , x ) . In this work , we propose to represent each observed condition C by inspecting how the model represents various central tokens given this condition , which gives the following collection of vector , i.e. , { G ( C , x ) |x ∈ V ocab } ( 1 ) where V ocab is the set of distinct tokens used by he tokenizer . However , the definition of set only provides rough description of the nature of contexts , which motivates as to introduce category for a rigorous definition , as described below . 4 CATEGORY THEORY FORMULATION AND NOTATIONS . As mentioned in Section 2 , in our problem formulation , “ context ” is a random variable . So here we first give the encoding of each observed condition ( i.e. , an actually appearing context ) using the formal language of probability theory ( Stroock , 2010 ) . Specifically , in BERT-style language model , we use a three-dimensional tensor ( corresponding to token species , position and type , respectively ) to define the sample space and take the σ-field of the sample space as the event space . Note that to describe the collection of probability distribution of different central tokens , we introduce a discrete random variable set K = { ki } and the distribution of each ki ’ s image corresponds to the distribution of context given a certain central token i ( Note that i is also defined as a triplet of species , position and type ) . Concretely , we have : Definition 4.1 . Let K be a finite discrete set whose each element ki : Ω → E is a random variable that projects from the discrete sample space of tokens Ω ∈ NN to its σ-field E , where N is set to three , corresponding to token species , token type and token position , respectively . The cardinal number of K is marked as n. Therefore , each element ki indicates the presence of a certain token ( with defined species , type and position ) as the central token that accompanies the context . Following the principle of distributional hypothesis , the representation of each central token represents the a probability density function of the context , which gives : Definition 4.2 An observed condition C is encoded as a set F ( C ) whose each element can be expressed as a probability density function fki conditioning on the observed condition C , i.e. , F ( C ) = { fki ( context = e|C , context ∈ E ) |i ∈ Ω } . Ultimately , one expects the representation of observed conditions to reveal a comprehensive characterization that transfers well to a general scope of downstream applications . In particular , the following two features are considered here , 1 ) for each observed condition , how it can affect the semantics of the central token , and 2 ) how different observed conditions modulate the semantics differently . This intuition is reminiscent of category theory that explicitly study the objects and their mappings ( i.e. , morphisms ) . Definition 4.3 ( Category of observed conditions ) To define the objects in a category of observed conditions , we have ob ( F ) = { F ( C ) |C ∈ E } ( 2 ) which contains the probability density of all possible context events . Naturally , we can define an abstract collection of morphisms that map from one object to another , i.e. , hom ( F ) = { hom ( Fi , Fj ) |Fi , Fj ∈ ob ( F ) } , ( 3 ) where hom ( Fi , Fj ) ≜ { Tij : Fi → Fj |Tij ( fk ( X = e|Ci , e ∈ E ) ) = fk ( X = e|Cj , e ∈ E ) } , ∀k ∈ K } . ( 4 ) Note that through the above definition , we introduce several constrains that reflect the nature of language . First , for each set of morphism between two objects , there only exists one morphism , assuming that the model should learn to capture one genuine relationship between two observed conditions . Second , this morphism is a bijection defined on F ( note the subscript k in Eq ( 4 ) ) , so each probability density of context should be projected to a unique partner with the same central token and vice versa ( i.e. , this morphism would not cause semantic collision when central tokens are different ) . We will discuss how our experimental results support these constrains later . After the definition of categories of observed conditions , we move on to its representation . Note that each observed condition corresponds to a token substitution set ( Definition 4.2 ) , and a model G is trained to perform injection from each probability density function in the set to a representation vector , i.e. , G ( C ) : F → M , where M consists of vectors in a d-dimensional Euclidean space . Therefore , we can similarly define the category of vectors as the representation of category of observed conditions . What is different here is we how have to take different representation models into consideration , which gives the following definition : Definition 4.4 ( Category of vector representations ) Suppose we have a set of G models : Gset = { Gm } , the objects of category of vector representations can be defined as ob ( M ) = { MCiGm } , ( 5 ) where MCiGm = { Vk|k ∈ K , Gm ∈ Gset , Vk ≜ Gm ( fk ( X = e|Ci , e ∈ E ) ) } , ( 6 ) The definition of morphism collections also considers the difference in model , which gives : hom ( M ) = { hom ( MCiGm , MCjGn ) |MCiGm , MCjGn ∈ ob ( M ) } ( 7 ) where hom ( MCiGm , MCjGn ) = DCiGm→CjGn : MCiGm → MCjGn , ( 8 ) DCiGm→CjGn ( Gm ( F ( Ci ) ) ) = Gn ( F ( Cj ) ) , ( 9 ) and Ci and Cj stand for two different observed conditions . Given the above definitions , we can then formulate the representation of context by a functor L from this category M to category F , i.e. , L ( Gm ( Ci ) ) = F ( Ci ) , ( 10 ) L ( DCiGm→CjGn ) = Tij , ∀DCiGm→CjGn ∈ hom ( MCiGm , MCjGn ) , ( 11 ) where Gm , Gn ∈ Gset .

Summary: This paper reports several empirical findings.  1. Replacing the [CLS] tokens after the fine-tuning does not change the output representation a lot.  2. By using UMAP to analyze the patterns of nearest neighbors in different Transformer layers, it finds that the topological structure in layer 1 is most similar to the structure in layer 2 and less similar to that in layer 24. This shows that each Transformer layer would change the topological structure of the context embeddings. 3. Let's denote the embeddings in the Transformer in this way (H1) -> ATT -> (A1) -> FFN -> (H2) -> ATT -> (A2) -> FFN -> (H3) ..., where (Hk) is the hidden states, ATT is the multi-head attention, (Ak) is the output of the attention layer, FFN is the feed-forward network. Figure 3 (c) shows that (Hk) and (Ak) are very different when k is small, but become similar when k is large. Compared to (Hk) and (Ak), (Hk) and (Hk-1) are more similar when k is small. Similarly, (Ak) and (Ak-1) are more similar than (Hk) and (Ak) when k is small.