A Bayesian approach to multi-task learning with network lasso

One important research topic is personalization, which is setting a different prediction model for each sample. For example, in the field of personalized medicine, disease should not be predicted by a single model for all patients but rather predicted using a different model for each patient. This issue is common to many research fields. To achieve this aim, multi-task learning (Evgeniou and Pontil, (2007); Argyriou et al., (2008); Lengerich et al., (2018)) is useful.

Hallac et al., (2015) proposed solving a multi-task learning problem through the regularized maximum likelihood method, which is also referred to as network lasso. A characteristic of network lasso is to estimate models for each sample and group the models simultaneously. Various applications of network lasso have been reported in the literature (e.g., Ambos et al., (2018); Jung, (2020); Tran et al., (2020)). The estimation of network lasso uses relational coefficients, which are quantitative expressions of the relationships among the samples. However, no estimating methods for the relational coefficients in network lasso have yet been reported.

To overcome this problem, we propose a Bayesian framework for network lasso that includes estimating the relational coefficients. We define the likelihood and the prior distribution of the regression coefficients, and consider network lasso from a Bayesian perspective. Under this framework, we assume the Dirichlet prior distribution for the relational coefficients. Models are estimated by the MCMC algorithm. As a result, our method provides more accurate estimation than existing methods, even if the relational coefficients of some of the samples are unknown. In addition, by using a prior distribution that induces variable selection, irrelevant features are excluded from the model. This should allow us to expand the applicable fields of data compared to the existing method.

The remainder of the paper is organized as follows. Section 2 describes regression modeling based on network lasso. In Section 3, we propose a Bayesian approach to multi-task learning with network lasso. Section 4 focuses on related work. Section 5 compares the performance of the proposed method with the existing method by Monte Carlo simulation, and Section 6 compares the performance by real data analysis. Concluding remarks are given in Section 7.

Suppose that we have observed data $\{(y_{i},\mbox{\boldmath$x$}_{i});i=1,2,\cdots,n\}$ for response variable $y$ and $p$-dimensional predictor variables $\mbox{\boldmath$x$}=(x_{1},x_{2},\cdots,x_{p})^{T}$. A linear regression model is assumed as follows:

where $\mbox{\boldmath$w$}_{i}=(w_{i1},w_{i2},\cdots,w_{ip})^{T}\ (i=1,2,\cdots,n)$ is a $p$-dimensional regression coefficient vector and $\epsilon_{i}$ is an error variable distributed as $\mbox{N}(0,\sigma^{2})$, in which $\sigma^{2}\ (>0)$ is a variance parameter. In a general regression problem, all regression coefficients are assumed to be the same (i.e., $\mbox{\boldmath$w$}_{1}=\mbox{\boldmath$w$}_{2}=\cdots=\mbox{\boldmath$w$}_{n}$). In (1), a different regression coefficient $\mbox{\boldmath$w$}_{i}$ is assumed for each sample. This can be regarded as considering a different model for each sample.

For the overall model, we consider solving the regression problem by minimizing as follows:

where $\lambda\ (>0)$ is a regularization parameter, $\mathcal{E}$ is a set whose elements are pairs of sample subscripts, and $r_{i_{1}i_{2}}$ is a relational coefficient. In this minimization problem, the second term allows us to estimate $\hat{\mbox{\boldmath$w$}}_{i_{1}}=\hat{\mbox{\boldmath$w$}}_{i_{2}}$. If $\hat{\mbox{\boldmath$w$}}_{i_{1}}=\hat{\mbox{\boldmath$w$}}_{i_{2}}$, then the $i_{1}$th sample and the $i_{2}$th sample belong to the same model. The relational coefficient $r_{i_{1}i_{2}}$ is the strength of the relationship between samples $(y_{i_{1}},\mbox{\boldmath$x$}_{i_{1}})$ and $(y_{i_{2}},\mbox{\boldmath$x$}_{i_{2}})$. The stronger the relationship is, the greater the value of $r_{i_{1}i_{2}}$ is. Hallac et al., (2015) referred to the minimization problem in (2) as network lasso.

Here we focus on the relational coefficient $r_{i_{1}i_{2}}$. The more correctly the relational coefficients are set, the better the accuracy of estimation of network lasso is. In network lasso, the relational coefficients are determined independently of (2). For example, if the samples have location information, then the relational coefficients are defined by the distances between the samples. On the other hand, if the data do not contain such information, the determination of the relational coefficients becomes difficult. Therefore, we propose an approach based on Bayesian sparse modeling. Adopting this approach allows us to estimate the model and the relational coefficients simultaneously.

