Scalable Bayesian Non-linear Matrix Completion

1 Introduction

In matrix completion-one of the most widely used approaches for collaborative filtering-the objective is to predict missing elements of a partially observed data matrix. Such problems are often characterized by large and extremely sparsely observed data sets. The classic linear solution to the problem is to find a factorization of the data matrix Y ∈ RN×D as a product of latent variables X ∈ RN×K and weights W ∈ RD×K (K ≪ N,D), from which elements of Y can be predicted as Y ≈ XWT . Probabilistic matrix factorization (PMF) [Salakhutdinov and Mnih, 2008b], formulates the problem as a probabilistic model, regularized by placing priors on X and W, and finds the solution as a maximum a posteriori (MAP) estimate of these matrices. Fully Bayesian matrix factorization [Salakhutdinov and Mnih, 2008a] expands this model by further placing priors on model hyperparameters, and marginalizing these along with X and W. Bayesian matrix factorization brings the advantages of automatic complexity control and better robustness to overfitting. Moreover, the solution comes with an uncertainty estimate, which is useful when the completed matrix is used for decision making. For instance, sparsely observed drug-target interaction matrices are used for deciding which interactions to measure next.

∗Contact Author. Currently at Institute for Molecular Medicine Finland (FIMM), University of Helsinki.

Lawrence and Urtasun [2009] generalized PMF using a Gaussian process latent variable model (GP-LVM) formulation, where the relationship between X and Y is given by Y ≈ f(X), with a GP-prior placed over f . The X is optimized to find its MAP solution. Note that this formulation also subsumes the linear model as a special case. Subsequently, a variational inference framework for fully Bayesian GP-LVM has been developed [Damianou et al., 2016; Titsias and Lawrence, 2010], building on sparse GP approximations [Qui˜nonero Candela and Rasmussen, 2005; Snelson and Ghahramani, 2006]. It parametrizes the covariance matrix implied by the GP kernel using a set of M ≪ N auxiliary inducing variables. While Bayesian GP-LVM has been successfully used for dimensionality reduction and extracting latent representations, less consideration has been given to its applicability in matrix completion tasks with extremely sparse data. Computationally, this is a much more demanding problem, because the variational updates have to be performed separately for each dimension of the data matrix, instead of being performed as a single operation.

Existing approaches for scaling up Bayesian GP-LVM make use of the insight that, conditional on the inducing variables, the data points can be decoupled for parallel computations. In this line of work, Gal et al. [2014] introduced a distributed version of Bayesian GP-LVM. Dai et al. [2014] proposed a similar framework, additionally using GPU acceleration to speed up local computations. Neither of the works demonstrated learning of latent variable models beyond moderately-sized data, nor have they been implemented for sparse matrices, which is needed for the problems considered in this paper. More importantly, current distributed solutions require the worker nodes to communicate with the master node in every iteration, which leads to an accumulating communication overhead as the number of worker units increased with the size of the problem. Vander Aa et al. [2016] reported such a phenomenon for their distributed MPI implementation of Bayesian linear matrix factorization. Finally, our experience indicates that existing distributed implementations may suffer from high memory consumption.

For GP-regression models, with X observed, Deisenroth and Ng [2015] proposed a framework with particularly good scaling properties and efficient use of memory. This framework utilizes a product-of-GP-experts (PoE) formulation [Cao and Fleet, 2014; Ng and Deisenroth, 2014; Tresp, 2000], which makes predictions using a product of independent local models, each operating on a subset of the data. These types of approximations are amenable to embarrassingly parallel computations, and can therefore be scaled up to arbitrarily large data sets, at least in principle. However, a direct application of PoE for nonlinear matrix completion may not produce satisfactory predictions for two reasons. First, since the target matrix is very sparsely observed, each local model has very limited information to learn an informative model without sharing information. Second, while local models could be combined into larger models to improve predictive performance, this is hindered by the general nonuniqueness of the latent variables in latent variable models.

