DeFINE: Deep Factorized Input Token Embeddings for Neural Sequence Modeling

1 INTRODUCTION . Neural models for NLP tasks , such as language modeling and machine translation , require large vocabularies for generality ( Chelba et al. , 2013 ; Bahdanau et al. , 2015 ; Luong et al. , 2015 ; Merity et al. , 2017 ) . These models often employ a similar architecture : tokens ( e.g. , words , sub-words , or characters ) , represented as one-hot vectors , are mapped to a dense continuous space ; they are then processed by a context model ; finally , the contextualized representations are mapped back to a vocabulary-sized vector for computing next-token probabilities . A language modeling example is shown in Figure 1a . The mapping in the first and last steps often uses a shared learned lookup table , referred to as an embedding layer , which takes every token in the vocabulary to a fixed m-dimensional vector . One drawback of this approach is that the number of parameters in the embedding layer increases as the vocabulary size grows , limiting us to small values of m over large vocabularies . Researchers have sought to improve the efficiency of the embedding layer by assigning lower frequency tokens smaller dimensional vectors , however , significant parameter reductions come at the cost of performance ( Morin & Bengio , 2005 ; Grave et al. , 2017a ; Baevski & Auli , 2019 ) . In all these approaches , token embedding is approximated with a linear function from tokens to vectors . In this work , we introduce Deep Factorized INput token Embeddings ( DeFINE ) for neural sequence modeling . DeFINE approximates the complicated token embedding function with far fewer parameters compared to standard methods . DeFINE allows for lower-dimensional input and output mappings in sequence models , reducing their computational burden without reducing performance . The representations produced by DeFINE are more powerful than those of other factorization techniques and even standard embedding layers . To accomplish this , DeFINE leverages a hierarchical group transformation ( HGT ) that learns deep representations efficiently and effectively . HGT connects different subsets of the input using sparse and dense connections . To improve the flow of information , DeFINE introduces a new skip-connection that establishes a direct link with the input layer at every level of its hierarchy , allowing gradients to flow back directly to the input via multiple paths . DeFINE replaces standard embedding layers , leaving the rest of the model untouched , and so it can be used with a wide variety of sequence modeling architectures and token-types , including words and sub-words . Figure 1 shows how we incorporate DeFINE with Transformer-XL ( Dai et al. , 2019 ) , a state-of-the-art Transformer-based language model , and the resulting reduction in total parameters . Our experiments show that both LSTM- and Transformer-based sequence models benefit from the use of DeFINE . Furthermore , our experiments with word-level language modeling and sub-word level machine translation tasks show that DeFINE can be used with different token types . On the Wikitext-103 dataset , an LSTM-based language model with DeFINE provides a 9 point improvement over a full capacity model while using half as many parameters . When combined with adaptive input ( Baevski & Auli , 2019 ) and output ( Grave et al. , 2017a ) representations , DeFINE improves the performance by about 3 points across LSTM-based ( see Table 1a ) and Transformer-XL-based ( see Table 2 ) language models with a minimal increase in training parameters . Computation time at inference is unaffected.1 Incorporating DeFINE into the popular AWD-LSTM language model ( Merity et al. , 2018b ) without finetuning results in a test perplexity of 54.2 on the Penn Treebank dataset , outperforming both the original and fine-tuned AWD-LSTM models as well as Transformer-XL and MoS ( Yang et al. , 2018 ) . For machine translation , DeFINE improves the efficiency of a Transformer model ( Vaswani et al. , 2017 ) by 26 % while maintaining translation quality . We provide substantive experiments which detail the impact of our architecture decisions and demonstrate the effectiveness of DeFINE across models of varying capacities . 2 RELATED WORK . Many sequence modeling tasks , including language modeling and machine translation , have a large vocabulary . As a consequence , the majority of a model ’ s parameters are located in the input ( or embedding ) and the output ( or classification ) layers . To reduce the computational load presented by these layers , Press & Wolf ( 2017 ) and Inan et al . ( 2017 ) introduce an effective mechanism called weight-tying that enables learning input and output representations jointly while significantly reducing the number of network parameters . To further reduce the computational load from these layers , factorization-based methods , such as projective embeddings ( Dai et al. , 2019 ) , grouped embeddings ( Chen et al. , 2018 ; Grave et al. , 2017a ; Goodman , 2001 ; Mnih & Hinton , 2009 ; Morin & Bengio , 2005 ) , and slim embeddings ( Li et al. , 2018 ) , have been proposed . Projective embeddings approximate a large embedding matrix with two smaller matrices while grouped embeddings cluster input tokens by frequency and assign different capacities to different clusters using projective embedding methods . We note that projective embeddings is a special case of grouped embeddings when the 1Embeddings learned using DeFINE can be cached , so DeFINE does not increase the computational cost at inference . number of clusters is one . The adaptive input method of Baevski & Auli ( 2019 ) generalizes projective and grouped embedding methods and proposes a factorization method that allows for faster , memory-efficient end-to-end training while providing similar or better benefits compared to existing post-training methods which require a pretrained embedding matrix ( Chen et al. , 2018 ) . Unlike projective and grouped embeddings , Li et al . ( 2018 ) extends group transformation ( Kuchaiev & Ginsburg , 2017 ; Mehta et al. , 2018 ) with the shuffling algorithm of Fisher & Yates ( 1943 ) to factorize these layers . Other techniques such as codebook learning ( Shu & Nakayama , 2017 ; Chen et al. , 2016 ; Acharya et al. , 2019 ) and quantization ( Rastegari et al. , 2016 ; Hubara et al. , 2017 ) can be used to further improve efficiency , especially in terms of storage requirements . DeFINE is orthogonal to these methods ; our empirical results in Section 4 show improved performance compared to these methods alone . Recent advances in sequence modeling , such as Transformers and multi-layer RNNs , demonstrate the power of deep architectures in NLP ( Jozefowicz et al. , 2016 ; Vaswani et al. , 2017 ; Merity et al. , 2018a ) . But while significant attention has been given to modeling the interactions between tokens with deep architectures ( e.g. , ELMo ( Peters et al. , 2018 ) and BERT ( Devlin et al. , 2019 ) ) , contextfree token representations are typically modeled with only corpus statistics ( Pennington et al. , 2014 ) or a single linear transformation ( Mikolov et al. , 2013 ; McCann et al. , 2017 ) . Character-level models ( Kim et al. , 2016 ) also effect deep representations of words as a convolution over characters , however these models often require more capacity to deliver performance comparable to word- or sub-wordlevel models ( Baevski & Auli , 2019 ) . Still , DeFINE can be used to learn deep representations of a variety of token types , including words , characters , or sub-words ( byte-pair encodings ) ( Sennrich et al. , 2015 ) . 3 DEFINE . Token embedding is often treated as simple function of a one-hot vector to a dense continuous space . The embedding layer can thus be thought of as a wide , shallow network consisting of a single linear transformation . At its heart , the function that this network approximates ( call it f ) takes a token from its orthographic form to a representation of those of its morphosyntactic and semantic properties which are relevant for modeling an arbitrary number of contexts in which the token can occur . Most NLP research assumes a simple embedding layer can sufficiently approximate the intractable function f . We hypothesize that , due to the complexity of f , a shallow network would require exceptional capacity to learn a good approximation . Time and data constraints prohibit learning such a high capacity shallow network . We propose , based on recent theoretical results of Liang & Srikant ( 2017 ) ,2 that a deeper network can approximate f with significantly fewer parameters than a shallow network . The validity of this assumption is evidenced by our experimental results in Section 4 . In this work , we introduce DeFINE , an effective way of learning deep token representations in high-dimensional space with a minimum of additional parameters . Our method is based on a MapExpand-Reduce ( MER ) principle , described in Section 3.1 , that first maps an input token to a low dimensional embedding vector , then transforms it to a high-dimensional space using a computationally efficient hierarchical group transformation ( HGT , Section 3.2 ) , which is sketched in Figure 2c . The resultant vector is then transformed to a low-dimensional space . Over the course of these transformations , we make use of a new connectivity pattern that establishes a direct link between the input and output layers ( Figure 3 ) , promoting feature reuse , and improving gradient flow ( Section 3.3 ) . The output layer of DeFINE can then be used in place of a traditional embedding as an input to sequence modeling tasks . We detail the various aspects of the architecture below . 3.1 THE MAP-EXPAND-REDUCE PRINCIPLE ( MER ) The first step in MER , Map , is similar to standard sequence models . Every input token in the vocabulary V is mapped to a fixed dimensional vector ei ∈ Rn×1 . However , in our case , the value of n is small ( say 64 or 128 , compared to typical dimensions of 400 or more ) . The next step , Expand , 2Liang & Srikant ( 2017 ) prove that , for a large class of functions , the number of neurons needed by a shallow network to approximate a function is exponentially larger than the corresponding number of neurons needed by a deep network . We make the assumption that f is in this class of functions . takes ei as an input and applies a hierarchical group transformation ( HGT ) to produce a very highdimensional vector êi ∈ Rk×1 , where k > > n. Unlike a stack of fully connected layers , HGT learns deep representations efficiently from different subsets of the input using sparse and dense connections . The last step , Reduce , projects the vector êi to a lower dimensional space to produce the final embedding vector eo ∈ Rm×1 for a given input token . The dimensions of eo can be matched to contextual representation models , such as LSTMs or Transformers , allowing DeFINE to serve as an input layer for these models . 3.2 HIERARCHICAL GROUP TRANSFORMATION ( HGT ) We introduce a hierarchical group transformation ( HGT ) , sketched in Figure 2c , to learn deep token representations efficiently . HGT comprises of a stack of N layers . At each layer , HGT uses a different number of groups that allows it learn representations from different subsets of input . HGT starts with gmax groups at the first layer and then subsequently decreases the number of groups by a factor of 2 at each level . This hierarchical grouping mechanism sparsifies the connections in fully connected ( or linear ) layers and allows us to learn representations efficiently with fewer parameters . Similar to a stack of fully connected layers , the N -th layer in HGT has access to every input element of the first layer through multiple paths , thereby , allowing it to learn effective representations . Group linear transformations ( GLT ) , originally introduced to improve the efficiency of the LSTM , also sparsify the connections in fully connected layers and significantly reduce computational costs ( Kuchaiev & Ginsburg , 2017 ; Mehta et al. , 2018 ) . However , if we stack multiple GLT layers , the outputs of a certain group are only derived from a small fraction of the input , thus learning weak representations . The hierarchical grouping mechanism in HGT allows the N -th layer to obtain input data from multiple paths , enabling HGT to learn stronger representations . A comparison of different transformations is given in Figure 2 . We can see that HGT is both efficient and has better access to the input . Note that linear and group linear transforms are special cases of HGT when gl = 1 and gl = g ( fixed ) , respectively . To transform ei ∈ Rn×1 to êi ∈ Rk×1 , HGT first samples the space between n and k linearly to construct N intermediate layers of increasing dimensionality . Therefore , the output vector produced by l+1-th layer will have higher dimensionality than the l-th layer . Assume that the linearly spaced vector dimensions are divisible by gmax , we transform ei to êi as follows : êli = { FG ( ei , W l , gl ) , l = 1 FG ( êl−1i , W l , gl ) , 1 < l ≤ N ( 1 ) where gl = max ( b gmax 2l−1 c , 1 ) , Wl are the weights learned at l-th layer , and FG is a group transformation function defined in Mehta et al . ( 2018 ) . Group transformation splits the input into g groups , each of which is processed independently using a linear transformation . The output of these groups are then concatenated to produce final output . See Section A.1 for details . 3.3 DEFINE UNIT The DeFINE unit is composed of HGT transformations that are designed using the MER principle . Though HGT layers are an efficient approximation to computationally expensive fully connected layers , they might impede training as the depth N of the DeFINE unit grows . Residual connections ( He et al. , 2016 ) have proved to be very effective at mitigating this issue , however , such connections are difficult to implement in HGT because the input and output dimensions of each layer are different . To maximize the flow of information and facilitate training with deeper DeFINE units , we introduce a simple new skip-connection that establishes a direct link between any layer in HGT with the input ei . Figure 3 visualizes the DeFINE unit with a depth of two ( N=2 ) . To enable the sparse connections in HGT to have access to the input ei and the output of the previous layer ( êl−1i ) , we chunk the input and the output into gl groups using a split layer . The chunked input and output vectors are then mixed such that the first chunk of the input and the first chunk of the l− 1-th layer ’ s output are put together as the input for the first group transformation in the l-th layer and so on until gl inputs have been constructed . The resultant vector is then fed to the l-th layer . This mechanism promotes input feature reuse efficiently . Additionally , it establishes a direct link with the input ei , allowing gradients to flow back to the input via multiple paths and resulting in improved performance . 3.4 DEFINE FOR SEQUENCE MODELING The DeFINE unit can be easily integrated with any new or existing sequence models . Sequence models typically consist of a stack of an input layer ( embedding or adaptive input layer ) , a contextual model ( e.g. , LSTM or Transformer ) , and a classification layer ( a fully-connected or adaptive softmax ) . Since DeFINE learns deep token representations , we can easily stack it immediately after the input . An example is shown in Figure 1 , where DeFINE is integrated with Transformer-XL , a state-of-the-art language model . DeFINE enables the use of relatively lower dimensions in the input layer , thus reducing network parameters . The input token representations , ei , êi , and eo , that a neural model learns for each token are independent of other tokens . This allows us to create another independent look-up table ( after training a model ) that caches the mapping between the input token and the output of the DeFINE unit ( eo ) , resulting in a mechanism that allows to skip the computations of the DeFINE unit at inference time .

Summary: This paper describes an approach to learn word embedding functions more efficiently and with fewer parameters. This is done by replacing the embedding lookup function which is typical in NLP tasks such as language modeling and machine translation with a hierarchical embedding model. This allows for a low dimensional embedding layer, reducing total parameters and training time. A novel skip-connections architecture is introduced as a part of the "embedding generation model". Experiments are conducted for language modeling and machine translation tasks and performance improvements are observed with a reduction in parameters and lesser training time.