In this section, we introduce multi-task learning using network lasso from the viewpoint of Bayesian sparse modeling. Bayesian sparse modeling, as represented by Bayesian lasso (Park and Casella, 2008), has advantages such as quantifying posterior uncertainty and estimating hyperparameters. The proposed method overcomes the problem described in the previous section.

We here summarize the research peripherally related to the proposed method. First, we describe localized lasso (Yamada et al., 2017), which performs grouping of individuals and variable selection simultaneously. Next, we describe another Bayesian approach for network lasso.

We demonstrated our proposed method with artificial data. We simulated data $(y_{i},\mbox{\boldmath$x$}_{i})\ (i=1,2,\cdots,n)$ from the following model:

where $\mbox{\boldmath$w$}^{*}_{k,i}$ is a vector of the $p$-dimensional true regression coefficient and $\epsilon_{i}$ is an error distributed as $\mbox{N}(0,1)$. In addition, $\mbox{\boldmath$x$}_{i}\ (i=1,2,\cdots,n)$ is generated from a uniform distribution on $[-1,1]$. We considered several settings: $n=30,\ 50$ and $p=5,\ 10$. When $p=10$, the features consist of five irrelevant features and five relevant features. The data have three groups, and the sample size for each group is the same. The true regression coefficients were set as follows:

• •

  When $p=5$,

  	$\displaystyle\mbox{\boldmath$w$}_{1,i}^{*}$	$\displaystyle=$	$\displaystyle(5.0,\ 1.0,\ -1.0,\ 0.0,\ 0.0)^{T},$
  	$\displaystyle\mbox{\boldmath$w$}_{2,i}^{*}$	$\displaystyle=$	$\displaystyle(0.0,\ 1.0,\ -5.0,\ 1.0,\ 0.0)^{T},$
  	$\displaystyle\mbox{\boldmath$w$}_{3,i}^{*}$	$\displaystyle=$	$\displaystyle(0.0,\ 0.0,\ 0.0,\ 0.5,\ -0.5)^{T}.$

• •

  When $p=10$,

  	$\displaystyle\mbox{\boldmath$w$}_{1,i}^{*}$	$\displaystyle=$	$\displaystyle(5.0,\ 1.0,\ -1.0,\ 0.0,\ 0.0,\ \mbox{\boldmath$0$}_{5}^{T})^{T},$
  	$\displaystyle\mbox{\boldmath$w$}_{2,i}^{*}$	$\displaystyle=$	$\displaystyle(0.0,\ 1.0,\ -5.0,\ 1.0,\ 0.0,\ \mbox{\boldmath$0$}_{5}^{T})^{T},$
  	$\displaystyle\mbox{\boldmath$w$}_{3,i}^{*}$	$\displaystyle=$	$\displaystyle(0.0,\ 0.0,\ 0.0,\ 0.5,\ -0.5,\ \mbox{\boldmath$0$}_{5}^{T})^{T}.$

The subscript set $\mathcal{E}$ was also generated by random numbers. Let $\mathcal{F}$ be the set of subscript pairs $(i,j)$ such that $\mbox{\boldmath$x$}_{i}$ and $\mbox{\boldmath$x$}_{j}$ belong to the same group and let $\mathcal{G}$ be the set of subscript pair $(i,j)$ such that $\mbox{\boldmath$x$}_{i}$ and $\mbox{\boldmath$x$}_{j}$ belong to the different groups. The subscript set $\mathcal{E}$ was given by $\mathcal{E}=\{f,g\}$, where $f\subset\mathcal{F}$ and $g\subset\mathcal{G}$. We assume that $\#f=\mbox{TR}\times\#\mathcal{F}$ and $\#g=\mbox{FR}\times\#\mathcal{G}$, where TR and FR are constants between zero and one.

As the estimation accuracy, we used the MSE given by

where

$\hat{\mbox{\boldmath$w$}}_{i}$ is an estimate of $\mbox{\boldmath$w$}_{i}$, and $\Sigma_{i}$ is the variance-covariance matrix of $\mbox{\boldmath$x$}_{i}$.

The dataset was generated 50 times. We computed the mean and standard deviation of MSE from the 50 repetitions. For each generated dataset, the estimates were obtained using 50,000 iterations of a Gibbs sampler. Candidates of the hyperparameters were set as $10^{l-i}\ (i=1,2,\cdots m)$. In order to make a fair comparison, the results were compared using the hyperparameter that maximized the MSE.

