Structured Embedding Models for Grouped Data

Word embeddings are a powerful approach for analyzing language, and exponential family embeddings (EFE) extend them to other types of data. Here we develop structured exponential family embeddings (S-EFE), a method for discovering embeddings that vary across related groups of data. We study how the word usage of U.S. Congressional speeches varies across states and party affiliation, how words are used differently across sections of the ArXiv, and how the co-purchase patterns of groceries can vary across seasons. Key to the success of our method is that the groups share statistical information. We develop two sharing strategies: hierarchical modeling and amortization. We demonstrate the benefits of this approach in empirical studies of speeches, abstracts, and shopping baskets. We show how S-EFE enables group-specific interpretation of word usage, and outperforms EFE in predicting held-out data.

section_name: Introduction
paragraphs: Word embeddings BIBREF0 , BIBREF1 , BIBREF2 , BIBREF3 , BIBREF4 , BIBREF5 , BIBREF6 are unsupervised learning methods for capturing latent semantic structure in language. Word embedding methods analyze text data to learn distributed representations of the vocabulary that capture its co-occurrence statistics. These representations are useful for reasoning about word usage and meaning BIBREF7 , BIBREF8 . Word embeddings have also been extended to data beyond text BIBREF9 , BIBREF10 , such as items in a grocery store or neurons in the brain. efe is a probabilistic perspective on embeddings that encompasses many existing methods and opens the door to bringing expressive probabilistic modeling BIBREF11 , BIBREF12 to the problem of learning distributed representations.
We develop sefe, an extension of efe for studying how embeddings can vary across groups of related data. We will study several examples: in U.S. Congressional speeches, word usage can vary across states or party affiliations; in scientific literature, the usage patterns of technical terms can vary across fields; in supermarket shopping data, co-purchase patterns of items can vary across seasons of the year. We will see that sefe discovers a per-group embedding representation of objects. While the naïve approach of fitting an individual embedding model for each group would typically suffer from lack of data—especially in groups for which fewer observations are available—we develop two methods that can share information across groups.
Figure FIGREF1 illustrates the kind of variation that we can capture. We fit an sefe to ArXiv abstracts grouped into different sections, such as computer science (cs), quantitative finance (q-fin), and nonlinear sciences (nlin). sefe results in a per-section embedding of each term in the vocabulary. Using the fitted embeddings, we illustrate similar words to the word 1.10intelligence. We can see that how 1.10intelligence is used varies by field: in computer science the most similar words include 1.10artificial and 1.10ai; in finance, similar words include 1.10abilities and 1.10consciousness.
In more detail, embedding methods posit two representation vectors for each term in the vocabulary; an embedding vector and a context vector. (We use the language of text for concreteness; as we mentioned, efe extend to other types of data.) The idea is that the conditional probability of each observed word depends on the interaction between the embedding vector and the context vectors of the surrounding words. In sefe, we posit a separate set of embedding vectors for each group but a shared set of context vectors; this ensures that the embedding vectors are in the same space.
We propose two methods to share statistical strength among the embedding vectors. The first approach is based on hierarchical modeling BIBREF13 , which assumes that the group-specific embedding representations are tied through a global embedding. The second approach is based on amortization BIBREF14 , BIBREF15 , which considers that the individual embeddings are the output of a deterministic function of a global embedding representation. We use stochastic optimization to fit large data sets.
Our work relates closely to two threads of research in the embedding literature. One is embedding methods that study how language evolves over time BIBREF16 , BIBREF17 , BIBREF18 , BIBREF19 , BIBREF20 , BIBREF21 . Time can be thought of as a type of “group”, though with evolutionary structure that we do not consider. The second thread is multilingual embeddings BIBREF22 , BIBREF23 , BIBREF24 , BIBREF25 ; our approach is different in that most words appear in all groups and we are interested in the variations of the embeddings across those groups.
Our contributions are thus as follows. We introduce the sefe model, extending efe to grouped data. We present two techniques to share statistical strength among the embedding vectors, one based on hierarchical modeling and one based on amortization. We carry out a thorough experimental study on two text databases, ArXiv papers by section and U.S. Congressional speeches by home state and political party. Using Poisson embeddings, we study market basket data from a large grocery store, grouped by season. On all three data sets, sefe outperforms efe in terms of held-out log-likelihood. Qualitatively, we demonstrate how sefe discovers which words are used most differently across U.S. states and political parties, and show how word usage changes in different ArXiv disciplines.
section_name: Model Description
paragraphs: In this section, we develop sefe, a model that builds on efe BIBREF10 to capture semantic variations across groups of data. In embedding models, we represent each object (e.g., a word in text, or an item in shopping data) using two sets of vectors, an embedding vector and a context vector. In this paper, we are interested in how the embeddings vary across groups of data, and for each object we want to learn a separate embedding vector for each group. Having a separate embedding for each group allows us to study how the usage of a word like 1.10intelligence varies across categories of the ArXiv, or which words are used most differently by U.S. Senators depending on which state they are from and whether they are Democrats or Republicans.
The sefe model extends efe to grouped data, by having the embedding vectors be specific for each group, while sharing the context vectors across all groups. We review the efe model in Section SECREF4 . We then formalize the idea of sharing the context vectors in Section SECREF8 , where we present two approaches to build a hierarchical structure over the group-specific embeddings.
section_name: Background: Exponential Family Embeddings
paragraphs: In exponential family embeddings, we have a collection of objects, and our goal is to learn a vector representation of these objects based on their co-occurrence patterns.
Let us consider a dataset represented as a (typically sparse) matrix INLINEFORM0 , where columns are datapoints and rows are objects. For example, in text, each column corresponds to a location in the text, and each entry INLINEFORM1 is a binary variable that indicates whether word INLINEFORM2 appears at location INLINEFORM3 .
In efe, we represent each object INLINEFORM0 with two sets of vectors, embeddings vectors INLINEFORM1 and context vectors INLINEFORM2 , and we posit a probability distribution of data entries INLINEFORM3 in which these vectors interact. The definition of the efe model requires three ingredients: a context, a conditional exponential family, and a parameter sharing structure. We next describe these three components.
Exponential family embeddings learn the vector representation of objects based on the conditional probability of each observation, conditioned on the observations in its context. The context INLINEFORM0 gives the indices of the observations that appear in the conditional probability distribution of INLINEFORM1 . The definition of the context varies across applications. In text, it corresponds to the set of words in a fixed-size window centered at location INLINEFORM2 .
Given the context INLINEFORM0 and the corresponding observations INLINEFORM1 indexed by INLINEFORM2 , the distribution for INLINEFORM3 is in the exponential family, DISPLAYFORM0 
 with sufficient statistics INLINEFORM0 and natural parameter INLINEFORM1 . The parameter vectors interact in the conditional probability distributions of each observation INLINEFORM2 as follows. The embedding vectors INLINEFORM3 and the context vectors INLINEFORM4 are combined to form the natural parameter, DISPLAYFORM0 
 where INLINEFORM0 is the link function. Exponential family embeddings can be understood as a bank of glm. The context vectors are combined to give the covariates, and the “regression coefficients” are the embedding vectors. In Eq. EQREF6 , the link function INLINEFORM1 plays the same role as in glm and is a modeling choice. We use the identity link function.
