Anchor & Transform: Learning Sparse Embeddings for Large Vocabularies

1 INTRODUCTION . Most machine learning models , including neural networks , operate on vector spaces . Therefore , when working with discrete objects such as text , we must define a method of converting objects into vectors . The standard way to map objects to continuous representations involves : 1 ) defining the vocabulary V = { v1 , ... , v∣V ∣ } as the set of all objects , and 2 ) learning a ∣V ∣ × d embedding matrix that defines a d dimensional continuous representation for each object . This method has two main shortcomings . Firstly , when ∣V ∣ is large ( e.g. , million of words/users/URLs ) , this embedding matrix does not scale elegantly and may constitute up to 80 % of all trainable parameters ( Jozefowicz et al. , 2016 ) . Secondly , despite being discrete , these objects usually have underlying structures such as natural groupings and similarities among them . Assigning each object to an individual vector assumes independence and foregoes opportunities for statistical strength sharing . As a result , there has been a large amount of interest in learning sparse interdependent representations for large vocabularies rather than the full embedding matrix for cheaper training , storage , and inference . In this paper , we propose a simple method to learn sparse representations that uses a global set of vectors , which we call the anchors , and expresses the embeddings of discrete objects as a sparse linear combination of these anchors , as shown in Figure 1 . One can consider these anchors to represent latent topics or concepts . Therefore , we call the resulting method ANCHOR & TRANSFORM ( ANT ) . The approach is reminiscent of low-rank and sparse coding approaches , however , surprisingly in the literature these methods were not elegantly integrated with deep networks . Competitive attempts are often complex ( e.g. , optimized with RL ( Joglekar et al. , 2019 ) ) , involve multiple training stages ( Ginart et al. , 2019 ; Liu et al. , 2017 ) , or require post-processing ( Svenstrup et al. , 2017 ; Guo et al. , 2017 ; Aharon et al. , 2006 ; Awasthi & Vijayaraghavan , 2018 ) . We derive a simple optimization objective which learns these anchors and sparse transformations in an end-to-end manner . ANT is ∗work done during an internship at Google . scalable , flexible , and allows the user flexibility in defining these anchors and adding more constraints on the transformations , possibly in a domain/task specific manner . We find that our proposed method demonstrates stronger performance with fewer parameters ( up to 40× compression ) on multiple tasks ( text classification , language modeling , and recommendation ) as compared to existing baselines . We further provide a statistical interpretation of our algorithm as a Bayesian nonparametric ( BNP ) prior for neural embeddings that encourages sparsity and leverages natural groupings among objects . Specifically , we show its equivalence to Indian Buffet Process ( IBP ; Griffiths & Ghahramani ( 2005 ) ) prior for embedding matrices . While such BNP priors have proven to be a flexible tools in graphical models to encourage hierarchies ( Teh & Jordan , 2010 ) , sparsity ( Knowles & Ghahramani , 2011 ) , and other structural constraints ( Roy et al. , 2016 ) , these inference methods are usually complex , hand designed for each setup , and non-differentiable . Our proposed method opens the door towards integrating priors ( e.g. , IBP ) with neural representation learning . These theoretical connections leads to practical insights - by asymptotically analyzing the likelihood of our model in the small variance limit using Small Variance Asymptotics ( SVA ; Roweis ( 1998 ) ) , we obtain a natural extension , NBANT , that automatically learns the optimal number of anchors to achieve a balance between performance and compression instead of having to tune it as a hyperparameter . 2 RELATED WORK . Prior work in learning sparse embeddings of discrete structures falls into three categories : Matrix compression techniques such as low rank approximations ( Acharya et al. , 2019 ; Grachev et al. , 2019 ; Markovsky , 2011 ) , quantizing ( Han et al. , 2016 ) , pruning ( Anwar et al. , 2017 ; Dong et al. , 2017 ; Wen et al. , 2016 ) , or hashing ( Chen et al. , 2015 ; Guo et al. , 2017 ; Qi et al. , 2017 ) have been applied to embedding matrices . However , it is not trivial to learn sparse low-rank representations of large matrices , especially in conjunction with neural networks . To the best of our knowledge , we are the first to present the integration of sparse low-rank representations , their non-parametric extension , and demonstrate its effectiveness on many tasks in balancing the tradeoffs between performance & sparsity . We also outperform many baselines based on low-rank compression ( Grachev et al. , 2019 ) , sparse coding ( Chen et al. , 2016b ) , and pruning ( Liu et al. , 2017 ) . Reducing representation size : These methods reduce the dimension d for different objects . Chen et al . ( 2016a ) divides the embedding into buckets which are assigned to objects in order of importance , Joglekar et al . ( 2019 ) learns d by solving a discrete optimization problem with RL , and Baevski & Auli ( 2019 ) reduces dimensions for rarer words . These methods resort to RL or are difficult to tune with many hyperparameters . Each object is also modeled independently without information sharing . Task specific methods include learning embeddings of only common words for language modeling ( Chen et al. , 2016b ; Luong et al. , 2015 ) , and vocabulary selection for text classification ( Chen et al. , 2019 ) . Other methods reconstruct pre-trained embeddings using codebook learning ( Chen et al. , 2018 ; Shu & Nakayama , 2018 ) or low rank tensors ( Sedov & Yang , 2018 ) . However , these methods can not work for general tasks . For example , methods that only model a subset of objects can not be used for retrieval because it would never retrieve the dropped objects . Rare objects might be highly relevant to a few users so it might not be ideal to completely ignore them . Similarly , task-specific methods such as subword ( Bojanowski et al. , 2017 ) and wordpiece ( Wu et al. , 2016 ) embeddings , while useful for text , do not generalize to general applications such as item and query retrieval . 3 ANCHOR & TRANSFORM . Suppose we are presented with data X ∈ V N , Y ∈ RN×c drawn from some joint distribution p ( x , y ) , where the support of x is over a discrete set V ( the vocabulary ) and N is the size of the training set . The entries in Y can be either discrete ( classification ) or continuous ( regression ) . The goal is to learn a d-dimensional representation { e1 , ... , e∣V ∣ } for each object by learning an embedding matrix E ∈ R∣V ∣×d where row i is the representation ei of object i . A model fθ with parameters θ is then used to predict y , i.e. , ŷi = fθ ( xi ; E ) = fθ ( E [ xi ] ) . At a high level , to encourage statistical sharing between objects , we assume that the embedding of each object is obtained by linearly superimposing a small set of anchor objects . For example , when the objects considered are words , the anchors may represent latent abstract concepts ( of unknown cardinality ) and each word is a weighted mixture of different concepts . More generally , the model assumes that there are some unknown number of anchors , A = { a1 , ... , a∣A∣ } . The embedding ei for object i is generated by first choosing whether the object possesses each anchor ak ∈ Rd . The selected anchors then each contribute some weight to the representation of object i . Therefore , instead of learning the large embedding matrix E directly , ANT consists of two components : Algorithm 1 ANCHOR & TRANSFORM algorithm for learning sparse representations of discrete objects . ANCHOR & TRANSFORM : 1 : Anchor : initialize anchor embeddings A . 2 : Transform : initialize T as a sparse matrix . 3 : Optionally + domain info : initialize domain sparsity ma- trix S ( G ) as a sparse matrix ( see Appendix F ) . 4 : for each batch ( X , Y ) do 5 : Compute loss L = ∑iDφ ( yi , fθ ( xi ; TA ) ) 6 : A , T , θ = UPDATE ( ∇L , η ) . 7 : T = max { ( T − ηλ2 ) ⊙ S ( G ) +T⊙ ( 1 − S ( G ) ) ,0 } . 8 : end for 9 : return anchor embeddings A and transformations T. 1 ) ANCHOR : Learn embeddings A ∈ R∣A∣×d of a small set of anchor objects A = { a1 , ... , a∣A∣ } , ∣A∣ < < ∣V ∣ that are representative of all discrete objects . 2 ) TRANSFORM : Learn a sparse transformation T from A to E. Each of the discrete objects is induced by some transformation from ( a few ) anchor objects . To ensure sparsity , we want nnz ( T ) < < ∣V ∣ × d. A and T are trained end-to-end for task specific representations . To enforce sparsity , we use an ` 1 penalty on T and constrain its domain to be non-negative to reduce redundancy in transformations ( positive and negative entries canceling out ) . min T≥0 , A , θ ∑ i Dφ ( yi , fθ ( xi ; TA ) ) + λ2∥T∥1 , ( 1 ) where Dφ is a suitable Bregman divergence between predicted and true labels , and ∥T∥1 denotes the sum of absolute values . Most deep learning frameworks directly use subgradient descent to solve eq ( 1 ) , but unfortunately , such an approach will not yield sparsity . Instead , we perform optimization by proximal gradient descent ( rather than approximate subgradient methods which have poorer convergence around non-smooth regions , e.g. , sparse regions ) to ensure exact zero entries in T : A t+1 , T t+1 , θ t+1 = UPDATE ( ∇∑ i Dφ ( yi , fθ ( xi ; T t A t ) ) , η ) , ( 2 ) T t+1 = PROXηλ2 ( T t+1 ) = max ( T t+1 − ηλ2 , 0 ) , ( 3 ) where η is the learning rate , and UPDATE is a gradient update rule ( e.g. , SGD ( Lecun et al. , 1998 ) , ADAM ( Kingma & Ba , 2015 ) , YOGI ( Zaheer et al. , 2018 ) ) . PROXηλ2 is a composition of two proximal operators : 1 ) soft-thresholding ( Beck & Teboulle , 2009 ) at ηλ2 which results from subgradient descent on λ2∥T∥1 , and 2 ) max ( ⋅,0 ) due to the non-negative domain for T. We implement this proximal operator on top of the YOGI optimizer for our experiments . Together , equations ( 2 ) and ( 3 ) give us an iterative process for end-to-end learning of A and T along with θ for specific tasks ( Algorithm 1 ) . T is implemented as a sparse matrix by only storing its non-zero entries and indices . Since nnz ( T ) < < ∣V ∣× d , this makes storage of T extremely efficient as compared to traditional approaches of computing the entire ∣V ∣×d embedding matrix . We also provide implementation tips to further speedup training and ways to incorporate ANT with existing speedup techniques like softmax sampling ( Mikolov et al. , 2013 ) or noise-contrastive estimation ( Mnih & Teh , 2012 ) in Appendix H. After training , we only store ∣A∣ × d + nnz ( T ) < < ∣V ∣ × d entries that define the complete embedding matrix , thereby using fewer parameters than the traditional ∣V ∣ × d matrix . General purpose matrix compression techniques such as hashing ( Qi et al. , 2017 ) , pruning ( Dong et al. , 2017 ) , and quantizing ( Han et al. , 2016 ) are compatible with our method : the matrices A and nnz ( T ) can be further compressed and stored . We first discuss practical methods for anchor selection ( §3.1 ) . In Appendix F we describe several ways to incorporate domain knowledge into the anchor selection and transform process . We also provide a statistical interpretation of ANT as a sparsity promoting generative process using an IBP prior and derive approximate inference based on SVA ( §3.2 ) . This gives rise to a nonparametric version of ANT that automatically learns the optimal number of anchors .

Summary: In this paper, the authors proposed a method to learn efficient representations of discrete tokens. They took a two step approach: in step 1, they learn "full fledged" embeddings for a subset of anchor tokens. In step 2, they learn a sparse matrix that is used to relate all tokens to the set of chosen anchors. This two-step approach reduced the overall number of parameters. The sparse matrix T can also encode domain knowledge (e.g. knowledge graphs). In the experiment section, the authors showed that their approach has good performance on several language tasks, with far fewer parameters.