The existing method to be compared was localized lasso (LL). We also compared a method that assumes fixed values for $r_{i_{1}i_{2}}$ and $\lambda_{1}$ instead of assuming prior distributions. We refer to this method as the Bayesian approach to multi-task learning with network lasso (BMN). The proposed method is referred to as the Bayesian approach to multi-task learning with network lasso using the Dirichlet–Laplace distribution (DLBMN). By comparing BMN and DLBMN, we can consider whether assuming a prior distribution for the hyperparameters increases the accuracy of the estimation.

Tables 1 to 4 show the means and standard deviations of the MSEs. Not surprisingly, all tables show that the larger the TR is and the smaller the FR is, the smaller the MSE is. Also, for all tables, regardless of TR, the increase in MSE with increasing FR is smaller for DLBMN than for LL and BMN. From the results for $n=30$ and $p=5$ in Table 1, when $\mbox{TR}=0.2$, the MSE is large regardless of the TF. The results for $n=30$ and $p=10$ in Table 2 also show that the MSE is large regardless of TF for $\mbox{TR}=0.2$. When FR is greater than $0.1$, LL and BMN have MSEs greater than 20, while the MSE of DLBMN is very small. In the results for $n=120$ in Table 3 and Table 4, the MSE is small when $\mbox{TR}=0.2$ and $\mbox{FR}=0.01$, unlike the results for $n=30$. In the results for $n=120$ and $p=5$, unlike the other results, DLBMN has a small MSE even with $\mbox{TR}=1.0$ and $\mbox{FR}=0.4$. DLBMN has a smaller MSE than LL and BMN in most of the results. In particular, the number of results such that the MSE was less than $1.0$ was 34 for DLBMN, compared to 19 for LL and 13 for BMN.

We applied our proposed method to a real dataset comprising Greater Sacramento area data (http://support.spatialkey.com/spatialkey-sample-csv-data/). This dataset was used by Hallac et al., (2015). This dataset is a listing of real estate transactions for one week of May 2008. It contains 985 samples of sales information, including latitude, longitude, number of bedrooms, number of bathrooms, square feet, and sales price. We considered a model that predicts the selling price from the numbers of bedrooms, bathrooms, and square feet. The data have some missing values, with 17% of the samples missing at least one of the above three features. We used the remaining 814 samples after excluding those with missing values. We determined $\mathcal{E}$ using the latitude and longitude coordinates of each house. The $\mathcal{E}$ is the set of subscript pairs $(i,j)$ corresponding to the five houses closest to the $i$-th house ($i=1,2,\cdots,n$) after removing the test set.

We considered the same three methods as in Section 5: LL, BMN, and DLBMN. The methods were compared using five-fold cross-validation. Accuracies of each method were evaluated using prediction squared error (PSE). The PSE corresponding to the $k$-th iteration of cross-validation is given by

where $\widetilde{\mbox{\boldmath$y$}}^{(k)}$ and $\widetilde{X}^{(k)}$ are the test data in the cross-validation, $n^{(k)}$ is the number of test samples, and $\hat{\mbox{\boldmath$\beta$}}^{(k)}$ is the vector of the regression coefficient estimated by the data other than the test data. We used the value of the hyperparameter that maximized the PSE.

All methods had a mean value and standard deviation of PSE of 0.37 and 0.09, respectively. We defined $\mathcal{E}$ by the distances between houses. It is assumed that the locations of properties and groups of individuals are highly related. Therefore, the accuracy of relational coefficients in estimating LL and BMN is expected to be high. In such a case, LL is known to estimate the model with high accuracy. We confirm by this study that DLBMN enjoys the same accuracy as LL, even in environments where LL can estimate with high accuracy.

We proposed an approach to multi-task learning with network lasso based on Bayesian sparse modeling. By grouping samples based on a Dirichlet–Laplace prior, we modeled the relationships among samples, which have not been considered by existing methods. Compared with the existing methods, the proposed method estimated the parameters with high accuracy, regardless of the estimation accuracy of the relationships among the samples. This result suggests that the fields of applicable data can be expanded compared to the existing methods.

S. K. was supported by JSPS KAKENHI Grant Numbers JP19K11854 and JP20H02227. Super-computing resources were provided by Human Genome Center (The Univ. of Tokyo).