Scalable Probabilistic Tensor Factorization for Binary and Count Data

1 Introduction

Tensor factorization methods [Kolda and Bader, 2009] offer a useful way to learn latent factors from complex multiway data. These methods decompose the original tensor data into a set of factor matrices (one for each mode or "way" of the tensor), which can be used as a latent feature representation for the objects in each of the tensor mode, and can be used for other tasks, such as tensor completion. Among tensor factorization methods, probabilistic approaches [Chu and Ghahramani, 2009; Xu et al., 2013; Rai et al., 2014] are especially appealing because of a proper generative model of the data, which allows modeling different data types and handling missing data in a natural way.

Real-world tensor data are often binary- or countvalued [Nickel et al., 2011; Chi and Kolda, 2012]. For example, a multirelational social network [Nickel et al., 2011] can be described as a three-way binary tensor with two modes denoting people and the third mode denoting the types of relationships. Likewise, from a database of research publications, one may construct a three-way (AUTHORS×WORDS× VENUES) count-valued tensor, where the three dimensions could be authors, words, and publication venues and each entry of the tensor denotes the number of time an author used a specific word at a specific venue. Tensor factorization on this multiway data can be used for topic modeling on such a publications corpus (the latent factors would correspond to topics). Another application could be in recommender systems; having learned the latent factors of authors and venues, one can use these factors for author-author recommendation (for potential co-authors) or author-venue recommendation (recommending the most appropriate venues for a given author).

Although several tensor factorization methods have been proposed in the recent years [Kolda and Bader, 2009; Chu and Ghahramani, 2009] and there has been a significant recent interest on developing scalable tensor factorization methods [Kang et al., 2012; Inah et al., 2015; Papalexakis et al., 2012; Beutel et al., 2014], most of these methods treat data as real-valued, and are therefore inappropriate for handling binary and count data; also see Related Work (Section 7). Motivated by the prevalence of such binary and count-valued tensors, we present a probabilistic tensor factorization framework which can handle binary and count-valued tensors, while being scalable for massive tensor data. Our starting point will be a conjugate, fully Bayesian model for both binary and count data, for which we develop an efficient Gibbs sampler. The framework is based on the P´olya-Gamma data augmentation strategy [Polson et al., 2012] which allows transforming likelihoods such as logistic (binary data) and negative binomial (count data) into a Gaussian. We then develop Expectation Maximization (EM), and, subsequently, an efficient online EM algorithm [Capp´e and Moulines, 2009] which allows scaling up to massive tensor data. Our online EM algorithm performs comparably or better than Gibbs sampling and EM, while being considerably faster to run (and requiring much less storage). Although here we focus on binary and count data, the framework can also handle other types of likelihoods, e.g., multinomial (for categorical data). We present experimental results on various data sets, including a publications corpus, and a massive transactions data set consisting of food-item purchase records from a large set of households.

2 Canonical PARAFAC (CP) Decomposition

Given a tensor Y of size n1×n2×· · ·×nK, with nk denoting the size of Y along the kth mode (or "way") of the tensor, the goal in a Canonical PARAFAC (CP) decomposition [Kolda and Bader, 2009] is to decompose Y into a set of K factor matrices U(1) , . . . ,U(K) where U(k) = [u(k) 1 , . . . ,u(k) R ], k = {1, . . . , K}, denotes the nk ×R factor matrix associated with mode k. In its most general form, CP decomposition expresses the tensor Y via a weighted sum of R rank-1 tensors

where f specifies the noise model and depends on the data being modeled (e.g., f can be Gaussian for realvalued, Bernoulli-logistic for binary-valued, Poisson for count-valued tensors); λr is the weight associated with the rth rank-1 component; the nk × 1 column vector u(k) r represents the rth latent factor of mode k; and ⊙ denotes vector outer product. We use subscript i = {i1, . . . , iK} to denote the K-dimensional index of the i-th entry in the tensor Y. The data are given as a set of N observations {yi}Ni=1 from the tensor and the goal is to perform CP decomposition (1) using these observations. Note that (1) can also be written, individually, for each observation yi as

where ψi = ∑Rr=1 λr ∏Kk=1 u(k) ikr is a real-valued number.

3 CP Decomposition for Logit Models