The third ingredient of the efe model is the parameter sharing structure, which indicates how the embedding vectors are shared across observations. In the standard efe model, we use INLINEFORM0 and INLINEFORM1 for all columns of INLINEFORM2 . That is, each unique object INLINEFORM3 has a shared representation across all instances.
The objective function. In efe, we maximize the objective function, which is given by the sum of the log-conditional likelihoods in Eq. EQREF5 . In addition, we add an INLINEFORM0 -regularization term (we use the notation of the log Gaussian pdf) over the embedding and context vectors, yielding DISPLAYFORM0 
Note that maximizing the regularized conditional likelihood is not equivalent to maximum a posteriori. Rather, it is similar to maximization of the pseudo-likelihood in conditionally specified models BIBREF26 , BIBREF10 .
section_name: Structured Exponential Family Embeddings
paragraphs: Here, we describe the sefe model for grouped data. In text, some examples of grouped data are Congressional speeches grouped into political parties or scientific documents grouped by discipline. Our goal is to learn group-specific embeddings from data partitioned into INLINEFORM0 groups, i.e., each instance INLINEFORM1 is associated with a group INLINEFORM2 . The sefe model extends efe to learn a separate set of embedding vectors for each group.
To build the sefe model, we impose a particular parameter sharing structure over the set of embedding and context vectors. We posit a structured model in which the context vectors are shared across groups, i.e., INLINEFORM0 (as in the standard efe model), but the embedding vectors are only shared at the group level, i.e., for an observation INLINEFORM1 belonging to group INLINEFORM2 , INLINEFORM3 . Here, INLINEFORM4 denotes the embedding vector corresponding to group INLINEFORM5 . We show a graphical representation of the sefe in Figure FIGREF1 .
Sharing the context vectors INLINEFORM0 has two advantages. First, the shared structure reduces the number of parameters, while the resulting sefe model is still flexible to capture how differently words are used across different groups, as INLINEFORM1 is allowed to vary. Second, it has the important effect of uniting all embedding parameters in the same space, as the group-specific vectors INLINEFORM4 need to agree with the components of INLINEFORM5 . While one could learn a separate embedding model for each group, as has been done for text grouped into time slices BIBREF16 , BIBREF17 , BIBREF18 , this approach would require ad-hoc postprocessing steps to align the embeddings.
When there are INLINEFORM0 groups, the sefe model has INLINEFORM1 times as many embedding vectors than the standard embedding model. This may complicate inferences about the group-specific vectors, especially for groups with less data. Additionally, an object INLINEFORM2 may appear with very low frequency in a particular group. Thus, the naïve approach for building the sefe model without additional structure may be detrimental for the quality of the embeddings, especially for small-sized groups. To address this problem, we propose two different methods to tie the individual INLINEFORM3 together, sharing statistical strength among them. The first approach consists in a hierarchical embedding structure. The second approach is based on amortization. In both methods, we introduce a set of global embedding vectors INLINEFORM4 , and impose a particular structure to generate INLINEFORM5 from INLINEFORM6 .
Hierarchical embedding structure. Here, we impose a hierarchical structure that allows sharing statistical strength among the per-group variables. For that, we assume that each INLINEFORM0 , where INLINEFORM1 is a fixed hyperparameter. Thus, we replace the efe objective function in Eq. EQREF7 with DISPLAYFORM0 
where the INLINEFORM0 -regularization term now applies only on INLINEFORM1 and the global vectors INLINEFORM2 .
Fitting the hierarchical model involves maximizing Eq. EQREF11 with respect to INLINEFORM0 , INLINEFORM1 , and INLINEFORM2 . We note that we have not reduced the number of parameters to be inferred; rather, we tie them together through a common prior distribution. We use stochastic gradient ascent to maximize Eq. EQREF11 .
Amortization. The idea of amortization has been applied in the literature to develop amortized inference algorithms BIBREF14 , BIBREF15 . The main insight behind amortization is to reuse inferences about past experiences when presented with a new task, leveraging the accumulated knowledge to quickly solve the new problem. Here, we use amortization to control the number of parameters of the sefe model. In particular, we set the per-group embeddings INLINEFORM0 to be the output of a deterministic function of the global embedding vectors, INLINEFORM1 . We use a different function INLINEFORM2 for each group INLINEFORM3 , and we parameterize them using neural networks, similarly to other works on amortized inference BIBREF28 , BIBREF29 , BIBREF30 , BIBREF31 . Unlike standard uses of amortized inference, in sefe the input to the functions INLINEFORM4 is unobserved and must be estimated together with the parameters of the functions INLINEFORM5 .
Depending on the architecture of the neural networks, the amortization can significantly reduce the number of parameters in the model (as compared to the non-amortized model), while still having the flexibility to model different embedding vectors for each group. The number of parameters in the sefe model is INLINEFORM0 , where INLINEFORM1 is the number of groups, INLINEFORM2 is the dimensionality of the embedding vectors, and INLINEFORM3 is the number of objects (e.g., the vocabulary size). With amortization, we reduce the number of parameters to INLINEFORM4 , where INLINEFORM5 is the number of parameters of the neural network. Since typically INLINEFORM6 , this corresponds to a significant reduction in the number of parameters, even when INLINEFORM7 scales linearly with INLINEFORM8 .
In the amortized sefe model, we need to introduce a new set of parameters INLINEFORM0 for each group INLINEFORM1 , corresponding to the neural network parameters. Given these, the group-specific embedding vectors INLINEFORM2 are obtained as DISPLAYFORM0 
 We compare two architectures for the function INLINEFORM0 : fully connected feed-forward neural networks and residual networks BIBREF32 . For both, we consider one hidden layer with INLINEFORM1 units. Hence, the network parameters INLINEFORM2 are two weight matrices, DISPLAYFORM0 
 i.e., INLINEFORM0 parameters. The neural network takes as input the global embedding vector INLINEFORM1 , and it outputs the group-specific embedding vectors INLINEFORM2 . The mathematical expression for INLINEFORM3 for a feed-forward neural network and a residual network is respectively given by DISPLAYFORM0 
 where we have considered the hyperbolic tangent nonlinearity. The main difference between both network architectures is that the residual network focuses on modeling how the group-specific embedding vectors INLINEFORM0 differ from the global vectors INLINEFORM1 . That is, if all weights were set to 0, the feed-forward network would output 0, while the residual network would output the global vector INLINEFORM2 for all groups.