In this work, we propose a distributed computational strategy which is almost as communication-efficient as embarrassingly parallel computation, but enables local models to share information and avoids the problem of non-uniqueness in aggregating local models. In a nutshell, one data subset is processed first, and the rest of the embarrassingly parallel computations are conditioned on the result. A similar idea was recently presented by [Qin et al., 2019] for Bayesian linear matrix completion. The remainder of the paper proceeds as follows: In Section 2 we first provide a brief review of GP-LVMs. Then, in Section 3, we present our framework for scalable Bayesian non-linear matrix completion. An empirical evaluation of the method, using simulations and a benchmark dataset, is given in Section 4. The paper ends with conclusions in Section 5.

2 Gaussian Process Latent Variable Models

A Gaussian process latent variable model (GP-LVM) [Lawrence, 2005] can be constructed from a non-linear multioutput regression model,

by placing a GP prior over the unknown functions f1 . . . , fd. Integrating over the space of functions with respect to a zeromean GP then yields the likelihood as

where K is an N × N kernel matrix defined by a GP covariance function k(xs,:, xt,:). We use θ to collectively denote all parameters, including the noise variance σ2 and the parameters of the covariance function.

When values are missing, as is the case in matrix completion, each factor of the likelihood (1) will only account for observed elements, thus we have

where nd denotes the set of indices of observed elements in column d. Furthermore, the symmetric matrices K and I will only include rows and columns corresponding to the indices nd.

2.1 Bayesian GP-LVM

For Bayesian GP-LVM, we complement the likelihood (2) by placing a prior on X. A standard choice is to set

The marginal likelihood of Bayesian GP-LVM is obtained by integrating the model with respect to p(X):

While this operation is intractable in general, Titsias and Lawrence [2010] introduced a variational framework, which leads to a tractable lower bound,

on the log of the marginal likelihood (4). For a detailed treatment of the framework, see [Damianou et al., 2016]. As a by-product of optimizing the lower bound (5), we get a variational approximation q(X) to the posterior p(X|Y) ∝ p(Y|X)p(X), for which we assume a factorized Gaussian form

In our current work, we use this approximation to share information between parallel computations (see Section 3.1).

2.2 Extension toMulti-view Settings

Manifold relevance determination (MRD) [Damianou et al., 2012] extends GP-LVM to a multi-view setting by reformulating the likelihood (1) as ∏v∈V p(Yv |X, θv ), where the elements of V index the data views, i.e. matrices conditionally independent given a single latent matrix X. In matrix completion problems, one of the views is typically the target in which prediction is carried out, while the other views constitute side-data. When predicting values in completely unobserved (or new) rows or columns in the target, predictions are effectively done 'outside' of the observed matrix. This can be done with the help of observed data in the side-views, and is referred to as out-of-matrix prediction.

3 Scalable Matrix Completion Using Bayesian GP-LVM

This section presents a computational strategy, which enables Bayesian GP-LVM to be scaled up for large-scale matrix completion problems using a product of experts (PoE). In brief, we first partition the observation matrix Y into I disjoint subsets Y(1) ,Y(2) . . . ,Y(I) for parallel processing. To avoid problems resulting from the unidentifiability of latent variable models, we couple the subsets as follows: in an initial stage only one subset is processed; in the second stage, the remaining subsets are processed in parallel using incremental learning (Section 3.1). For each subset, we use an implementation of the variational inference framework for Bayesian GP-LVM [Titsias and Lawrence, 2010]. Finally, with a set of independently learned Bayesian GP-LVMs, we use PoE to predict unobserved matrix elements (Section 3.2).

The proposed strategy is summarized in Figure 1. In Section 3.3, we present a further variant of the method, which uses intermediate aggregation of the submodels, prior to PoE, to improve predictive performance. The scalability of our method is briefly discussed in Section 3.4.

3.1 Coupling Parallel Inferences Using Incremental Learning

To couple the parallel inferences over subsets in the second stage of our method, we use a probabilistic variant of the incremental learning algorithm introduced by [Yao et al., 2011] for online learning of (non-Bayesian) non-linear latent variable models. Let Y(1) be the submatrix processed in the initial stage. Furthermore, denote by Y(i) aug = [Y(1) ,Y(i) ] ∈ R(N1+Ni)×D , i = 2, . . . , I, the combined submatrix obtained by augmenting Y(i) with Y(1) . The corresponding combined latent matrix is denoted by X(i) aug = [X(1) ,X(i) ] ∈ R(N1+Ni)×K .

