Reconstructing Word Embeddings via Scattered $k$-Sub-Embedding

1 INTRODUCTION . Embedding is useful for mapping high-dimensional vector to a low-dimensional space for various neural network models such as language models and graph neural networks . Especially in modern language models , a contextual embedding is represented respectively for each contextual unit . For instance , some language models like word2vec ( Mikolov et al. , 2013 ) which deal with a whole word as one embedding vector , suffer from the Out-of-Vocabulary ( OOV ) problem due to the variety of words . Later language models ( Bojanowski et al. , 2017 ) subdivide words into multiple segmented words . To deal with the subdividing , various tokenizers ( Sennrich et al. , 2016 ; Kudo & Richardson , 2018 ) are proposed to construct common tokens based on word corpuses . However , each token contains a different context , and a good criterion for dividing the token is not clear . There are some cases that the atomic-level tokens share notable common properties . For example , two tokens 〈A〉 and 〈a〉 are assigned to different embeddings , but they both are vowels and the same alphabet . In the case of combined characters ( i.e. , Hangul of Korean character ) or ideograms ( i.e. , Chinese and Japanese characters ) , some tokens might include more common meanings than other languages . The tokens of similar meanings tend to locate closely in embedding space when the language model is trained with those tokens . In this paper , we propose a novel embedding structure which replaces a word embedding with several sub-embeddings . If a token consists of several contextual elements , a sub-embedding can be assigned to each element to constitute the original embedding . An embedding vector shares its learning parameters with other embedding vectors because of the nested allocated sub-embeddings . Various scattered sub-embedding structures can be generated depending on structural topology of the sub-embeddings . Firstly , we sequentially allocate sub-embeddings through module operation , and conduct Cartesian product with sub-embeddings to build embedding vectors ( Figure 1 ) . This could robustly map the sub-embeddings , but each sub-embedding does not reflect the context of each token . To address this problem , we modify the scattering algorithm by clustering the tokens in the embedding space using a pretrained network . Notable factors of this paper are as follows : 1 . The number of learning parameters in embedding part is dramatically reduced by Cartesian product . 2 . The language models are free from OOV problem through sub-embedding . 3 . The proposed embedding structure can be applied to most existing language models easily by replacing only the input embedding part . We evaluate our sub-embedding structure on English and multilingual downstream tasks . We borrow RoBERTa ( Liu et al. , 2019 ) structure for the standard English downstream tasks . In addition , XLMR ( Conneau et al. , 2020 ) is used to test the performance of sub-embedding in multilingual tasks . We show that the network replaced by the k-sub-embedding compresses the embedding over 99.9 % while the accuracy on multilingual benchmark drops only 1.4 % p . Finally , we demonstrate how the sub-embedding is trained by visualizing the distribution of k-sub-embeddings . 2 RELATED WORK . Distributed Representations of Word Embedding . Word embedding is commonly used in the language models to represent the word in the latent space . Distributed representation is mainly used instead of a one-hot vector for the efficiency . Despite the difficulty of learning distributed representation , word2vec ( Mikolov et al. , 2013 ) is able to map words to embedding space successfully and shows context analogies of each word embedding . In the case of word2vec , the vocabulary is composed with the words in the input corpus , it is easily suffered from OOV problem . To deal with this problem , some language models ( Pennington et al. , 2014 ; Bojanowski et al. , 2017 ) split the words into subtokens and learn the co-occurrence of words . Furthermore , ELMo ( Peters et al. , 2018 ) gathers hidden states and the embeddings of the bidirectional language models to represent contextual embeddings . Around the same time , language models ( Devlin et al. , 2019 ; Liu et al. , 2019 ; Lan et al. , 2020 ) based on transformer ( Vaswani et al. , 2017 ) are proposed to learn the context of entire input sequence with the attention mechanism . In this paper , we split the contextual embeddings into sub-embeddings that contain common context of the embeddings . Relationship Between Tokenizers and Language Models . Language Models usually split a sequence into tokens through tokenizers . Classical tokenizers are designed to divide a sequence into contextual units such as words , characters , or morphs . Those tokenizers are intuitive to implement , but they are easily suffered from OOV problem and need special knowledge to split into morphs . To alleviate the problem mentioned above , some works propose that the tokenizers learn the input corpus to construct concrete vocabularies . Byte-Pair Encoding ( Sennrich et al. , 2016 ) iteratively merge tokens to a larger token , byte-level tokens are used to cover all input cases ( Radford et al. , 2019 ) . On the other hand , Kudo ( 2018 ) ; Kudo & Richardson ( 2018 ) iteratively reduce meaningless tokens from a large token set . We can build up the vocabulary robustly using these tokenizers , but the contextual unit can be different to each token . Typical embedding scheme maps the token to each embedding vector , CharformerTay et al . ( 2021 ) pinpoints this inflexible problem . They propose the gradient-based method to figure out latent subword representation . In contrasts to other static tokenizers , the embedding of token is changed by subword block . In addition , Clark et al . ( 2021 ) alleviates the tokenization framework , the tokenizer is changed to character-level without human knowledge . Xue et al . ( 2021 ) proposes token-free models which encode input sequence with contextual unit , not tokens . In this paper , we propose the embedding structure regardless of tokenizers . Unlike token-free models ( Xue et al. , 2021 ) , the our k-sub-embedding structure needs a tokenizer , but the networks do not suffered anymore from the number of tokens . Since the embeddings act a prominent role in the language models , a series of studies come up with simplifying the embeddings . One of the lightening methods is to compress the embedding weight with matrix factorization . Lan et al . ( 2020 ) conducts matrix factorization to the embedding part , and the high-dimensional embedding vector could be compressed suceessfully to 128-dimensional vector . Our proposed structure also reconstruct the embedding through lookup operation without additional computations . 3 PROPOSED METHODS . We investigate the requirements of each sub-embedding , and suggest ( 1 ) randomly scattered k-subembedding with naive approach ; ( 2 ) clustered k-sub-embedding using contextual knowledge . 3.1 PRELIMINARIES . Modern language models suffer from massive parameters of embeddings because the number of embeddings is associated directly with the diversity of output tokens . Some studies tried to compress the embedding part ; Lan et al . ( 2020 ) lightened the embedding weight via matrix factorization . We horizontally split embedding vectors into sub-embedding vectors which are shared with some other embedding vectors . In the proposed embedding structure , sub-embeddings tend to highly correlate . We suggest the following requirements to verify the modified embeddings perform like the original embedding . 1 . ( Uniqueness of divided embeddings ) Let ei , ej be embedding vectors , and { eki } Kk=1 , { ekj } Kk=1 are their sub-embeddings , ∀i , j ∈ { 1 , . . . , N } ∃k∗ ∈ { 1 . . .K } s.t . ek ∗ i 6= ek ∗ j and i 6= j . 2 . ( Contextual mapped sub-embedding ) If a pair of tokens have similar contextual meanings , their sub-embeddings share more parts than an arbitrary pair . We try to find the function to map the origin embedding vectors into sub-embedding vectors . In this work , we suggest the bijective function F : N → M×· · · × M for converting the embedding to each sub-embedding where N = { 1 , . . . , N } ⊂ N is a set of the embedding index and M = { 1 , . . . , M } ⊂ N is a set of each sub-embedding index . To expand the function to allocate sub-embedding in several ways , we generalize the mapping function using Cartesian product of functions : F ( n ) = ( f1 × f2 ×· · · × fk ) ( n , . . . , n︸ ︷︷ ︸ k ) , ( 1 ) where the function fk : N → M denotes k-th sub-embedding index . Since the cardinality of k-ary Cartesian product is Mk , the embeddings can be covered only N 1 k sub-embeddings . This makes the number of embedding parameters dramatically reduce with log scale . We suppose that the dimension of the embedding d is distributed equally to sub-embedding M = d/k . Then , the number of the embedding parameters N × d is scaled down to k ×M × ( d/k ) = M × d. Finally , the representation of the embedding would be replaced by en = K⊕ k=1 efk ( n ) , where ⊕ is a concatenate operator , and en , efk are the corresponding embedding vector of the embedding index . We suggest Algorithm 1 , 2 by modifying fk that satisfies the requirements . 3.2 RANDOMLY SCATTERED k-SUB-EMBEDDING We construct naively fk to combine sub-embeddings through a Cartesian product which can generate up to Mk embeddings . To ensure generated embeddings to be distinct , the number of each sub-embedding M should be larger than N1/k . Randomly allocated sub-embedding algorithm ( Algorithm 1 ) shows that whole embeddings are generated by repeatedly applying the modulo operation . We set M = dN 1K e to extract the most compressed sub-embeddings . The modulo operation is derived with the perspective of M -base number . Converting to M -base number is bijective , and we can easily set the index of each sub-embedding as each digit of the radix . Algorithm 1 Randomly allocated Sub-Embedding 1 : Input : Number of the embeddings N , embedding dimension d , sub-embedding sets K 2 : M ← dN 1K e . number of each sub-embedding 3 : Initialize k-th M sub-embedding vectors { ekm ∈ R d k } Mm=1 for all k ∈ { 1 , . . . , K } 4 : for n = 1 , 2 , . . . , N do 5 : for k = 1 , 2 , . . . , K do 6 : fk ( n ) ≡ ( n/Mk−1 ) mod Mk , . k-th digit of the M -base number 7 : end for 8 : en = K⊕ k=1 ekfk ( n ) 9 : end for 10 : Output : The combined embedding vectors { en } Nn=1 . 3.3 CLUSTERED k-SUB-EMBEDDING The language models based on transformer learn the whole context of the input sequence , and each embedding vector of a token is mapped to the embedding space reflecting its context . Mikolov et al . ( 2013 ) recognizes that arbitrary two closely mapped word vectors tend to have similar context . If the context of each token is given , we can enhance the allocation heuristic using the contexts . In the case of tokens that have a similar meaning , we assume that the two tokens can be classified with smaller differences . In this case , we can allocate more sub-embeddings to be shared . To estimate the similarity of each pair of the tokens , a pretrained network is used to get each L2 distance between each embedding vector . Assuming that all sub-embeddings are independently initialized randomly , the tokens that share more sub-embedding are expected to have less L2 distance . The method of allocating sub-embedding to similar tokens is based on k-means ( Arthur & Vassilvitskii , 2007 ) algorithm . Each embedding vector is treated as an instance of the k-means algorithm , thus the algorithm is adjusted recursively to each sub-embedding space . The recursive clustering algorithm aims to separate the instances which are allocated in some identical sub-embeddings . According to Algorithm 2 , the mapped sub-embeddings can satisfy the second requirement due to the k-means algorithm based on L2 space .

Summary: This paper proposes an embedding compression method which replaces an $N$ x $d$ embedding table with $k$ distinct $M$ x $(d/k)$ embedding tables (where $M = N^{1/k}$). Specifically, the $d$-dimensional embedding for each token in the vocabulary is replaced by $k$ different sub-embedding vectors $(e_i)_{i=1}^k$, each of dimension $d/k$, where $e_i$ is drawn from the $i^{th}$ codebook. It proposes two ways of selecting these sub-embeddings: 1) arbitrarily assigning the sub-embeddings to each token (based on its index), 2) using k-means to do this assignment procedure.  It shows that very high compression rates (e.g., 99.4-99.9% compression) can be attained using this method, while typically only sacrificing 0%-4% absolute accuracy.