The objective function under amortization is given by DISPLAYFORM0 
We maximize this objective with respect to INLINEFORM0 , INLINEFORM1 , and INLINEFORM2 using stochastic gradient ascent. We implement the hierarchical and amortized sefe models in TensorFlow BIBREF33 , which allows us to leverage automatic differentiation.
Example: structured Bernoulli embeddings for grouped text data. Here, we consider a set of documents broken down into groups, such as political affiliations or scientific disciplines. We can represent the data as a binary matrix INLINEFORM0 and a set of group indicators INLINEFORM1 . Since only one word can appear in a certain position INLINEFORM2 , the matrix INLINEFORM3 contains one non-zero element per column. In embedding models, we ignore this one-hot constraint for computational efficiency, and consider that the observations are generated following a set of conditional Bernoulli distributions BIBREF2 , BIBREF10 . Given that most of the entries in INLINEFORM4 are zero, embedding models typically downweigh the contribution of the zeros to the objective function. BIBREF2 use negative sampling, which consists in randomly choosing a subset of the zero observations. This corresponds to a biased estimate of the gradient in a Bernoulli exponential family embedding model BIBREF10 .
The context INLINEFORM0 is given at each position INLINEFORM1 by the set of surrounding words in the document, according to a fixed-size window.
Example: structured Poisson embeddings for grouped shopping data. efe and sefe extend to applications beyond text and we use sefe to model supermarket purchases broken down by month. For each market basket INLINEFORM0 , we have access to the month INLINEFORM1 in which that shopping trip happened. Now, the rows of the data matrix INLINEFORM2 index items, while columns index shopping trips. Each element INLINEFORM3 denotes the number of units of item INLINEFORM4 purchased at trip INLINEFORM5 . Unlike text, each column of INLINEFORM6 may contain more than one non-zero element. The context INLINEFORM7 corresponds to the set of items purchased in trip INLINEFORM8 , excluding INLINEFORM9 .
In this case, we use the Poisson conditional distribution, which is more appropriate for count data. In Poisson sefe, we also downweigh the contribution of the zeros in the objective function, which provides better results because it allows the inference to focus on the positive signal of the actual purchases BIBREF10 , BIBREF2 .
section_name: Empirical Study
paragraphs: In this section, we describe the experimental study. We fit the sefe model on three datasets and compare it against the efe BIBREF10 . Our quantitative results show that sharing the context vectors provides better results, and that amortization and hierarchical structure give further improvements.
Data. We apply the sefe on three datasets: ArXiv papers, U.S. Senate speeches, and purchases on supermarket grocery shopping data. We describe these datasets below, and we provide a summary of the datasets in Table TABREF17 .
ArXiv papers: This dataset contains the abstracts of papers published on the ArXiv under the 19 different tags between April 2007 and June 2015. We treat each tag as a group and fit sefe with the goal of uncovering which words have the strongest shift in usage. We split the abstracts into training, validation, and test sets, with proportions of INLINEFORM0 , INLINEFORM1 , and INLINEFORM2 , respectively.
Senate speeches: This dataset contains U.S. Senate speeches from 1994 to mid 2009. In contrast to the ArXiv collection, it is a transcript of spoken language. We group the data into state of origin of the speaker and his or her party affiliation. Only affiliations with the Republican and Democratic Party are considered. As a result, there are 83 groups (Republicans from Alabama, Democrats from Alabama, Republicans from Arkansas, etc.). Some of the state/party combinations are not available in the data, as some of the 50 states have only had Senators with the same party affiliation. We split the speeches into training ( INLINEFORM0 ), validation ( INLINEFORM1 ), and testing ( INLINEFORM2 ).
Grocery shopping data: This dataset contains the purchases of INLINEFORM0 customers. The data covers a period of 97 weeks. After removing low-frequency items, the data contains INLINEFORM1 unique items at the 1.10upc (Universal Product Code) level. We split the data into a training, test, and validation sets, with proportions of INLINEFORM2 , INLINEFORM3 , and INLINEFORM4 , respectively. The training data contains INLINEFORM5 shopping trips and INLINEFORM6 purchases in total.
For the text corpora, we fix the vocabulary to the 15k most frequent terms and remove all words that are not in the vocabulary. Following BIBREF2 , we additionally remove each word with probability INLINEFORM0 , where INLINEFORM1 is the word frequency. This downsamples especially the frequent words and speeds up training. (Sizes reported in Table TABREF17 are the number of words remaining after preprocessing.)
Models. Our goal is to fit the sefe model on these datasets. For the text data, we use the Bernoulli distribution as the conditional exponential family, while for the shopping data we use the Poisson distribution, which is more appropriate for count data.
On each dataset, we compare four approaches based on sefe with two efe BIBREF10 baselines. All are fit using sgd BIBREF34 . In particular, we compare the following methods:

Figures and tables:
caption: Figure 1: (a) INTELLIGENCE is used differently across the ArXiv sections. Words with the closest embedding to the query are listed for 5 sections. (The embeddings were obtained by fitting an amortized S-EFE.) The method automatically orders the sections along the horizontal axis by their similarity in the usage of INTELLIGENCE. See Section 3 additional for details. (b) Graphical representation of S-EFE for data in S categories. The embedding vectors ρ(s)v are specific to each group, and the context vectors αv are shared across all categories.
caption: Table 1: Group structure and size of the three corpora analyzed in Section 3.
caption: Table 2: Test log-likelihood on the three considered datasets. S-EFE consistently achieves the highest held-out likelihood. The competing methods are the global EFE, which can not capture group variations, and the separate EFE, which cannot share information across groups.
caption: Table 3: List of the three most different words for different groups for the Congressional speeches. S-EFE uncovers which words are used most differently by Republican Senators (red) and Democratic Senators (blue) from different states. The complete table is in the Appendix.

Questions and answers:
question: Do they evaluate on English only datasets?
answers: answer: unanswerable: yes
answer: unanswerable: no
yes_no: no
question: What experiments are used to demonstrate the benefits of this approach?
answers: answer: unanswerable: no
extractive_spans: On each dataset, we compare four approaches based on sefe with two efe BIBREF10 baselines. All are fit using sgd BIBREF34 . In particular, we compare the following methods:
answer: unanswerable: no
free_form_answer: Calculate test log-likelihood on the three considered datasets
question: What hierarchical modelling approach is used?
answers: answer: unanswerable: no
extractive_spans: the group-specific embedding representations are tied through a global embedding
question: How do co-purchase patterns vary across seasons?
answers: answer: unanswerable: yes
question: Which words are used differently across ArXiv?
answers: answer: unanswerable: no
extractive_spans: intelligence