The objective of incremental learning is to learn the joint latent matrix X(i) aug without extensive relearning of X(1) , while still allowing it to be updated. When learning X(i) aug, Yao et al. [2011] added a regularizer to the log-likelihood to prevent X(1) from deviating too much from its initial estimate, and to speed up learning. The original incremental learning algorithm used the Frobenius norm ‖X(1) - ˆX(1) ‖2F to regularize the updated estimate of X(1) . In our current work, we use the KL-divergence KL (q(X(1) )||ˆ q(X(1) )) to ensure that the updated variational posterior approximation q(X(1) ) remains close to the initial approximation ˆ q(X(1) ). For X(i) , we use the default prior given in Equation (3). Thus, the variational inference for the incremental learning of Bayesian GP-LVM follows the procedure introduced in Section 2.1, with the KL terms in the lower bound of Eq. (5) taking the following form:

For the augmented subsets, the inducing points and kernel hyperparameters are initialized to values obtained in the initial stage. For initialization of latent variables, we use posterior means for X(1) and nearest neighbors for X(i) [Yao et al., 2011].

3.2 Prediction with Product of Experts

Product of experts (PoE) prediction for Gaussian process regression [Deisenroth and Ng, 2015] uses the simple idea of combining predictions made by independent GP-models (i.e. 'experts') as a product:

where x∗ is a given test input and y∗ is the corresponding output value to be predicted. Under this model, a prediction is proportional to a Gaussian with parameters

With latent variable models, the essential difference to the above is that the inputs are unobserved and must therefore be inferred. In matrix completion, we wish to predict the missing part of a partially observed test point y∗ = (yO ∗ , yU ∗ ) ∈ RD , where yO ∗ are the observed elements (or side views) and yU ∗ are the missing values (or the unobserved target view) to be reconstructed. The prediction task can be finished in two steps. First, we infer the latent input x∗ of the test point, which involves maximizing the variational lower bound on the marginal likelihood

to obtain the approximate posterior q(X, x∗) = q(X)q(x∗). The lower bound has the same form as the learning objective function in Equation (5), but for its maximization, the variational distribution q(X) over latent variables for training data and parameters θ remains fixed during test time. After obtaining q(x∗), making predictions for yU ∗ is approached as GP prediction with uncertain inputs [Damianou et al., 2016; Girard et al., 2003].

In our distributed setting, the experts in PoE correspond to submodels learned from the augmented training subsets formed in the incremental learning phase. To correct for the initial subset Y(1) being used in I - 1 training sets, we formulate a corrected PoE as follows:

Finally, denoting the means and variances of the local predictive distributions as ˆµi and ˆσ2 i , respectively, we compute the aggregated statistics of the corrected PoE predictions as:

In their distributed GP framework, Deisenroth and Ng [2015] used a re-weighted variant of PoE, which they coined the robust Bayesian committee machine (rBCM). Although rBCM has been shown to outperform the basic PoE for GPregression, in our current setup, we have not observed any advantage of it over PoE. We have therefore formulated our framework using standard PoE but note that the extension to rBCM is straightforward.

3.3 Improved Solution with Intermediate Aggregation

PoE aggregates predictions from local submodels learned on data subsets, effectively using a block-diagonal approximation of the full-data covariance matrix. With larger submodels, PoE provides a closer approximation to the full covariance matrix, which can be expected to result in better predictive accuracy. Here we introduce an intermediate aggregation strategy, by which submodels are aggregated for improved predictive performance, while the initial training of submodels is still done on smaller subsets with lower computational cost. While latent variable models are in general non-identifiable, making a direct aggregation of local models difficult to carry out in a meaningful way, the incremental learning introduced in Section 3.1, encourages identifiability among local models, alleviating the problem.

The aggregation of submodels involves (i) stacking together local variational distributions, which are assumed to be independent across subsets, (ii) concatenating the corresponding data subsets, and finally (iii) aggregating the hyperparameters of the models. The model parameters can be approximated using suitable statistics (e.g. mode, median or mean) of the distributions. In our implementation, we use the mode to approximate the kernel and Gaussian noise variance parameters, and use averaging to estimate inducing variables.