In this paper, we develop efficient tensor decomposition methods for data where the likelihood (2) is of the form

This general form subsumes several likelihood models [Polson et al., 2012]. For example, for mi = 1, we get the logistic likelihood for binary data; and for mi = yi + ξ, the negative binomial likelihood for count data (where ξ is the overdispersion parameter). In this exposition, we will focus on these two types of data (binary and count), which are the most prevalent in applications of tensor decomposition methods, but the framework we develop can be straightforwardly extended to handle other data types such as categorical data via multinomial likelihood [Polson et al., 2012].

A key thing to notice about the likelihood models of the form (3) is that, using the P´olya-Gamma augmentation scheme [Polson et al., 2012] allows "Gaussian-ifying" the likelihood given in (3). Specifically, for each observation yi, introducing a P´olya-Gamma random variable ωi [Polson et al., 2012] results in the following

where κi = yi - mi/2 and ωi ∼ PG(mi, ψi) denotes a draw from a P´olya-Gamma distribution. Since the form of the likelihood on the r.h.s. of (4) is Gaussian, this results in a conjugate model when the model parameters also have Gaussian prior distributions. In particular, in probabilistic CP decomposition models of the form (1), the model parameters u(k) r and λr are usually given Gaussian prior distributions [Xiong et al., 2010; Rai et al., 2014], and therefore the P´olya-Gamma augmentation strategy naturally leads to fully conjugate Bayesian analysis even though the original likelihood 3 is non-conjugate. This property was recently exploited in [Rai et al., 2014]; however, their work has the following limitations: (1) only the case of binary data was considered, and (3) inference in the model was done using batch Gibbs sampling which can be infeasible for massive tensors.

In contrast, in this paper, we present a scalable framework for CP decomposition, based on P´olya-Gamma augmentation that applies to not just binary but also count data (and extendable to other data types such as categorical). For both these cases, we present scalable inference algorithms to handle massive tensor data, for which Gibbs sampling [Rai et al., 2014] methods do not scale. Our inference algorithms are based on online Expectation Maximization [Capp´e and Moulines, 2009] which allows processing data in minibatches, and easily scales to massive tensors.

4 P´olya-Gamma CP for Binary &Count Data

We first describe the general CP decomposition model with P´olya-Gamma augmentation for likelihoods of the form

where ψi = ∑Rr=1 λr ∏Kk=1 u(k) ikr is a real-valued number. We further assume a Gaussian prior distribution on each of the factor matrices U(k) = [u(k) 1 , . . . ,u(k) R ]

and assume that the coefficients λr's are drawn from the multiplicative gamma process (MGP) prior [Bhattacharya and Dunson, 2011; Rai et al., 2014], with a truncation-level R over the rank of the decomposition. The MGP prior shrinks the λr's for the unnecessary factors to close to zero, thereby inferring the effective tensor rank from the data, and has the form

Using the fact that (5) can be re-expressed as a Gaussian using the transformation (4), and the fact that the priors on the model parameters u(k) r 's and the λr's are also Gaussian, we can derive a fully Bayesian inference procedure using Gibbs sampling. In [Rai et al., 2014], the Gibbs sampler for binary tensors was described, so we skip it here for brevity; therefore, and in Section 4.1, we focus only on the Gibbs sampler for the count data case, which was not considered in [Rai et al., 2014]. We will refer to our model as PGCP.

4.1 Gibbs Sampler for PGCP with Count Data

We first re-express ψi = ∑Rr=1 λr ∏Kk=1 u(k) ikr as

where c(k) ikr and d(k) ikr are defined as c(k) ikr = λr ∏k′ ̸=k u(k′ ) ik′ r and d(k) ikr = ∑r′ ̸=r λr′ ∏Kk=1 u(k) ikr′ . Using (8), (4), uikr ∼ Nor(0, 1), and using results from Gaussian-Gaussian conjugacy, the Gibbs sampling updates for u(k) nr , 1 ≤ n ≤ nk, for the case of count-valued tensor data, are given as u(k) nr ∼ Nor(α(k) nr , τ (k) nr -1 ), with the precision and mean defined as

