A NCHOR & T RANSFORM : L EARNING S PARSE E MBEDDINGS FOR L ARGE V OCABULARIES
A BSTRACT
Learning continuous representations of discrete objects such as text, users, movies, and URLs lies at the heart of many applications including language and user modeling
When using discrete objects as input to neural networks, we often ignore the underlying structures (e.g natural groupings and similarities) and embed the objects independently into individual vectors
As a result, existing methods do not scale to large vocabulary sizes.
In this paper , we design a simple and efﬁcient embedding algorithm that learns a small set of anchor embeddings and a sparse transformation matrix.
We call our method A NCHOR & T RANSFORM (ANT) as the embeddings of discrete objects are a sparse linear combination of the anchors, weighted according to the transformation matrix.
ANT is scalable, ﬂexible, and end-to-end trainable.
We further provide a statistical interpretation of our algorithm as a Bayesian nonparametric prior for embeddings that encourages sparsity and leverages natural groupings among objects.
By deriving an approximate inference algorithm based on Small Variance Asymptotics, we obtain a natural extension that automatically learns the optimal number of anchors instead of having to tune it as a hyperparameter.
On text classiﬁcation, language modeling and movie recommendation benchmarks, we show that ANT is particularly suitable for large vocabulary sizes and demonstrates stronger performance with fewer parameters (up to × compression) as compared to existing compression baselines
I NTRODUCTION
Most machine learning models, including neural networks, operate on vector spaces.
Therefore, when working with discrete objects such as text, we must deﬁne a method of converting objects into vectors.
The standard way to map objects to continuous representations involves: 1) deﬁning the vocabulary
V = { v 1 ,...,v ∣ V ∣ } as the set of all objects, and 2) learning a ∣ V ∣ ×
d embedding matrix that deﬁnes 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
A NCHOR & T RANSFORM (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 scalable, ﬂexible, and allows the user ﬂexibility in deﬁning these anchors and adding more constraints on the transformations, possibly in a domain/task speciﬁc manner.
We ﬁnd that our proposed method
Conﬁdential + Proprietary “good” “the” pretrained space e.g. GloVe/co-occurrence
A d |A|
Frequency + Clustering
Conﬁdential + Proprietary d |V| E =
A d |V| |A| |A| =
eg. frequency, clustering
T
E x
[Table]
2. Transform
[Table]
+wcNE6easjqNRaxbITFMcMXqwEGwVqIZkaFgzXB4M6k375k2PFY1GCUskKSveMQpAWt13eMOsAfQMiuX8HjKhmY1TZSJYi3HXbfolbyp8DL4cyiiuapd96vTi2kqmQIqiDFt30sgyIgGTgUbFzqpYQmhQ9JnbYuKSGaCbHrHGJ9ap4ftYvsU4Kn7eyIj0piRDG2nJDAwi7WJ+V+tnUJ0HWRcJSkwRWeLolRgiPEkFNzjmlEQIwuEam
+O0Ak6Qz66QhV0i6qojih6RM/oFb05T86L8+58zFpzznzmEP2R8/kD6eSZ8A==</latexit> demonstrates stronger performance with fewer parameters (up to 40 × compression) on multiple tasks (text classiﬁcation, 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.
Speciﬁcally, we show its equivalence to Indian Buffet Process (IBP; Grifﬁths & Ghahramani (2005))
prior for embedding matrices.
While such BNP priors have proven to be a ﬂexible 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, NB ANT, 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.
R ELATED W ORK
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 ﬁrst 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 difﬁcult to tune with many hyperparameters.
Each object is also modeled independently without information sharing.
Task speciﬁc methods include learning embeddings of only common words for language modeling (Chen et al., 2016b;
Luong et al., 2015), and vocabulary selection for text classiﬁcation (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 cannot work for general tasks.
For example, methods that only model a subset of objects cannot 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-speciﬁc 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.
A NCHOR & T RANSFORM
Suppose we are presented with data X ∈ V N , Y ∈ R N × 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 (classiﬁcation) or continuous (regression).
The goal is to learn a d -dimensional representation { e 1 ,..., e ∣ V ∣ } for each object by learning an embedding matrix
E ∈ R ∣ V ∣×
d where row i is the representation e i of object i .
A model f θ with parameters θ is then used to predict y , i.e., ˆ y i
=
f θ ( x i ; E ) =
f θ ( E [ x i ]) .
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 = { a 1 ,..., a ∣ A ∣ } .
The embedding e i for object i is generated by ﬁrst choosing whether the object possesses each anchor a k ∈ R d .
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:
[Equation]
[Equation]
1) A NCHOR : Learn embeddings A ∈ R ∣ A ∣× d of a small set of anchor
objects A = { a 1 ,..., a ∣ A ∣ } , ∣ A ∣ <<
∣ V ∣ that are representative of all discrete objects.
2) T RANSFORM : 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 speciﬁc representations.
To enforce sparsity, we use an (cid:96) 1 penalty on T and constrain its domain to be non-negative to reduce redundancy in transformations (positive and negative entries canceling out).
[Equation]
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 :
[Equation]
where η is the learning rate, and U PDATE is a gradient update rule (e.g., SGD (Lecun et al., 1998),
A DAM (Kingma & Ba, 2015), Y OGI (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 Y OGI 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 speciﬁc 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 efﬁcient 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 deﬁne 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 ﬁrst 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.
3. A NCHOR : S ELECTING THE A NCHORS A
Inspired by research integrating initialization strategies based on clustering (Teh et al., 2007) and
Coresets (Bachem et al., 2015) with Bayesian nonparametrics, we describe several practical methods to select anchor objects that are most representative of all objects (refer to Appendix D for a comparison of initialization strategies.).
Frequency and TF-IDF:
For tasks where frequency or TF-IDF (Ramos, 1999) are useful for prediction, the objects can simply be sorted by frequency and the most common objects selected as the anchor points.
While this might make sense for tasks such as language modeling (Luong et al., 2015; Chen et al., 2016b), choosing the most frequent objects might not cover rare objects that are not well represented by common anchors.
[Figure]
Clustering:
To ensure that all objects are close to some anchor, we use k -means ++ initialization (Arthur & Vassilvitskii, 2007).
Given a feature space representative of the relationships between objects, such as Glove (Pennington et al., 2014) for words or a co-occurrence matrix (Haralick et al., 1973) for more general objects, k -means ++ initialization picks cluster centers to span the entire space.
This can augment other strategies, such as initializing anchors using frequency followed by clustering to complete remaining anchors (see Figure 2).
Random basis vectors: Initialize A to a set of random basis vectors.
This simple yet powerful method captures the case where we have less knowledge about the objects (i.e., without access to any pretrained representation/similarity space).
3. S TATISTICAL I NTERPRETATION AS A B AYESIAN N ONPARAMETRIC P RIOR
To provide a statistical interpretation of ANT, we ﬁrst analyze a generative process for discrete representations that is consistent with our algorithm.
Given a set of anchors, A = { a 1 ,..., a ∣ A ∣ } , we use a binary latent variable z ik ∈ { 0 , 1 } to indicate whether object i possesses anchor k and a positive latent variable w ik ∈ R ≥ 0 to denote the weight that anchor k contributes towards object i .
Therefore, the representation e i is given by e i
=
∑ k w ik z ik a k .
Ideally, we want the vector z i to be sparse for efﬁcient learning and storage.
More formally, suppose there are K ∶= ∣ A ∣ anchors, then:
• Z ∈ R ∣ V ∣× K ∼ IBP ( a,b ) ; A ∈ R K × d ∼ P ( A ) =
N( 0 , 1 ) ; W ∈ R ∣ V ∣× K ∼ P ( W ) =
Exp ( 1 ) • for i = 1 , ⋯ ,N - ˆ y i
=
f θ ( x i ; ( Z
[Equation]
In this generative process, the selection matrix Z follows a two-parameter Indian Buffet Process (IBP; Grifﬁths & Ghahramani (2005))
prior (Ghahramani et al., 2007).
Not only does this BNP prior allow for a potentially inﬁnite number of anchors, but it also encourages each object to only select a small subset of anchors, resulting in a sparse z i (see Appendix A for details).
We place a standard
Gaussian prior on the continuous anchors embeddings a k and an exponential prior on the weights
W which give the actual non-negative transformation weights for the non-zero entries deﬁned in
Z .
E =
( Z ○ W )
A is the ﬁnal embedding learnt by our model which represents a d -dimensional continuous representation { e 1 ,..., e ∣ V ∣ } for each discrete object where row i is the representation e i of object i .
Finally, a neural model f θ with parameters θ is used to predict y i given the embedded representations, i.e., ˆ y i
=
f θ ( x i ; ( Z ○ W ) A ) =
f θ (( Z ○ W )
A [ x i ]) .
Likelihood Model/Loss: We assume that the ﬁnal emission model y i ∣
ˆ y i belongs to the exponential family.
Since exponential family distributions have a corresponding Bregman divergence (Banerjee et al.
(2005); see Appendix C for examples), we choose D φ ( y i , ˆ y i ) as the corresponding Bregman divergence between predicted and true labels.
Appropriate choices for D φ recover cross-entropy and
MSE losses.
b φ ( y i ) does not depend on any learnable parameter or variable and can be ignored.
Joint likelihood: Under the generative model as deﬁned above, the joint likelihood is given by:
[Equation]
However, calculating the posterior or MAP estimate is hard, especially due to the presence of the non-linear deep network f θ .
Approximate inference methods such as MCMC, variational inference, or probabilistic programming would be computationally and statistically inefﬁcient since it would involve sampling, evaluating, or training the model multiple times.
To tackle this problem, we perform approximate inference via Small Variance Asymptotics (SVA), which captures the beneﬁts of rich latent-variable models while providing a framework for scalable optimization (Broderick et al., 2013a;
Jiang et al., 2012; Roychowdhury et al., 2013).
Approximate Inference via SVA:
To use SVA, we introduce a scaling variable β and shrink the variance of the emission probability by taking β → ∞ .
The scaled probability emission becomes
[Equation]
Following Broderick et al.
(2013a), we modulate the number of features in the large- β limit by choosing constants λ 1 > λ 2 > 0 and setting the IBP hyperparameters a = exp (− βλ 1 ) and b = exp ( βλ 2 ) .
This prevents a limiting objective function that favors a trivial cluster assignment (every data point assigned to its own separate feature).
Maximizing the asymptotic joint likelihood (after taking limits, i.e., lim β →∞ 1 β log p ( Y , Z , W , A ∣ X ) ) results in the following objective function:
[Equation]
where we have combined the variables Z and W with their constraints into one variable T .
The exponential prior for W results in a non-negative domain for T .
Please refer to Appendix B for derivations.
Note that eq (5) suggests a natural objective function in learning representations that minimize the prediction loss D φ ( y i ,f θ ( x i ; TA )) while ensuring sparsity of T as measured by the (cid:96)
-norm and using as few anchors as possible ( K ).
Therefore, optimizing eq (5) gives rise to a nonparametric version of ANT, which we call NB ANT, that automatically learns the optimal number of anchors.
To perform optimization over the number of anchors, our algorithm starts with a small ∣
A ∣ = 10 and either adds anchors (i.e., adding a new row to A and a new column to T ) or deletes anchors to minimize eq (5) at every epoch depending on the trend of the objective evaluated on validation set.
We outline the exact algorithm in Appendix G along with more implementation details.
Analogously, we can derive the ﬁnite case objective function for a ﬁxed number of anchors K :
[Equation]
which, together with a (cid:96) 1 penalty on T as a convex relaxation for the (cid:96) 0 penalty, recovers the objective function in eq (1).
The solution for this ﬁnite version along with K yields the Pareto front.
Different values of λ 1 in eq (5) can be used for model selection along the front as elucidated in Appendix L.
E XPERIMENTS
To evaluate ANT, we experiment on text classiﬁcation, language modeling, and movie recommendation tasks.
Experimental details are in Appendix J and full results are in Appendix K.
4. T EXT C LASSIFICATION
Setup: We follow the setting in Chen et al.
(2019) with four datasets: AG-News ( V = 62 K) (Zhang et al., 2015), DBPedia ( V = 563 K) (Lehmann et al., 2015), Sogou-News ( V = 254 K) (Zhang et al., 2015), and Yelp-review ( V = 253 K)
(Zhang et al., 2015).
We use a CNN for classiﬁcation (Kim, 2014).
ANT is used to replace the input embedding and domain knowledge is derived from WordNet and co-occurrence in the training set.
We record test accuracy and number of parameters used in the embedding only.
For ANT, num params is computed as ∣ A
∣ × d
+ nnz ( T ) .
[Table] On top of the CNN , we compare to the following compression approaches.
Vocabulary selection methods: 1) F REQUENCY where only embeddings for most frequent words are learnt (Chen et al., 2016b;
Luong et al., 2015), 2) TF-IDF which only learns embeddings for words with high
TF-IDF score (Ramos, 1999), 3) GL (group lasso) which aims to ﬁnd underlying sparse structures in the embedding matrix via row-wise (cid:96) 2 regularization (Liu et al., 2015; Park et al., 2016; Wen et al., 2016), 4) VVD (variational vocabulary dropout) which performs variational dropout for vocabulary selection (Chen et al., 2019).
We also compare to 5) S PARSE VD (sparse variational dropout) which performs variational dropout on all parameters (Chirkova et al., 2018), 6) S PARSE VD-V OC which uses multiplicative weights for vocabulary sparsiﬁcation (Chirkova et al., 2018), and 7) a S PARSE
C ODE model that learns a sparse code to reconstruct pretrained word representations (Chen et al., 2016b).
All CNN architectures are the same for all baselines with details in Appendix J.1.
Results on AG-News are in Table 1 and results for other datasets are in Appendix K.1.
We observe that restricting T ≥ 0 using an exponential prior is important in reducing redundancy in the entries.
Domain knowledge from WordNet and co-occurrence also succeeded in reducing the total (non-zero) embedding parameters to 0 .
40 M, a compression of 40 × and outperforming the existing approaches.
4. L ANGUAGE M ODELING
Setup: We perform experiments on word-level Penn Treebank (PTB) ( V = 10 K) (Marcus et al., 1993) and
WikiText-103 ( V = 267 K) (Merity et al., 2017) with LSTM (Hochreiter & Schmidhuber, 1997) and AWD-LSTM (Merity et al., 2018).
We use ANT as the input embedding tied to the output embedding.
Domain knowledge is derived from WordNet and co-occurrence on the training set.
We record the test perplexity and the number of (non-zero) embedding parameters.
Baselines: We compare to S PARSE VD and S PARSE VD-V OC , as well as low-rank ( LR ) and tensortrain ( TT ) model compression techniques (Grachev et al., 2019).
Note that the application of variational vocabulary selection to language modeling with tied weights is non-trivial since one is unable to predict next words when words are dynamically dropped out.
We also compare against methods that compress the trained embedding matrix as a post-processing step before evaluation:
P OST -S PARSE H ASH (post-processing using sparse hashing)
(Guo et al., 2017) and P OST -S PARSE
H ASH + k -SVD (Awasthi & Vijayaraghavan, 2018; Guo et al., 2017) which uses k -SVD (which is the basis of dictionary learning/sparse coding) (Aharon et al., 2006) to solve for a sparse embedding matrix, instead of adhoc-projection in (Guo et al., 2017).
Comparing to these post-processing methods demonstrates that end-to-end training of sparse embeddings is superior to post-compression.
Results:
On PTB (Table 2), we improve the perplexity and compression as compared to previously proposed methods.
We observe that sparsity is important: baseline methods that only perform lowerrank compression with dense factors (e.g., LR LSTM) tend to suffer in performance and use many parameters, while ANT retains performance with much better compression.
ANT also outperforms post-processing methods (P OST -S PARSE H ASH ), we hypothesize this is because these post-processing methods accumulate errors in both language modeling as well as embedding reconstruction.
Using an anchor size of 500 / 1 , 000 reaches a good perplexity/compression trade-off: we reach withinpoints perplexity with 5 × reduction in parameters and within 7 points perplexity with 10 × reduction.
Using AWD-LSTM, ANT with 1 , 000 dynamic basis vectors is able to compress parameters by 10 × while achieving 72 . 0 perplexity.
Incorporating domain knowledge allows us to further compress the parameters by another 10 × and achieve 70 .
perplexity, which results in 100 × total compression.
[Table] On WikiText- [Table] we train using sampled softmax (Bengio & Senecal, 2008) (due to large vocabulary) for 500 , 000 steps.
To best of our knowledge [Table] we could not ﬁnd literature on compressing language models on WikiText-103.
We tried general compression techniques like low rank tensor and tensor train factorization (Grachev et al., 2019) [Table] but these did not scale.
As an alternative
[Table]
we consider a
H ASH E MBED baseline that retains the frequent k words and hashes the remaining words into 1 ,
OOV buckets (Svenstrup et al., 2017).
We vary k ∈ { 1 × 10 5 , 5 × 10 4 , 1 × 10 4 } (details in Appendix J.3).
From Table 2 (bottom)
[Table]
we reach within 3 perplexity with ∼ 16 × reduction in parameters and withinperplexity with ∼ 80 × reduction, outperforming the frequency and hashing baselines.
We observe that
ANT’s improvement over post-compression methods (P OST -S PARSE H ASH ) is larger on WikiText than PTB, suggesting that ANT is particularly suitable for large vocabularies.
4. R ECOMMENDER S YSTEMS
Setup: We perform experiments on both movie and product recommendation tasks.
For movie recommendations, we follow Ginart et al.
and we experiment on MovieLens 25 M (Harper & Konstan, 2015) with 126 K users and 59 K movies.
We also present results for MovieLens 1M in Appendix K.3.
On product recommendation, we show that ANT scales to Amazon Product reviews (Ni et al., 2019), the largest existing dataset for recommender systems with 233 M reviews spanning 43 .
5 M users and 15 .
2 M products. Following Wan et al. , we ensured that the users and products in the test set have appeared in the training data for generalization. Baselines: We compare to a baseline Matrix Factorization ( MF ) model (Koren et al., 2009) with full embedding matrices for movies and users and to Mixed Dimension ( M IX D IM ) embeddings (Ginart et al., 2019), a compression technique that assigns different dimension to different users/items based on popularity.
We also compare to S PARSE CBOW (Sun et al., 2016) which learns sparse E by placing an (cid:96) 1 penalty over all entries of E and optimizing using online subgradient descent, and S LIMMING
(Liu et al., 2017), which performs subgradient descent before pruning small weights by setting them to 0 .
Such methods learn embeddings for objects independently without statistical strength sharing among related objects.
We also test NB ANT using the algorithm derived from the Bayesian nonparametric interpretation of ANT.
Results: From Table 3, ANT outperforms standard matrix factorization and dense mixed dimensional embeddings for performance and compression.
NB ANT is also able to automatically select an optimal [Table] [Table] number of anchors ( 6 / 8 ) to achieve solutions along the performance-compression Pareto front.
In Figure 3, we plot the value of eq (5) across values of ∣ A ∣ after a comprehensive hyperparameter sweep on ANT across 1000 settings.
In comparison, NB ANT optimizes ∣ A ∣ and reaches a good value of eq (5) in a single run without having to tune ∣ A ∣ as a hyperparameter, thereby achieving best balance between performance and compression.
Please refer to Appendix K.3 for more results and discussion on NB ANT.
For product recommendation, we ﬁrst experiment on a commonly used subset of the data, Amazon
Electronics (with 9 .
84 M users and 0 .
76 M products), to ensure that our results match published baselines (Wan et al., 2020), before scaling our experiment to the entire dataset. From Table 4, we ﬁnd that ANT compresses embeddings by 25 × on Amazon Electronics while maintaining performance, and 10 × on the full Amazon reviews dataset.
Online NB ANT: Since NB ANT automatically grows/contracts ∣ A ∣ during training, we can further extend NB ANT to an online version that sees a stream of batches without revisiting previous ones (Bryant & Sudderth, 2012).
We treat each batch as a new set of data coming in and train on that batch until convergence, modify ∣ A ∣ as in Algorithm 2, before moving onto the next batch.
In this signiﬁcantly more challenging online setting, NB ANT is still able to learn well and achieve a MSE of 0 .
875 with 1 .
25 M non zero parameters.
Interestingly this online version of NB ANT settled on a similar range of ﬁnal user ( 8 ) and item ( 8 ) anchors as compared to the non-online version (see
Table 3), which conﬁrms the robustness of NB ANT in ﬁnding relevant anchors automatically.
In
Appendix K.3 we discuss more observations around online NB ANT including ways of learning ∣ A ∣ .
4. D ISCUSSION AND O BSERVATIONS
Here we list some general observations regarding the importance of various design decisions in ANT:
1) Sparsity is important: Baselines that compress with dense factors (e.g., LR, TT) suffer in performance while using many parameters, while ANT retains performance with better compression.
Choice of A : We provide results on more clustering initializations in Appendix D. In general, performance is robust w.r.t.
choice of A .
While frequency and clustering work better, using a dynamic basis also performs well. Thus, it
is beneﬁcial to use any extra information about the discrete objects (
e.
g.,
domain
knowledge or having a good representation space like GloVe to perform clustering)
[Table]
3) Anchors and sparse transformations learned: We visualize the important transformations (large entries) learned between anchors and non-anchors in Table 5.
Left, we show the most associated non-anchors for a given anchor word and ﬁnd that the induced non-anchors are highly plausible: stock accurately contributes to bonds, certiﬁcates, securities , and so on.
Right, we show the largest (non-anchor, anchor) pairs learned, where we ﬁnd related concepts such as (billion, trillion) and (government, administration) .
On MovieLens, for each anchor, we sort the movies according to the magnitude of their transformation coefﬁcients which automatically discovers movie clusters based on underlying genres .
We obtain a genre purity ratio of 61 .
7% by comparing automatically discovered movie clusters with the true genre tags provided in MovieLens.
4) Zero transformations learned: For MovieLens, we ﬁnd that ANT assigns 2673 out ofmovies to an entire zero row, of which 84 % only had 1 rating (i.e., very rare movies).
Therefore, compression automatically discovers very rare objects ( 1 labeled point).
On WikiText-103, rare words (e.g., Anarky, Perl, Voorhis, Gaudí, Lat, Bottomley, Nescopeck ) are also automatically assigned zero rows when performing high compression ( 54 . 2 ppl with 0 .
4 M params).
Certain rare words that might be predictive, however, are assigned non-zero rows in T , such as: sociologists, deadlines, indestructible, causeways, outsourced, glacially, heartening, unchallenging, roughest .
5) Choice of λ 1 ,λ 2 : Tuning λ 1 allows us to perform model selection by controlling the trade-off between ∣ A ∣ (model complexity) and performance.
By applying eq (5) on our trained models in
Table 2, choosing a small λ 1 = 2 × 10 − 5 prefers more anchors ( ∣ A ∣ = 1 , 000 ) and better performance (ppl = 79 . 4 ), while a larger λ 1 = 1 × 10 − 1 selects fewer anchors ( ∣ A ∣ = 100 ) with a compromise in performance (ppl
= 106 . 6 ).
Tuning λ 2 allows us to control the tradeoff between sparsity and performance (see details in Appendix L).
6) Convergence:
In Figure 4, we plot the empirical convergence of validation loss across epochs.
ANT converges as fast as the (non-sparse) MF baseline, and faster than compression baselines MixDim (Ginart et al., 2019) and Sparse CBOW (Sun et al., 2016).
ANT also converges to the best validation loss.
7) Scalability:
In addition to fast convergence, ANT also works effectively on large datasets such as Movielens 25 M ( 162 K users, 59 K movies, 25 M examples) and WikiText- 103 ( 267 K unique words, 103 M tokens).
For each epoch on Movielens 25 M, standard MF takes 165 s on a GTX 980 Ti GPU while
ANT takes 176 s for ∣ A ∣ = 5 and 180 s for ∣ A ∣ = 20 .
ANT also scales to the largest recommendation dataset, Amazon reviews, with 25 M users and 9 M products.
C ONCLUSION
This paper presented A NCHOR & T RANSFORM to learn sparse embeddings of large vocabularies using a small set of anchor embeddings and a sparse transformation from anchors to all objects.
We also showed a statistical interpretation via integrating IBP priors with neural representation learning.
Asymptotic analysis of the likelihood using SVA yields an extension that automatically learns the optimal number of anchors.
On text classiﬁcation, language modeling, and recommender systems,
ANT outperforms existing approaches with respect to accuracy and sparsity.
A PPENDIX A I NDIAN B UFFET P ROCESS WITH T WO P ARAMETERS
In this section we provide a more detailed treatment of the Indian Buffet Process (IBP) (Grifﬁths &
Ghahramani, 2005; 2011; Paisley et al., 2010) as well its two-parameter generalization (Ghahramani et al., 2007).
We begin with describing the single parameter version, which essentially deﬁnes a probability distribution over sparse binary matrices with a ﬁnite number of rows and an unbounded number of columns.
Under IBP prior with hyperparameter a , to generate such a sparse binary random matrix Z with ∣ V ∣ rows, we have the following process:
[Equation]
It can be shown from this construction that a given matrix Z will have non-zero probability under the
IBP prior if and only if the number of columns containing non-zero entries is ﬁnite, albeit a random quantity (Grifﬁths & Ghahramani, 2005).
Also note that b k keeps diminishing as 0 < v j < 1 , thus most z ik will be 0, thereby giving rise to the desired sparsity.
Moreover, it can be shown that number of number of non-empty columns would be O ( log ∣ V ∣) in expectation.
Like most Bayesian nonparametric models, it is best understood with an analogy.
Consider of a set of customers (objects) queued up to take dishes (features/anchors) from a buffet arranged in a line.
The ﬁrst customer starts at beginning of the buffet and takes a serving of all of the ﬁrst Poisson ( a ) dishes.
The remaining customers are more picky.
The i th customer samples dishes in proportion to their popularity (i.e., with probability m k / i ), where m k is the number of previous customers who have sampled a dish.
Having reached the end of all previous sampled dishes, the i th customer then tries
Poisson ( a / i ) new dishes.
The result of this process for the entire vocabulary V is a binary matrix Z with ∣ V ∣ rows and inﬁnitely many columns where z
ik = 1 if the i th customer sampled the k th dish.
Using either description of the IBP, we can ﬁnd the distribution of the sparse binary matrix Z with ∣ V ∣ rows and K non-empty columns to be given by (Grifﬁths & Ghahramani, 2005; Broderick et al., 2013b):
[Equation]
where m k denotes number of customers (objects) who selected dish (anchor) k , H ∣ V ∣ is the ∣ V ∣ -th
Harmonic number H ∣ V ∣ =
∑ ∣ V ∣ j = 1 1 j , and K h is the number of occurrences of the non-zero binary vector h among the columns in Z .
However, the number of features per object and the total number of features are directly coupled through the single parameter a .
The two-parameter generalization of the IBP allows one to independently tune the average number of features for each object and the overall number of features used across all V objects (Ghahramani et al., 2007).
In particular, we now have an additional hyperparameter b along with a .
The ﬁrst customer, as before, samples Poisson ( a ) dishes.
However, the i -th customer now samples in proportion to their popularity with probability m k /( i + b ) , where m k is the number of previous customers who have sampled a dish.
Having reached the end of all previously sampled dishes, the i th customer tries Poisson ( ab /( i + b )) new dishes.
The parameter b is introduced in such a way as to preserve the expected number of features per object to be still a , but the expected overall number of features is now O ( ab log ∣ V ∣) .
The total number of features used thus increases as b increases, thus providing more control on sparsity of Z .
This two parameter IBP prior for a binary matrices Z with ∣ V ∣ rows and K non-empty columns is given by:
[Equation]
where m k denotes number of customers (objects) who selected dish (anchor) k and H ∣ V ∣ =
∑ ∣ V ∣ j =1 b + j − 1 .
This distribution is suitable for use as a prior for Z in models that represent objects using a sparse but potentially inﬁnite array of features.
Historically, IBP has been used as a prior in latent feature models, where the binary entries z ik of a random matrix encode whether feature k is used to explain observation i .
The IBP can be further combined with a simple observation model p ( y i ∣ z i, ∶ ,θ ) from the exponential family like the Gaussian distribution (see e.g. Grifﬁths & Ghahramani (2011)).
The complexity of inference, using MCMC or variational methods, for such binary factor analysis models has limited the adoption of more complicated observation models.
In this work, we break this barrier and, to the best of our knowledge, are the ﬁrst to integrate IBP with deep representation learning of discrete objects by employing an efﬁcient SVA based inference.
Thus, our approach combines the representation capabilities of neural networks with desirable sparsity properties.
B D ERIVATION OF O BJECTIVE F UNCTION VIA SVA
In this section we derive our objective function using Small Variance Asymptotics (SVA) (Jiang et al., 2012).
Recall that the generative process in our model is given by:
• Z ∈ R ∣ V ∣× K ∼ IBP ( a,b ) • A ∈ R K × d ∼ P ( A ) =
N( 0 , 1 ) • W ∈ R ∣ V ∣× K ∼ P ( W ) =
Exponential ( 1 ) • for i = 1 , ⋯ ,N - ˆ y i
=
f θ ( x i ; ( Z ○ W ) A ) -
y i ∼ p ( y i ∣ x i ; Z , W , A ) =
exp {− D φ ( y i , ˆ y i )}
b φ ( y i )
The joint log-likelihood under our generative model above is therefore:
[Equation]
To use SVA, an approximate objective function for ﬁnding point estimates is obtained by taking the limit of the emission probability variances down to zero.
We begin by introducing a scaling variable β and shrinking the variance of the emission probability to 0 by taking β → ∞ .
The scaled probability emission becomes
[Equation]
Following Broderick et al.
(2013a), we modulate the number of features in the large- β limit by choosing constants λ 1 > λ 2 > 0 and setting the IBP hyperparameters with β as follows:
[Equation]
This prevents a limiting objective function that favors a trivial cluster assignment (every data point assigned to its own separate feature).
We now take the limit of the log-likelihood term by term:
[Equation]
• lim β →∞ 1 β log p
[Equation]
y i
[Equation]
x i
[Equation]
Z
[Equation]
• lim β →∞ 1 β log p
( Z ) =
− λ 2 ∥
Z ∥ 0 − ( λ 1 − λ 2 ) K , see box below.
• lim β →∞ 1 β log p
W ) = 0 , if W ≥ 0 else −∞
[Equation]
• lim β →∞ 1 β log p
( A ) =
as log p
( A ) =
O ( 1 ) .
For convenience
[Table]
we re-write the limit of the IBP prior as
[Table]
[Equation]
[Equation]
For part a (cid:13) :
[Equation]
For part b (cid:13) : lim
[Equation]
For part c (cid:13) : lim
[Equation]
We know that m k is the number of objects which uses anchor k which counts the number of non-zero entries in the k -th column of Z .
When we sum over all k , it just becomes the number of non-zero entries in Z , which is equivalent to the L 0 norm of Z , i.e., ∥ Z ∥ 0 .
Therefore, the MAP estimate under SVA as given by
[Equation]
is equivalent to optimizing the following objective function:
[Equation]
where the exponential prior for W resulted in a limiting domain for W to be positive.
Note that we can combine the optimizing variables Z and W with their constraints into one variable T ≥ 0 .
Also we can switch from a maximization problem to a minimization problem by absorbing the negative sign.
Finally we arrive at the desired objective:
[Equation]
C E XPONENTIAL F AMILY D ISTRIBUTIONS AS B REGMAN D IVERGENCES
In this section we provide some results that relate exponential families distributions and Bregman divergences.
As a result, we can relate likelihood models from Sec.
3.2 to appropriate Bregman divergences.
Thus, a probabilistic observation model can be translated to a loss functions minimizing the Bregman divergence, which are more amenable to deep network training using gradient based methods.
We begin by deﬁning the Bregman divergence below and stating the relationship formally in Theorem 1.
Deﬁnition 1.
(Bregman, 1967) Let φ ∶ S → R , S = dom ( φ ) be a strictly convex function deﬁned on a convex set S
⊂ R d
such that φ is differentiable on ri ( S ) , assumed to be non-empty.
The Bregman divergence D φ ∶ S × ri ( S ) → [ 0 , ∞) is deﬁned as
[Equation]
22) where ∇ φ ( y ) represents the gradient vector of φ evaluated at y .
Theorem 1.
(Banerjee et al., 2005)
There is a bijection between regular exponential families and regular Bregman divergences.
In particular, for any exponential family distribution p ( x ∣ θ ) =
p 0 ( x ) exp (⟨ x , θ ⟩ − g ( θ )) can be written as p ( x ∣ µ ) =
exp (− D φ ( x , µ ))
b φ ( x ) where φ is the Legendre dual of the log-partition function g ( θ ) and µ = ∇ θ g ( θ ) .
From Theorem 1, we can see that maximizing log-likelihood log p ( x ∣ θ ) is same as minimizing the
Bregman divergence D φ ( x , µ ) .
Note that we can ignore b φ ( x ) as it depends only on observed data and does not depend on any parameters.
We now illustrate some common examples of exponential families (like Gaussian and categorical), derive their corresponding Bregman divergences, and connect to usual loss functions used in deep networks (like MSE and cross-entropy).
Example 1: Gaussian distribution.
(Banerjee et al., 2005)
We start with the unit variance spherical
Gaussian distributions with with mean µ , which have densities of the form:
[Equation]
Using the log-partition function for Gaussian distribution, we can calculate that φ ( x ) = 12 ∥ x ∥ 2 , which yields Bregman divergence equal to:
[Equation] [Equation]
[Equation]
Thus, D φ ( x , µ ) along with constant b φ ( x ) given by
[Equation]
27) recovers the Gaussian density p ( x ) = exp (− D φ ( x , µ ))
b φ ( x ) .
Therefore, when we assume that labels have a Gaussian emmission model, the corresponding Bregman divergence D φ ( x , µ ) = 12 ∥
x − µ ∥recovers the squared loss commonly used for regression.
Example 2: Multinomial distribution.
(Banerjee et al., 2005)
Another exponential family that is widely used is the family of multinomial distributions:
[Equation]
where x j ∈ Z + are frequencies of events, ∑ dj = 1 x j = N and q j ≥ 0 are probabilities of events, ∑ d j = 1 q j = 1 .
The multinomial density can be expressed as the density of an exponential distribution in x = { x j } d − 1 j = 1 with natural parameter θ = log ( q j q d ) d − 1 j = 1 , cumulant function g ( θ )
=
− N log q d , and expectation parameter µ = ∇ g ( θ ) =
[ Nq j ] d − 1 j = 1 .
The Legendre dual φ of g is given by
[Equation]
As a result, the multinomial density can be expressed as a Bregman divergence equal to:
[Equation]
[Equation]
and constant b φ ( x ) given by
[Equation]
which recovers the multinomial density p ( x ) = exp (− D φ ( x , µ ))
b φ ( x ) .
Therefore, when the labels are generated from a multinomial distribution, the corresponding Bregman divergence D φ ( x , µ ) =
−∑ d j = 1 x j log µ j + constant recovers the cross-entropy loss commonly used for classiﬁcation.
D L EARNING THE A NCHOR E MBEDDINGS A
Here we provide several other strategies for initializing the anchor embeddings:
• Sparse lasso and variational dropout (Chen et al., 2019).
Given the strong performance of sparse lasso and variational dropout as vocabulary selection methods, it would be interesting to use sparse lasso/variational dropout to ﬁrst select the important task-speciﬁc words before jointly learning their representations and their transformations to other words.
However, sparse lasso and variational dropout require ﬁrst training a model to completion unlike frequency and clustering based vocabulary selection methods that can be performed during data preprocessing.
• Coresets involve constructing a reduced data set which can be used as proxy for the full data set, with provable guarantees such that the same algorithm run on the coreset and the full data set gives approximately similar results (Phillips, 2016; Har-Peled & Mazumdar, 2004).
Coresets can be approximately computed quickly (Bachem et al., 2017) and can be used to initialize the set of anchors A .
In general, there is a trade-off between how quickly we can choose the anchor objects and their performance.
Randomly picking anchor objects (which is equivalent to initializing the anchor embeddings with dynamic basis vectors) becomes similar to learning a low-rank factorization of the embedding matrix (Sedov & Yang, 2018), which works well for general cases but can be improved for task-speciﬁc applications or with domain knowledge.
Stronger vocabulary selection methods like variational dropout and group lasso would perform better but takes signiﬁcantly longer time to learn.
We found that intermediate methods such as frequency, clustering, with WordNet/co-occurrence information works well while ensuring that the preprocessing and training stages are relatively quick.
In Appendix K we provide more results for different initialization strategies including those based on clustering initializations.
In general, performance is robust with respect to the choice of A among the ones considered (i.e., random, frequency, and clustering).
While frequency and clustering work better, using a set of dynamic basis embeddings still gives strong performance, especially when combined with domain knowledge from WordNet and co-occurrence statistics.
This implies that when the user has more information about the discrete objects (e.g., having a good representation space to perform clustering), then the user should do so.
However, for a completely new set of discrete objects, simply using low-rank basis embeddings with sparsity also work well.
E T RANSFORM : L EARNING A S PARSE T
In addition to a simple sparse linear transformation, we describe some extensions that improve sparsity and expressitivity of the learned representations.
Reducing redundancy in representations: To further reduce redundancy in our sparse representations, we perform orthogonal regularization of dynamic basis vectors A by adding the loss term
Conﬁdential + Proprietary
A frequencyclustering basis
A2A3A x M x M ... d |V|
E
T1T2 T4T
[Figure]
L( A ) =
∑ i ≠ j ∣ a ⊺ i a j ∣ to the loss function in eq (1).
This ensures that different basis vectors a i and a j are orthogonal instead of being linear combinations of one another which would lead to redundancies across different learnt entries in T .
Mixture of anchors: In general, different initialization strategies may bring about different advantages.
For example, using a mixture of random basis vectors has been shown to help model multisense embeddings (Athiwaratkun et al., 2018; Nguyen et al., 2017).
One can deﬁne a set of M anchor embeddings A 1 ,..., A M each initialized by different strategies and of possibly different sizes.
Nonlinear mixture of transformations:
To complement learning multiple sets of anchor embeddings A 1 ,..., A M , the straightforward extension of the T RANSFORM step would be to learn a separate linear transformation for each anchor embedding and summing the result:
E =
∑ Mm = 1 T m A m .
However, the expressive power of this linear combination is equivalent to one set of anchor embeddings equal to concatenating A 1 ,..., A M and one linear transformation.
To truly exhibit the advantage of multiple anchors, we transform and combine them in a nonlinear fashion, e.g.,
E =
∑ Mm = 1 softmax ( T m )
A m (softmax over the rows of T m , Figure 5).
Different transformations can be learned for different initializations of anchors.
This is connected with the multi-head attention mechanism in the Transformer (Vaswani et al., 2017), where softmax ( T m ) are the softmax-activated (sparse) attention weights and
A m the values to attend over.
The result is an embedding matrix formed via a nonlinear mixture of anchors (each initialized with different strategies) and sparse transformations.
F I NCORPORATING D OMAIN K NOWLEDGE
ANT also allows incorporating domain knowledge about object relationships.
Suppose we are given some relationship graph G
=
( V,E ) where each object is a vertex v ∈ V and an edge ( u,v )
∈ E exists between objects u and v if they are related.
Real-world instantiations of such a graph include 1) WordNet (Miller, 1995) or ConceptNet (Liu & Singh, 2004) for semantic relations between words, 2) word co-occurrence matrices (Haralick et al., 1973), and 3) Movie Clustering datasets (Leskovec & Krevl, 2014).
From these graphs, we extract related positive pairs P = {( u,v ) ∈ E } and unrelated negative pairs N = {( u,v ) ∉ E } .
We incorporate domain information as follows (see Figure 6 for a visual example):
Positive pairs: To incorporate a positive pair ( u,v ) , we do not enforce sparsity on T u,v .
This allows
ANT to freely learn the transformation between related objects u and v without being penalized for sparsity.
On the other hand, transformations between negative pairs will be sparsely penalized.
In other words, before computing the (cid:96) 1 -penalty, we element-wise multiply T with a domain sparsity matrix S ( G )
where S ( G ) u,v = 0 for ( u,v ) ∈ P (entries not (cid:96) 1 -penalized) and S ( G ) u,v = 1 otherwise (entries are (cid:96) 1 -penalized), resulting in the following modiﬁed objective:
[Equation]
Since we perform proximal GD, this is equivalent to only soft-thresholding the entries between unrelated objects, i.e., T = max {( T − ηλ 2 ) ⊙ S ( G ) + T ⊙ ( 1 − S ( G )) , 0 } .
Note that this strategy is applicable when anchors are selected using the frequency method.
Negative pairs: For negative pairs, we add an additional constraint that unrelated pairs should not share entries in their linear combination coefﬁcients of the anchor embeddings.
In other words, we
Conﬁdential + Proprietary “good” “the”
[Table]
pretrained space e.g
[Table]
GloVe/co
[Table]
occurrence
Conﬁdential + Proprietary “good” “the” Clustering steppretrained space e.g. GloVe/co-occurrence
Conﬁdential + Proprietary “good” “the” Clustering steppretrained space e.g. GloVe/co-occurrence Conﬁdential + Proprietary “good” “the” Clustering steppretrained space e.g. GloVe/co-occurrence
Conﬁdential
[Table]
Proprietary
[Table]
[Figure]
add the loss term
[Equation]
[Equation]
to the loss in eq (1), where each inner sum discourages t u and t v from sharing similar entries.
This strategy can used regardless of the way anchors are selected.
We acknowledge that there are other ways to incorporate domain knowledge as well into the general ANT framework, and we only serve to give some initial examples of such methods.
G N ONPARAMETRIC A NCHOR & T RANSFORM
In this section we provide details for our non-parametric extension of ANT.
Recall that our full objective function derived via small variance asymptotics is given by:
[Equation]
(34) which suggests a natural objective function in learning representations that minimize the prediction loss D φ ( y i ,f θ ( x i ; TA )) while ensuring sparsity of T as measured by the (cid:96) 0 -norm and using as few anchors as possible ( K ).
Therefore, optimizing eq (5) gives rise to a nonparametric version of
ANT, which we call NB ANT, that automatically learns the optimal number of anchors.
To perform optimization over the number of anchors, our algorithm starts with a small initial number of anchors
K = ∣ A ∣ = 10 and either adds ∆ K anchors (i.e., adding ∆ K new rows to A and ∆ K new sparse columns to T ) or deletes ∆ K anchors to minimize eq (34) at every epoch depending on the trend of the objective evaluated on the training set.
We detail the full Algorithm 2, and highlight the main changes as compared to ANT.
Practically, this algorithm involves the same number of training epochs and batches through each training epoch as the vanilla ANT method.
To enable sharing of trained anchors, we change the indices from where A and T are read from so that the partially trained removed anchors are still stored in case more anchors need to be added again.
H E FFICIENT L EARNING AND I NFERENCE
The naive method for learning E from anchor embeddings A and the sparse transformations T still scales linearly with ∣ V ∣ × d .
Here we describe some tips on how to perform efﬁcient learning and inference of the anchor embeddings A and the sparse transformations T :
• Store T as a sparse matrix by only storing its non-zero entries and indices.
From our
experiments
, we have
shown that nnz ( T ) <<
∣ V ∣ ×
d which makes storage efﬁcient.
• For inference, use sparse matrix multiply as supported in TensorFlow and PyTorch to compute E
= TA (or its non-linear extensions).
This decreases the running time from scaling by ∣ V
∣ ×
d to only scaling as a function of nnz ( T ) .
For training, using inbuilt sparse
Compute eq using current value
K
A
T the validation set
10: if eq is on a decreasing trend then 11:
12: else if eq is on an increasing trend then 13:
K
K columns
14: else 15: keep current values of K, A , T .
16: end if 17: end for 18: return anchor embeddings A transformations
representation of most deep learning frameworks
like
PyTorch
or Tensorﬂow
is not
optimal, as
they do not support changing non-zero locations in sparse matrix and apriori its not easy to ﬁnd optimal set of non-zero locations
• During training, instead, implicitly construct E from its anchors and transformations.
In fact, we can do better: instead of constructing the entire E matrix to embed a single datapoint x ∈ R 1 ×∣ V ∣ , we can instead ﬁrst index x into T , i.e., xT ∈ R 1 ×∣
A ∣ before performing a sparse matrix multiplication with A , i.e., ( xT ) A ∈ R 1 × d .
We are essentially taking advantage of the associative property of matrix multiplication and the fact that xT is a simple indexing step and ( xT ) A is an effective sparse matrix multiplication.
To enable fast row slicing into sparse matrix, we just storing the matrix in adjacency list or CSOO format.
(We move away from CSR as adding/deleting a non-zero location is very expensive.)
When gradient comes back, only update the corresponding row in T .
The gradient will be sparse as well due to the L1-prox operator.
• Above trick solves the problem for tasks where embedding is used only at the input, e.g., classiﬁcation.
For tasks like language model, where embedding is used at output as well one can also use above mentioned trick with speedup techniques like various softmax sampling techniques (Bengio & Senecal, 2008; Mikolov et al., 2013) or noise-contrastive estimation (Gutmann & Hyvarinen, 2010; Mnih & Teh, 2012), which will be anyway used for large vocabulary sizes.
To elaborate, consider the case of sampled softmax (Bengio & Senecal, 2008).
We normally generate the negative sample indices, and then we can ﬁrst index into T using the true and negative indices before performing sparse matrix multiplication with A .
This way we do not have to instantiate entire E by expensive matrix multiplication.
• When training is completed, only store the non-zero entries of T or store T as a sparse matrix to reconstruct E for inference.
• To save time when initializing the anchor embeddings and incorporating domain knowledge, precompute the necessary statistics such as frequency statistics, co-occurrence statistics, and object relation statistics.
We use a small context size of 10 to measure co-occurrence of two words to save time.
When using WordNet to discover word relations, we only search for immediate relations between words instead of propagating relations across multiple steps (although this could further improve performance).
• In order to incorporate domain knowledge in the sparsity structure, we again store 1 − S ( G ) using sparse matrices.
Recall that S ( G ) has an entry equal to 1 for entries representing unrelated objects that should be (cid:96) 1 -penalized, which makes S ( G ) quite dense since most anchor and non-anchor objects are unrelated.
Hence we store 1 −
S ( G ) instead which consists few non-zero entries only at (non-anchor, anchor) entries for related objects.
Element-wise multiplications are also replaced by sparse element-wise multiplications when computing
T ⊙ S ( G ) and T ⊙ ( 1 − S ( G )) .
[Table]
• Finally, even if we want to use our ANT framework with full softmax in language model, it is possible without blowing up memory requirements.
In particular, let g ∈ R 1 ×∣ V ∣ be the incoming gradient from cross-entropy loss and h ∈ R d × 1 be the vector coming from layers below, like LSTM.
The gradient update is then
[Equation]
The main issue is computing the huge ∣ V ∣ × ∣ A ∣ outer product as an intermediate step which will be dense.
However, note that incoming gradient g is basically a softmax minus an offset corresponding to correct label.
This should only have large values for a small set of words and small for others.
If we carefully apply the (cid:96) 1 -prox operator earlier, which is nothing but a soft-thresholding, we can make this incoming gradient sparse very sparse.
Thus we need to only calculate a much smaller sized outer product and touch a small number of rows in T .
Thus, making the approach feasible.
I G ENERALITY OF ANT
We show that under certain structural assumptions on the anchor embeddings and transformation matrices, ANT reduces to the following task-speciﬁc methods for learning sparse representations: 1) Frequency (Chen et al., 2016b),
TF-IDF, Group Lasso (Wen et al., 2016), and variational dropout (Chen et al., 2019)
based vocabulary selection, 2) Low-rank factorization (Grachev et al., 2019), and 3)
Compositional code learning (Shu & Nakayama, 2018; Chen et al., 2018).
Hence, ANT is general and uniﬁes some of the work on sparse representation learning done independently in different research areas.
Frequency-based vocabulary selection (Luong et al., 2015; Chen et al., 2016b):
Initialize A with the ∣ A ∣ most frequent objects and set T a,a = 1 for all a ∈ A , T = 0 otherwise.
Then E
= TA consists of embeddings of the ∣ A ∣ most frequent objects with zero embeddings for all others.
During training, gradients are used to update A but not T (i.e., only embeddings for frequent objects are learned).
By changing the selection of A , ANT also reduces to other vocabulary selection methods such as
TF-IDF (Ramos, 1999), Group Lasso (Wen et al., 2016), and variational dropout (Chen et al., 2019)
Low-rank factorization (Acharya et al., 2019; Markovsky, 2011; Grachev et al., 2019):
Initialize
A by a mixture of random basis embeddings (just 1 anchor per set)
A 1 ,..., A M ∈ R 1 ×
d
and do not enforce any sparsity on the transformations T 1 ,..., T M ∈ R ∣ V ∣× 1 .
If we further restrict ourselves to only linear combinations E =
∑ Mm = 1 T m A m , this is equivalent to implicitly learning the M low rank factors a 1 ,..., a M , t 1 ,..., t M that reconstruct embedding matrices of rank at most M .
[Table]
Compositional code learning (Shu & Nakayama, 2018; Chen et al., 2018):
Initialize A by a mixture of random basis embeddings A 1 ,..., A M , initialize transformations T 1 ,..., T M , and apply a linear combination E
=
∑ Mm = 1 T m A m .
For sparsity regularization, set row i of S ( G ) mi as a reverse one-hot vector with entry d mi = 0 and all else 1.
In other words, index d mi of row row T mi is not regularized, and all other entries are (cid:96) 1 -regularized with extremely high λ 2 such that row T mi essentially becomes an one-hot vector with dimension d mi = 1 .
This results in learning a codebook where each object in
V is mapped to only one anchor in each mixture.
Therefore, ANT encompasses several popular methods for learning sparse representations, and gives further additional ﬂexibility in deﬁning various initialization strategies, applying nonlinear mixtures of transformations, and incorporating domain knowledge via object relationships.
J E XPERIMENTAL D ETAILS
Here we provide more details for our experiments including hyperparameters used, design decisions, and comparison with baseline methods.
We also include the anonymized code in the supplementary material.
J. T EXT C LASSIFICATION
Base CNN model: For all text classiﬁcation experiments, the base model is a CNN (Lecun et al., 1998) with layers of 2 D convolutions and 2 D max pooling, before a dense layer to the output softmax.
The code was adapted from https://github.com/wenhuchen/
Variational-Vocabulary-Selection and the architecture hyperparameters are provided in Table 6.
The only differences are the output dimensions which is 4 for AG-News, 14 for DBPedia, 5 for Sogou-News, and 5 for Yelp-review.
Anchor: We experiment with dynamic, frequency, and clustering initialization strategies.
The number of anchors ∣ A ∣ is a hyperparameter that is selected using the validation set.
The range of ∣ A ∣ is in { 10 , 20 , 50 , 80 , 100 , 500 , 1 , 000 } .
Smaller values of ∣ A ∣ allows us to control for fewer anchors and smaller transformation matrix T at the expense of performance.
Transformation: We experiment with sparse linear transformations for T .
λ 2 is a hyperparameter that is selected using the validation set.
Larger values of λ 2 allows us to control for more sparse entries in T at the expense of performance.
For experiments on dynamic mixtures, we use a softmax-based nonlinear combination E =
∑ M m = 1 softmax ( T m )
A m where softmax is performed over the rows of
[Table]
T m .
Note that applying a softmax activation to the rows of T m makes all entries dense so during training, we store T m as sparse matrices (which is efﬁcient since T m has few non-zero entries) and implicitly reconstruct E .
Domain knowledge: When incorporating domain knowledge in ANT, we use both WordNet and cooccurrence statistics.
For WordNet, we use the public WordNet interface provided by NLTK http: //www.nltk.org/howto/wordnet.html .
For each word we search for its immediate related words among its hypernyms, hyponyms, synonyms, and antonyms.
This deﬁnes the relationship graph.
For co-occurrence statistics, we deﬁne a co-occurrence context size of 10 on the training data.
Two words are deﬁned to be related if they co-occur within this context size.
A note on baselines: Note that the reported results on S PARSE VD and S PARSE VD-V OC (Chirkova et al., 2018) have a different embedding size: 300 instead of 256 .
This is because they use pre-trained word2vec or GloVe embeddings to initialize their model before compression is performed.
[Equation]
[Table]
ODELING ON PTB
Base LSTM model:
Our base model is a 2 layer LSTM with an embedding size of 200 and hidden layer size of 200 .
The code was adapted from https://github.com/salesforce/ awd-lstm-lm and the full table of hyperparameters is provided in Table 7.
Base AWD-LSTM model:
In addition to experiments on an vanilla LSTM model as presented in the main text, we also performed experiments using a 3 layer AWD-LSTM with an embedding size of 400 and hidden layer size of 1 , 150 .
The full hyperparameters used can be found in Table 8.
Anchor: We experiment with dynamic, frequency, and clustering initialization strategies.
The number of anchors ∣ A ∣ is a hyperparameter that is selected using the validation set.
The range of ∣ A ∣ is in { 10 , 20 , 50 , 80 , 100 , 500 , 1 , 000 } .
Smaller values of ∣ A ∣ allows us to control for fewer anchors and smaller transformation matrix T at the expense of performance.
Domain knowledge: When incorporating domain knowledge in ANT, we use both WordNet and cooccurrence statistics.
For WordNet, we use the public WordNet interface provided by NLTK http: //www.nltk.org/howto/wordnet.html .
For each word we search for its immediate related words among its hypernyms, hyponyms, synonyms, and antonyms.
This deﬁnes the relationship graph.
For co-occurrence statistics, we deﬁne a co-occurrence context size of 10 on the training data.
Two words are deﬁned to be related if they co-occur within this context size.
[Table]
A note on baselines:
We also used some of the baseline results as presented in Grachev et al.
Their presented results differ from our computations in two aspects: they include the LSTM parameters on top of the embedding parameters, and they also count the embedding parameters twice since they do not perform weight tying (Press & Wolf, 2017)
(see equation (6) of Grachev et al.
To account for this, the results of S PARSE VD and S PARSE VD-V OC (Chirkova et al., 2018), as well as the results of various LR and TT low rank compression methods (Grachev et al., 2019) were modiﬁed by subtracting off the LSTM parameters ( 200 × 200 × 16 ).
This is derived since each of the 8 weight matrices W i,f,o,c ,U i,f,o,c in an LSTM layer is of size 200 × 200 , and there are a 2 LSTM layers.
We then divide by two to account for weight tying.
In the main text, we compared with the strongest baselines as reported in Grachev et al.
(2019): these were the methods that performed low rank decomposition on both the input embedding ( ∣ V ∣ × d ), output embedding ( d × ∣ V ∣ ), and intermediate hidden layers of the model.
For full results, please refer to Grachev et al.
Note that the reported results on S PARSE VD and S PARSE VD-V OC (Chirkova et al., 2018) have a different embedding size and hidden layer size of 256 instead of 200 , although these numbers are close enough for fair comparison.
In our experiments we additionally implemented an LSTM with an embedding size of 256 and hidden layer size of 256 so that we can directly compare with their reported numbers.
For baselines that perform post-processing compression of the embedding matrix, P OST -S PARSE
H ASH (post-processing using sparse hashing)
(Guo et al., 2017) and P OST -S PARSE H ASH + k -SVD (improving sparse hashing using k -SVD)
(Guo et al., 2017; Awasthi & Vijayaraghavan, 2018), we choose two settings: the ﬁrst using 500 anchors and 10 nearest neighbors to these anchor points, and the second using 1 , 000 anchors and 20 nearest neighbors.
The ﬁrst model uses 500 × d + ∣ V ∣ ×non-zero embedding parameters while the second model uses 1 ,
×
d + ∣ V ∣ × 20 parameters.
For
AWD-LSTM on PTB, this is equivalent to 0 .
3 M and 0 .
6 M embedding parameters respectively which is comparable to the number of non-zero parameters used by our method.
J. L ANGUAGE M ODELING ON W IKI T EXT -
Base AWD-LSTM model:
Our base model is a 4 layer AWD-LSTM with an embedding size of 400 and hidden layer size of 2 , 500 .
The code was adapted from https://github.com/ salesforce/awd-lstm-lm and the hyperparameters used can be found in Table 9.
[Table]
A note on baselines: While Baevski & Auli (2019) adapt embedding dimensions according to word frequencies, their goal is not to compress embedding parameters and they use 44 .
9 M (dense) parameters in their adaptive embedding layer, while we use only 2 M. Their embedding parameters are calculated by their reported bucket sizes and embedding sizes (three bands of size 20 K ( d = 1024 ), 40 K ( d = 256 ) and 200 K ( d = 64 )).
Their perplexity results are also obtained using a Transformer model with 250 M params while our AWD-LSTM model uses 130 M params.
For the H ASH E MBED baseline that retains the frequent k words and hashes the remaining words into 1 , 000 OOV buckets (Svenstrup et al., 2017), We vary k ∈ { 1 × 10 5 , 5 × 10 4 , 1 × 10 4 } to obtain results across various parameter settings.
J. M OVIE R ECOMMENDATION ON M OVIE L ENS
Base MF model: We show the hyperparamters used for the MF model in Table 10.
We use the Yogi optimizer (Zaheer et al., 2018) to learn the parameters.
ANT and NB ANT :
We build ANT on top of the MF model while keeping the base hyperparamters constant.
For ANT, we apply compression to movie and user embedding matrices individually.
NB ANT involves deﬁning the starting value of K = ∣ A ∣ , and a ∆ K value which determines the rate of increase or decrease in K .
For Movielens 25 M we use a larger initial ∣ A ∣ and ∆ K since it is a larger dataset and also takes longer to train, so we wanted the increase and decrease in anchors to be faster (see Table 10).
Beyond this initial setting, we found that performance is robust with respect to the initial value of K and ∆ K , so we did not tune these parameters.
In practice, we tie the updates of the number of user anchors and movie anchors instead of optimizing over both independently.
Therefore, we start with the same number of initial user and movie anchors before incrementing or decrementing them by the same ∆ K at the same time.
We found that this simpliﬁcation did not affect performance and NB ANT was still able to ﬁnd an optimal number of anchors for a good trade-off between performance and compression.
K M ORE R ESULTS
In the following sections, we provide additional results on learning sparse representations of discrete objects using ANT.
[Table]
K. T EXT CLASSIFICATION
Extra results: We report additional text classiﬁcation results on AG-News, DBPedia, Sogou-News, and Yelp-review in Table 11.
For AG-News, using a mixture of anchors and transformations also achieves stronger performance than the baselines using 5 anchors per mixture, although the larger number of transformations leads to an increase in parameters.
Our approach with different initializations and domain knowledge achieves within 1% accuracy with 21 × fewer parameters on DBPedia, within 1% accuracy with 10 × fewer parameters on Sogou-News
[Table]
and within 2% accuracy with 22 × fewer parameters on Yelp-review.
Different initialization strategies: Here we also presented results across different initialization strategies and ﬁnd that while those based on frequency and clustering work better, using a set of dynamic basis embeddings still gives strong performance, especially when combined with domain knowledge from WordNet and co-occurrence statistics.
This implies that when the user has more information about the discrete objects (e.g., having a good representation space to perform clustering),
[Table] [Table]
then the user should do so.
However, for a completely new set of discrete objects, simply using low-rank basis embeddings with sparsity also work well.
Incorporating domain knowledge: We ﬁnd that from WordNet and co-occurrence helps to further reduce the total embedding parameters while maintaining task performance.
K. L ANGUAGE MODELING
Extra results on PTB: We report additional language modeling results using AWD-LSTM on PTB in Table 12.
ANT with 1 , 000 dynamic basis vectors is able to compress the embedding parameters by 10 × while achieving 72 . 0 test perplexity.
By incorporating domain knowledge, we further compress
[Table]
NB ANT Auto → 11 ∗ Auto → 11 ∗ Random ( λ 1 = 0 .
01 ,λ 2 = 1 × 10 − 4 ) (cid:51) (cid:51)
Auto → 11 ∗ Auto → 11 ∗ Random ( λ 1 = 0 .
01 ,λ 2 = 2 × 10 − 4 ) (cid:51) (cid:51)
the embedding parameters by another 10 × and achieve 70 . 0 test perplexity, which results in 100 × total compression as compared to the baseline.
We also perform more controlled experiments with different embedding dimension sizes 200 and 250 where we also outperform the baselines.
Extra results on WikiText-103: We also report full results using AWD-LSTM on WikiTextin Table 13, where we reach within 3 points perplexity with ∼ 16 × reduction and within 13 points perplexity with ∼ 80 × reduction, outperforming the frequency (H ASH E MBED ) and post-processing hashing (S PARSE H ASH ) baselines.
K. M OVIE R ECOMMENDATION
Extra results on MovieLens 1M: We also report results on MovieLens 1M in Table 14, where we also observe improvements in accuracy and compression.
We also run NB ANT on Movielens 1M.
Despite the small size of the dataset, NB ANT is able to optimize for ∣ A ∣ quickly and achieve a good trade-off between performance and compression.
Extra results on MovieLens 25 M: Finally, we provide a 3D version of the 2D plot shown in Figurewhere we simpliﬁed the plot by showing grid points with an equal number of user and movie anchors.
We provide several orientation views of the full 3D plot across user anchors ( x -axis), movie anchors ( y -axis), and objective value of eq (34) ( z -axis) in Figure 7.
NB ANT reaches the objective value indicated by the shaded red plane which is close to the optimal objective value as computed over all grid search experiments for user and movie anchors.
Therefore, NB ANT can efﬁciently optimize ∣ A ∣ and reach a good value of eq (5) from just one run.
Online NB ANT: Since NB ANT automatically grows/contracts ∣ A ∣ during training, we can further extend NB ANT to an online version that sees a stream of batches without revisiting previous ones Bryant & Sudderth (2012).
This further enables NB ANT to scale to large datasets that cannot all ﬁt in memory.
We treat each batch as a new set of data coming in and train on that batch until convergence, modify ∣ A ∣ as in Algorithm 2, before moving onto the next batch.
In this signiﬁcantly more challenging online setting, NB ANT is still able to learn well and achieve a MSE of 0 .
875 with 1 .
25 M non zero parameters.
From Figure 8, we found that initially ∣ A ∣ grew steadily from 10 up to 26 as online batches were seen.
As even more batches were seen, the number of clusters decreased steadily fromto 10 , and oscillated between 8 and 10 .
This means initially some connections/groupings between online batches were not clear, but with more data, natural clusters merged together.
Interestingly this online version of NB ANT settled on a similar range of ﬁnal user ( 8 ) and item ( 8 ) anchors as compared to the non-online version (see Table 3), which conﬁrms the robustness of NB ANT in ﬁnding relevant anchors automatically.
L E FFECT OF λ 1 AND λ
In this section we further study the effect of the hyperparameters λ 1 and λ 2 .
Recall that our full objective function derived via small variance asymptotics is given by:
[Equation]
The role of hyper-parameter λ 2 is clear: For a ﬁxed λ 1 and integer valued variable K , tuning λcontrols the trade-off between sparsity of T and model performance (see Table 13).
The role of hyper-parameter λ 1 is more subtle.
It can be considered as a weighing fraction for scalarizing an underlying multi-objective optimization problem.
To elaborate, one can consider our goal as a multi-objective problem of minimizing the predictive loss while simultaneously using a minimal number of anchors ( K ).
Then the hyperparameter λ 1 can be used to select a solution along the Pareto front.
In other words, tuning the hyperparameter λ 1 allows us to perform model selection by controlling the trade-off between the number of anchors used and prediction performance.
We apply eq (36) on the trained models in Table 2 and report these results in Table 15.
Choosing a small λ 1 = 2 × 10 − 5 selects a model with more anchors ( ∣ A ∣ = 1 , 000 ) and better performance (ppl = 79 . 4 ), while a larger λ 1 = 1 × 10 − 1 selects the model with fewest anchors ( ∣ A ∣ = 100 ) with a compromise in performance (ppl = 106 . 6 ).
[Table]
.25 K
[Table]
[Table]