Since the first subset Y(1) is used multiple times through incremental learning, the corresponding variational distribution q(X(1) ) is obtained through the following aggregation:

where

Above, each of the variational distributions ˆ qi(X(1) ) is Gaussian, of the form given by Equation (6).

Note that after intermediate aggregation, each training subset is used only once to make predictions, and we may use the ordinary PoE formulation in Equation (7) for prediction.

3.4 Computational Cost

Our method aims to leverage the scaling properties of sparse GP for training and those of PoE for prediction. Thus, for data partitioned into subsets of size Ni, i = 1, . . . , I, and assuming that a sufficient number of parallel workers is available, the time complexity for training is O (maxi(NiM 2 ) · D), where M < Ni is the number of inducing points and D reflects the fact that variational updates have to be performed separately for each dimension of sparsely observed data. For prediction, the cost is O (maxi(N 2 i )). For incremental learning and intermediate aggregation, Ni refers to the size of the concatenation of multiple subsets. By intermediate aggregation of submodels, we are able to trade off prediction cost against accuracy.

4 Experiments

In this section, we evaluate the predictive performance of the proposed method for out-of-matrix prediction problems on simulated and real-world chemogenomic data, and compare it with two alternative approaches: (i) the embarrassingly parallel or subset of data (SoD) approach, which has been widely studied to scale up Gaussian Process regression models, and (ii) Macau, Bayesian multi-relational factorization with side information [Simm et al., 2015], supporting out-ofmatrix prediction. Macau is implemented with highly optimized C libraries, and is available for experiments on largescale data. The comparison with the SoD approach shows the advantage of our method in sharing information among submodels, while the comparison with Macau shows the benefit of using Bayesian non-linear matrix factorization. We emphasize, however, that the choice and effectiveness of a model always depends on the problem at hand.

Simulated data. We generated synthetic data using nonlinear signals corrupted with Gaussian noise, using matern data generator available in the GPy 1 software. The data has three views Y = {Yv : v = 1, 2, 3}, the dimension of the views is as follows: N = 25,000, D1 = 150, D2 = 100, D3 = 150. As the task is to perform out-of-matrix prediction, we randomly selected 20% of the rows as a test set, using the remaining rows as the training set. In addition, 80% of the data in the first view were masked as missing, to simulate the sparsely observed target data in a real-world application. The other two views were regarded as side information and were fully observed. Unlike Bayesian GP-LVM, Macau cannot handle missing values in the side data.

Real-world data. We performed the experiments on ExCAPE-DB data [Sun et al., 2017], which is an aggregation of public compound-target bioactivity data and describes interactions between drugs and targets using the pIC50 2 measure. It has 969,725 compounds and 806 targets with 76,384,863 observed bioactivities. The dataset has 469 chem2vec features as side information which are generated from ECFP fingerprint features for the compounds using word2vec software. We used 3-fold cross validation to split the training and test set, where about 30% of the rows or compounds were chosen as test set in each fold.

4.1 Experimental Setup

The experimental setting for MRD models is: number of inducing points 100, optimization through scaled conjugate gradients (SCG) with 500 iterations. For the SoD approach, the latent variables were initialized with PPCA method. We ran Macau with Gibbs sampling for 1200 iterations, discarded the first 800 samples as burn-in and saved every second of the remaining samples yielding in total 200 posterior samples. We set the dimension of latent variables K=10 for ExCAPEDB data, K=5 for simulated data for all methods.

For the proposed and SoD methods, we partitioned the simulated data into 10x1 subsets and ExCAPE-DB data into 400x1 subsets. Other partitions are also possible; we have chosen the size of subsets such that Bayesian inference could be performed for the subsets in reasonable time on a single CPU. Notice that the views with missing values are generally sparsely observed in many real-world applications, which makes it challenging to learn informative models for such data. Following Qin et al. [2019], we reordered the rows and columns of training data in descending order according to the proportion of observations in them. This makes the first subset the most densely observed block, thus making the resulting submodel informative and facilitating the parallel inferences in the following stages.