In a similar manner, writing ψi = (∏Kk=1 u(k) ikr )λr + (∑r′ ̸=r λr′ ∏Kk=1 u(k) ikr′) = ari λr + bri , using (4) and the fact λr ∼ N(0, τ -1 r ), the Gibbs sampling updates for λr, 1 ≤ r ≤ R are λr ∼ Nor(ˆµr, ˆ τ -1 r ), with the precision and mean defined as

The other MGP parameters (the δl's) can be sampled easily from a Gamma posterior, following [Rai et al., 2014]. We skip the details here for brevity. Sampling the overdispersion parameter ξ: In the negative binomial model for counts, we also need to estimate the overdispersion parameter ξ. We assume a prior ξ ∼ Gamma(a0, 1/h) and, following [Zhou et al., 2012b], the Gibbs sampling updates for ξ are

where Li ∼ Poisson(ξ log(1 + exp(1 + exp(ψi)).

5 Inference via ExpectationMaximization

Although the Gibbs sampler described in Section 4.1 is efficient for moderate-sized count-valued tensors, running it on massive tensor data can be infeasible, both due to high storage as well as high computational costs.

In this section, we first describe an ExpectationMaximization (EM) algorithm for parameter estimation in the PGCP model (for both binary- as well as count-valued tensor data), and then describe an online EM algorithm that can process data in mini-batches and scale to massive-sized tensors.

To derive the EM algorithm for the PGCP model, we need to write down the expected complete data (log)-likelihood for each of the model parameters. The latent variables here are the P´olya-Gamma variables ωi's. Luckily, the expectation of a PG random variable ω ∼ PG(b, c) is available in closedform [Polson et al., 2012], and is given by

In the case of binary data where ωi ∼ PG(1, ψi), the expectation is given by E[ωi] = (0.5/ψi) tanh(ψi/2), whereas in the case of count data where ωi ∼ PG(yi + ξ, ψi), the expectation is given by E[ωi] = (0.5(yi + ξ)/ψi) tanh(ψi/2). In the rest of this section, we will abuse notation and simply use ωi in place of E[ωi]. Also, in what follows, κi is defined as κi = yi -mi/2, where mi = 1 for binary data and mi = yi + ξ for count data. With these ingredients, the EM algorithm for the PGCP model then proceeds as follows: EStep: Compute expectations of ω1, . . . , ωN. M-Step: Maximize the expected complete data log-likelihood w.r.t. each model parameter. For the u(k) nr (n = 1, . . . , nk), it follows from (8) that the expected complete data log-likelihood Q(u(k) nr ) is given by the following quadratic

Further, we can vectorize the above and write the complete data log-likelihood for u(k) r = [u(k) i1r , . . . , u(k) nkr]⊤ as

where S(r,k) is a diagonal matrix, and S(r,k) and t(r,k) ∈ Rnk are defined as S(r,k) nn = ∑i:ik=n c2ikr ωi and t(r,k) n = ∑i:ik=n (κi -dikrωi)cikr. Note that since S(r,k) is diagonal, maximizing (13) w.r.t. u(k) r is very efficient with no matrix inversions needed. Also note that, it is possible to incorporate a prior distribution on u(k) r (e.g., Gaussian, sparsity, etc.)

Alternatively, one can solve a constraint optimization problem for the objective (13) by placing any desired constraint on u(k) r (e.g., sparsity, non-negativity, etc.).

Likewise, the expected complete data log-likelihood for the λ = (λ1, . . . , λR) is

where Sλ is again an R×R diagonal matrix with diagonal entry given by Sλr = ∑i ari 2ωi and tλ ∈ RR is defined as tλr = ∑i (κi - bri ωi)ari .

In the case of count data, we need to estimate the overdispersion parameter ξ of the negative binomial. This can be done using the Newton-Raphson procedure in the M-step.

6 Online ExpectationMaximization

The batch EM algorithm presented in Section 5 is more efficient than the batch Gibbs sampling presented in Section 4.1. In many cases, however, the data can be too massive to handle even for batch EM. In this section, we present an online EM algorithm [Capp´e and Moulines, 2009; Scott and Sun, 2013] that works by processing data in minibatches and therefore has the capability to scale to massive tensor data (in terms of the tensor dimensions and/or in terms of the number of tensor observations).

Note that the batch EM algorithm presented in Section 5 operates on the sufficient statistics (e.g., S(r,k) , t(r,k) , Sλ, tλ, etc.) of the model parameters. These sufficient statistics are computed using all the observations from the tensor. Noting that these statistics are defined as summation over the individual observations from the tensor, it is possible to recursively update these sufficient statistics as more and more data arrives. This eliminates the need of having to store all of the tensor observations in the memory and computing these statistics using the entire data in each round of the EM algorithm. Instead, with the online EM algorithm [Capp´e and Moulines, 2009], these statistics can be computed using the current mini-batch of the data, and added to the previous statistics. We leverage this idea to develop an efficient, online EM algorithm for the proposed PGCP model.

In the online EM algorithm, we express the sufficient statistics at time t as a convex combination of the old statistics from time t - 1 and contribution from the new mini-batch of data. For example, for the parameters u(k) r , the estimates of the sufficient statistics S(r,k,t) nn and t(r,k,t) n after seeing a new minibatch It of data at round t is given by

The sufficient statistics of the other parameters can also be updated in a similar manner, and used in the M step.

We note that, as shown in [Capp´e and Moulines, 2009], updating the sufficient statistics recursively as above, and using these in the M step, is equivalent to a stochastic gradient ascent-style updates of the parameters; e.g., for the parameter u(k) r , it would be equivalent to the following update

where I(.) denotes the Fisher Information Matrix. The learning rate γt ∝ (t + t0)-c for some c ∈ (0.5, 1] is chosen to ensure that ∑∞t=1 γt = ∞ and ∑∞t=1 γ 2 t < ∞; these conditions are necessary for the convergence of online EM [Capp´e and Moulines, 2009]. Following [Capp´e and Moulines, 2009], we set c to be close to 0.5, which worked well in practice in our experiments. Moreover, as recommended in [Capp´e and Moulines, 2009; Scott and Sun, 2013], we also smooth the parameter estimates using PolyakRuppert averaging where the parameters are averaged over the final T - T0 where T0 is some large enough value such that the algorithm has settled down.

Finally, note that, although motivated differently, our online EMmirrors stochastic variational inference [Hoffman et al., 2013] methods, where the parameters of the variational approximation of the posterior distribution are updated recursively in a similar fashion with each incoming minibatch of data.

7 RelatedWork

Tensor factorization methods have become increasingly prevalent nowadays, thanks to the rise of multirelational and multiway data sets in a wide variety of application domains. While much of the work on tensor factorization methods assumes the tensor entries to be real-valued, some generalizations to data from exponential family distributions [Hayashi et al., 2010], generalized linear models [Yılmaz et al., 2011], or specifically for binary-valued tensors [Nickel et al., 2011; Rai et al., 2014] have been also proposed. However, these models can be computationally infeasible to run on massive tensors (e.g., the method in [Hayashi et al., 2010] requires tensor unfolding operations which can easily render it impractical for large tensors) as they all require batch processing. For count-valued tensor data, [Chi and Kolda, 2012] proposed a Poisson tensor factorization method, but this method cannot handle missing values in the tensor, and can therefore not be used for tensor completion. Moreover, these methods also have other limitations such as the need to pre-specify the rank of decomposition a priori, or to select it via crossvalidation.

With the advent of massive-scale tensor data in applications such as knowledge-bases and social networks, there has been a considerable recent interest on scaling up tensor factorization methods by developing distributed and parallel methods [Kang et al., 2012; Inah et al., 2015; Papalexakis et al., 2012; Beutel et al., 2014]. Most of these methods, however, have one or more of the following limitations: (1) they assume a parallel or distributed setting, (2) missing data cannot be easily handled/predicted, (3) they lack a proper generative model of the data (most of these assume real-valued data), and are not designed for a Bayesian setting, and (4) the rank of the decomposition needs to be chosen via crossvalidation. In contrast, our proposed framework does not have any of these limitations.

In [Hayashi et al., 2012], the authors proposed an EM algorithm for exponential family tensor factorization, using Laplace and Gaussian process approximations, and extended it to online EM; however, the proposed algorithm, based on tensor unfolding, has computational complexity that is cubic in the sum of the tensor dimensions, and is prohibitive even for moderate-size tensors. In contrast, our algorithms have per-iteration complexity that is linear in the tensor dimension as well as linear in the number of observations (or in the minibatch size in the case of online EM) in the tensor, and therefore can scale up nicely even for massive tensor data.

Finally, scalable Bayesian tensor decomposition is a relatively under-explored area. Some recent attempts include [Zhe et al., 2015; Rai et al., 2014; Zhe et al., 2013]. However, these are limited to real-valued and/or binaryvalued tensors (except for the work in [Zhou et al., 2014], which assumes categorical data), and scale poorly w.r.t the size of the tensor and/or rely on batch inference. In contrast, our online EM based framework applies to binary and count data (real-valued data can be trivially handled without the need for P´olya-Gamma augmentation), and extends easily to other data types, such as categorical data (via multinomial likelihood), or data with binomial likelihoods.

8 Experiments

We experiment with several real-world data sets that consist of both binary and count tensors. These include:

• Kinship: This is a 104×104×26 binary tensor [Nickel et al., 2011] consisting of 26 types of relations among a set of 104 individuals.
• Digg: This is a 581×124×48 binary tensor [Zhe et al., 2013] with 0.024% nonzeros.
• DBLP: This is a 10, 000 × 200 × 10, 000 binary tensor [Zhe et al., 2015] constructed from the DBLP database, and contains author-conference-keyword relations (a tensor entry is 1 if the relation holds). This data has 0.001% nonzeros.
• GDELT: This is a 220×220×20×52 count tensor [Leetaru and Schrodt, 2013] consisting of country-country interactions based on 20 types of actions; we use a subset of this data collected over a period of one year (year 2012).
• Scholars: This is a 2370×8663×4066 count tensor constructed from a publications database of Duke University 1 , where each entry of the tensor denotes the number of times an authors uses a word at a publication venue.
• Transactions: This is a 117, 054×438×67, 095 count tensor constructed from a database of transactions data of food-item purchases at various stores in the US 2 ; the three tensor modes correspond to households, stores, and items purchased.

In the experiments, we refer to our model PGCP (P´olyaGamma CP decomposition) with Gibbs, EM and online EM inference, as PGCP-Gibbs, PGCP-EM, and PGCP-Online EM, respectively. For binary data, we compare with MGPCP [Rai et al., 2014], and (2) Infinite Tucker Factorization [Xu et al., 2013]. For count data, we compare with: (1) PTF [Chi and Kolda, 2012] 3 which assumes a Poisson factorization model, and (2) LRA-NTD[Zhou et al., 2012a] which is a non-negative tensor factorization method but data is assumed to be real-valued.

Experimental settings: All the methods use random initialization. For the online EM, we set the mini-batch size to onetenth of the total data size. For our methods, the upper bound R on the tensor rank was set to 10 for DBLP data and Kinship data; Digg data, R was set to 5 (larger values yield similar results, albeit at a greater computational cost). For GDELT, Scholar, and Transactions data sets, R was set to 100. Unless specified otherwise, the Gibbs sampler was run for 1000 iterations with 500 burn-in iterations. Unless specified otherwise, the online EM and EM algorithms were run for 1000 iterations. All experiments are done on a standard desktop with Intel i7 3.4GHz processor and 24GB RAM.

1 https://scholars.duke.edu/

2 The data is provided to us by USDA.

3 Since the method in [Chi and Kolda, 2012] cannot handle missing data, we implement a fully Bayesian version of the model and use it in our experiments

8.1 Tensor Completion

We first experiment on the task of tensor completion where we compare all the methods for their ability to predict missing entries in the tensor. For binary tensors, we split the data into 80% training and 20% test, while ensuring that the nonzeros are also split the same way between the training and test data. For count data, we use 95% data as the training data and use the remaining 5% as the heldout data (in Section 8.4, we also show an additional experiment with varying fractions of the training data). The tensor completion results are shown in Table 1 for binary tensor data, and Table 2 for count tensor data. Each result is averaged over five random splits of the training and test data. On both binary as well as count data, the PGCP (all the three variants) outperform the other baselines (except Digg where Inf-Tucker baseline does marginally better). Moreover, the EM and online EM variants of PGCP perform at par or better than Gibbs sampling, suggesting that on massive data sets, the Gibbs sampler, in addition to running slower, may also be mixing poorly (in Section 8.2, we compare the three variants of PGCP in terms of running times).

8.2 Scalability

In this section, we assess the performance of the online EM as compared to the batch counterparts (batch EM and batch Gibbs) in terms of its convergence speed vs the wall-clock run time. Each method was run for a specified amount of time (long enough that it roughly converges to a stable AUC or log-likelihood score). The results from various data sets are shown in Figures (1) and (2). As these figures show, the online EM algorithm reaches reasonably good quality solutions fairly quickly as compared to batch EM and batch Gibbs. In most cases, it even attains better quality solutions than its batch counterparts. In particular, this can be attributed to the fact that the Gibbs sampler may exhibit slow mixing on large data sets. This experiment demonstrates that the PGCP model with the online EM algorithm can be a viable alternative to Gibbs sampling and batch EM on massive tensor data sets.

DBLP Data PGCP-Gibbs PGCP-EM

PGCP-Online EM

10 1

10 2

10 3

Time in seconds (on log scale)

Figure 1: (Left) Scalability on Kinship. (Right) Scalability on DBLP. Blue: PGCP-Gibbs, Green: PGCP-EM, Red: PGCP-Online EM

8.3 Topic Modeling Results

Our model can also be used for topic modeling of multiway data. On the Scholar data (AUTHORS × WORDS × VENUES tensor), we do an experiment with PGCP-Online EM by imposing non-negativity constraint on the factor matrices and examine the factor matrix of the words dimension of the tensor. Table 3 shows three of the factors and the top 8 words

Test AUC Score

0.95

0.9

0.85

0.8

0.75

0.7

0.65

0.6

0.55

0.5

0.45

in each of these factors (based on their factor scores). Looking at the top words in these factors suggests that these factors correspond to topics machine learning & statistics (ML), immunology, and cardiac imaging, which are representative subject areas of the Scholar publications corpus.

8.4 Varying Amounts of Training Data

We also investigate the behavior of PGCP with varying amounts of training data. For this experiment, we first split the Scholar data into 80% training and 20% test data, and then use a varying fraction of the training data (here we only report the results with PGCP-online EM). In Figure 3, we show results for (1) heldout log-likelihood vs fraction of training data, and (2) running-time vs fraction of training data.

As Figure 3 (left) shows, the heldout log-likelihood expectedly gets better with increasing training data; however, even with small fractions of the training data, the model achieves reasonably high log-likelihoods, suggesting its robustness.

Figure 3 (right) shows the total running time (summed over the total number of iterations) for each fraction of the training data. As shown, the total running time scales roughly linearly in the number of observations in the training data.

9 Conclusion and FutureWork

We have presented a scalable framework for tensor factorization that can handle binary and count data via a conjugate model obtained using a data augmentation strategy. Leveraging this conjugacy, our framework can perform a fully Bayesian analysis via closed-form Gibbs sampling, as well as allows efficient inference via EM and online EM algorithms. Our work opens doors to the a scalable and fully Bayesian analysis of binary- and count-valued tensors, which are now becoming increasingly prevalent in several applications. The work here can be extended in several interesting directions. Although, here, we only considered binary and count data, our framework, which is based on the P´olya-Gamma augmentation, can also be straightforwardly extended to handle categorical data, or data with binomial likelihoods.

Another interesting extension can be incorporating sources of auxiliary information for the tensor for improved tensor decomposition. For example, often data sources with side-information about each of the tensor modes are available [Yılmaz et al., 2011; Ermis¸ et al., 2013; Beutel et al., 2014; Rai et al., 2015; Khan and Kaski, 2014]. By sharing the factor matrices of such modes across these multiple data sources, such side-information can be effectively leveraged. To further scale up our framework, many steps of the Gibbs sampler in Section 4.1, as well as of the proposed EM and online EM algorithms, can be parallelized/distributed.

Acknowledgments

The research reported here was supported in part by ARO, DARPA, DOE, NGA and ONR. Any opinions, findings, recommendations, or conclusions are those of the authors and do not necessarily reflect the views of the Economic Research Service, U.S. Department of Agriculture. The analysis, findings, and conclusions expressed in this paper also should not be attributed to either Nielsen or Information Resources, Inc. (IRI). This research was conducted in collaboration with USDA under a Third Party Agreement with IRI.