1 https://sheffieldml.github.io/GPy/

2 IC50 (units in µM) is the concentration of drug at which 50% of the target is inhibited. The lower the IC50 of the drug, the less likely the drug will be to have some off-target effect (e.g. potential toxicity) that is not desired. pIC50 = -log10 (IC50).

We evaluated performance by predictive accuracy and the quality of prediction for downstream ranking tasks. Root mean squared error (RMSE) is a common performance measure for matrix completion. In real-world applications, such as item recommendation or drug discovery, we are more interested in the performance of the ranking task, for instance how many of the recommended items the user actually clicks or buys, how many drugs recommended by models actually have the desired effect for the disease. For this purpose, we regard matrix completion as a classification task (of whether a prediction is relevant or not at a given threshold), use F1and AUC-ROC score as performance metrics for ExCAPEDB. Furthermore, following Qin et al. [2019], we use the wall-clock time 3 to measure the speed-up achieved by parallelization. For our method, the reported wall-clock time is calculated by summing the maximum wall-clock times of submodels for each inference stage plus the wall-clock time of making prediction.

For compound activity prediction tasks, we use a pIC50 cutoff (a.k.a. affinity level) at 5 and 6, corresponding to concentrations of 10µM and 1µM, respectively. The test set was further filtered by only keeping targets having at least 100 compounds, at least 25 active compounds, and 25 inactive compounds, to ensure a minimum number of positive and negative data points.

Macau 4 was run on compute nodes with 20 CPUs; all the other methods were run on a single CPU. Our implementation is based on the GPy1 package.

4.2 Results

The results for simulated and ExCAPE-DB data are given in Table 1 and 2, respectively. In Table 1, column 'Full posterior' refers to the performance of MRD learned from the full data; column 'Intermediate aggregation' refers to our method which works by first aggregating multiple submodels into a model with reasonable size (as long as the compute node can still accommodate the model to make predictions) and then perform predictions by aggregating predictions from multiple experts with PoE.

It is clear from Table 1 that the model with full posterior performs better than other methods in terms of predictive performance; our intermediate aggregation method achieves competitive results while being much faster than the full posterior. The intermediate aggregation method also performs better than the SoD approach and the variant of our method without the intermediate aggregation step. With a linear kernel, all MRD models perform better than Macau.

3 Wall-clock time measures the real time between the start and the end of a program. For parallel processes, we use the wall-clock time of the slowest process.

4 We ran the Macau version available in SMURFF software: https://github.com/ExaScience/smurff.

For ExCAPE-DB data, our intermediate aggregation method (by aggregating 10 or 20 submodels to obtain larger models for prediction) performs much better than all the other methods in all performance metrics for different affinity levels. At both affinity levels, all versions of our proposed method perform better than the SoD method in terms of RMSE, F1-score and AUC-ROC score. Again, we observed that all MRD methods perform better than Macau in all performance metrics. In both tables, the wall-clock times of our methods are larger than that of the SoD approach. This is due to the two-stage parallel inferences of the proposed scheme.

To summarise, the proposed method with an intermediate aggregation step achieves a better trade-off between predictive accuracy and computation time. The proposed method performs better than the embarrassingly parallel approaches for scalable Gaussian process models and a state-of-the-art highly optimized implementation of linear Bayesian matrix factorization with side information.

5 Conclusion

In this paper, we have introduced a scalable approach for Bayesian non-linear matrix completion. We have argued that a key factor in constructing distributed solutions for massivescale data is to limit the communication required between computational units. To this end, we have introduced a computational scheme which leverages embarrassingly parallel techniques developed for Gaussian process regression by suitably adapting them for Bayesian Gaussian process latent variable models. The resulting framework is almost as communication-efficient as embarrassingly parallel computation, adding only one additional stage of communication, while achieving accuracy close to the non-distributed full data solution.

Acknowledgements

The authors gratefully acknowledge the computational resources provided by the Aalto Science-IT project and support by the Academy of Finland (Finnish Center for Artificial Intelligence, FCAI, and projects 319264, 292334).