In this work, we develop formulations that simultaneously learn: (a) the parameters of multi-response/task regression models, and (b) the grouping structure in the parameter matrix (row or column or both) that reflects the group relationship between inputs and outputs. We present optimization approaches to efficiently solve the resulting convex problems, and show their numerical convergence. We describe and justify several hyperparameter choices we make during this optimization. Our methods are validated empirically on synthetic data and on real-world datasets concerning phenotypes and genotypes of plant varieties. From the synthetic data experiments, we find that our methods provide a much better and more stable (i.e., lesser standard error) recovery of the underlying group structure. In real-world data experiments, our approaches reveal natural groupings of phenotypes and checkerboard patterns of phenotype-SNP groups that inform us of the joint relationship between them.

We emphasize that the parameters as well as the grouping structures are fully unknown a-priori, and inferring them simultaneously is our major contribution. This is in contrast to the naive way of estimating the parameters first and then clustering. This naive approach has the danger of propagating the estimation error into clustering, particularly in high dimensions, where the estimator is usually inaccurate due to lack of sufficient samples. Moreover, the clustering step of the naive approach does not use the full information of the data. The joint estimation-clustering procedure we propose naturally promotes sharing of information within groups. Our formulations adopt the convex bi-clustering cost function (Chi et al., 2014) as the regularizer to encourage groupings between columns (tasks) and rows (features) in the parameter matrix. Note that, Chi et al. (2014) assume that the data matrix to be used for bi-clustering is known a-priori, which is obviously not the case for our setting. As a result, Chi et al. (2014) can deal only with fixed data and cannot estimate unknown model parameters, while our approaches can simultaneously estimate parameters and discover the clustering structure in them.

To the best of our knowledge, this is the first method that can simultaneously cluster and estimate the parameters efficiently in a unified optimization. We emphasize that our main goal is to discover the underlying parameter bi-cluster structure without compromising estimation accuracy. Experiments show that our clusterings are better than other methods while the estimation accuracy is no worse or sometimes even better.

The premise in multi-task learning is that appropriate sharing of information can benefit all the tasks (Caruana, 1998; Obozinski et al., 2006; Yu et al., 2018). Assuming all tasks to be closely related can be excessive as it ignores the underlying specificity of the mappings. There have been several extensions to multi-task learning that address this problem. The authors in Jalali et al. (2010) propose a dirty model for feature sharing among tasks, wherein a linear superposition of two sets of parameters—one that is common to all tasks, and one that is task-specific—are used. Kim and Xing (2010) leverages a predefined tree structure among the output tasks (e.g., using hierarchical agglomerative clustering) and imposes group regularizations on the task parameters based on this tree. The approach proposed in Kumar and Daume (2012) learns to share by defining a set of basis task parameters and posing the task-specific parameters as a sparse linear combination of these. Jacob et al. (2009) and Kang et al. (2011) assume that the tasks are clustered into groups and proceed to learn the group structure along with the task parameters using a convex and an integer quadratic program, respectively. However, these approaches do not consider joint clustering of the features. In addition, the mixed integer program of Kang et al. (2011) is computationally intensive and greatly limits the maximum number of tasks that can be considered. Another pertinent approach is the Network Lasso formulation presented in Hallac et al. (2015). This formulation, however, is limited to settings where only clustering among the tasks is needed.

As mentioned before, convex bi-clustering method (Chi et al., 2014) aims at grouping observations and features in a data matrix; while our approaches aim at discovering groupings in the parameter matrix of multi-response regression models while jointly estimating such a matrix, and the discovered groupings reflect groupings in features and responses.

In section 2, we will discuss the proposed joint estimation-clustering formulations; in section 3, we will present the optimization approaches. The choice of hyperparameters used and their significance is discussed in section 4. We illustrate the solution path for one of the formulations in section 4.3. We will provide results for estimation with synthetic data, and two case studies using multi-response GWAS with real data in sections 5, 6, respectively. We conclude in section 7. Additional details and convergence proofs for the optimization approaches are provided in the Supplementary Material.

We begin with the simplest formulation, which, as we shall see, is a special case of the latter one.

Here L(Z; Θ) is the loss function, R(Θ) is a regularizer, and ΩW(Θ)=∑i<jwij‖Θ·i-Θ·j‖2 and Θ_·i is the ith column of Θ. Ω_W(Θ) is inspired by the convex bi-clustering objective (Chi et al., 2014, Equation 1) and it encourages sparsity in differences between columns of Θ. Similarly, ΩW(ΘT) encourages sparsity in the differences between the rows of Θ. When the overall objective is optimized, we can expect to see a checkerboard pattern in the model parameter matrix. Note that W and W~ are non-negative weights that reflect our prior belief on the closeness of the rows and columns of Θ.

The degree of sharing of parameters and hence that of bi-clustering, is controlled using the tuning parameter λ₂. When λ₂ is small, each element of Θ will be its own bi-cluster. As λ₂ increases, more elements of Θ fuse together, the number of rectangles in the checkerboard pattern will reduce. See Figure 2 for the change of the checkerboard structure as λ₂ increases. Further, by varying λ₂ we get a solution path instead of just a point estimate of Θ (see section 4.3). In the rest of the paper, we will use the same design matrix X across all tasks for simplicity, without loss of generality.

For sparse multi-task linear regression, we have L(Z; Θ) = L(X, Y; Θ) and formulation 1 can be instantiated as,

Here the rows of Θ correspond to the features, i.e., the columns of X, and the columns of Θ correspond to the tasks, i.e., the columns of Y. Therefore, the checkerboard pattern in Θ provides us insights on the groups of features that go together with the groups of tasks.

Formulation 1 is natural and simple, but it forces the parameters belonging to the same row or column cluster to be equal, and this may be limiting. To relax this requirement, we introduce a surrogate parameter matrix Γ that will be used for bi-clustering. This will be mandated to be close to Θ. For multitask regression this yields the objective

Remark 1. To interpret this more carefully, let us assume that Θ·i=Θ·i¯+Γ·i in Equation (3). In other words, Θ_·i has a global component Θ·i¯, and the component Γ_·i that participates in the clustering. As λ₂ → ∞, Θ·i¯→0, and hence Θ_·i → Γ_·i. Now, formulation 2 reduces to formulation 1. Further, if λ₁ and λ₂ are held constant while only λ₃ increases, Θ·i→Θ·i¯+Γ, since Γ_·i → Γ for all i. The key difference between formulation 2 and 1 is the presence of a task-specific global component Θ·i¯, which lends additional flexibility in modeling the individual tasks even when λ₃ → 0. Whereas, in Equation (2), when λ₂ → ∞, all the components of Θ_·i take the same value for all i, and the tasks are forced to share the same coefficients without any flexibility.

Remark 2. In certain applications, it might make sense to cluster together features/tasks whose effects have the same amplitude but different signs. This can be accommodated by considering ΩW(Θ)=∑i<jwij‖Θ·i-cijΘ·j‖2 where c_i,j ∈ {−1, 1} are predefined constants reflecting whether the features or tasks are expected to be negatively or positively correlated.

For our formulation 1 we use the proximal decomposition method introduced in Combettes and Pesquet (2008). This is an efficient algorithm for minimizing the sum of several convex functions. Our general objective function (1) involves three such functions: f₁ being L(X, Y; Θ), f₂ being R(Θ), and f₃ being the term that multiplies λ₂. At a high level, the algorithm iteratively applies proximal updates with respect to these functions until convergence.

We stack the regression matrix Θ into a column vector (Θ1;…;Θk)∈ℝpk. The proximal operator is given by: proxfb=argmina(f(a)+12‖b-a‖22), where a and b are pk-dimensional vectors. The proximal operator of the regularized loss can be computed according to the specific L and R functions. The overall optimization procedure is given in Algorithm 1 and with the following update rules.

For our formulation 2 we use an alternating minimization method on Θ and Γ; i.e., we alternatively minimize over Θ and Γ with the other fixed. The first alternating step is to estimate Θ while fixing Γ. This minimization problem is separable for each column and each sub-problem can be easily written as a standard Lasso problem:

by defining y~i=]yi,λ2Γ·i] and X~=[X,λ2Ip] and hence can be solved efficiently and in parallel for each column. In the second step, we fix Θ and optimize for Γ. The optimization is

which is a standard bi-clustering problem on Θ and can be solved efficiently using the COnvex BiclusteRing Algorithm (COBRA) introduced in Chi et al. (2014), and described in Algorithm 3 (Supplementary Material) for completeness. The overall procedure is given in Algorithm 2.

We establish the following convergence result for our algorithms, when the loss function L(X, Y; Θ) is convex in Θ. The proofs are given in the Supplementary Material.

Proposition 1. The algorithm described in section 3.1 converges to the global minimizer.

Proposition 2. The algorithm described in section 3.2 converges to the global minimizer.

For both formulations 1 and 2, computational complexity is dominated by the use of the COBRA algorithm. COBRA solves a sequence of convex clustering problems. The subroutine used to solve each convex clustering subproblem scales in storage and computational operations as O(kpq), where k is the number of tasks, p is the number of features and q is the number of non-zero weights. In our case q is much smaller than p² and k². Indeed as we shall see in section 4.1, our weights are based on κ = 5 nearest neighbors.

The choice of the column and row similarity weights W and W~ can affect the quality of the clustering results and we follow the suggestion in Chi et al. (2014) to set these. However, in their case the data matrix to be clustered is fixed and known, but in our case, the coefficient matrix Θ we want to cluster is not known. We will get a rough estimate Θ^ by just minimizing the regularized loss function L(Z; Θ) + λ₁R(Θ). For example, with multi-task regression in Equations (2) and (3), we can solve

where λ₁ is tuned by cross-validation (CV) and re-used in the rest of the algorithm. From our multi-task regression experiment, we find that the clustering results are robust to the choice of λ₁.

With the estimated Θ^ we are ready to compute W and W~. The weights for the columns i and j are computed as wij=1ijk·exp(-ϕ‖Θ^·i-Θ^·j‖22) where 1ijk is 1 if j is among i's κ-nearest-neighbors or vice versa and 0 otherwise. Here ϕ is non-negative and ϕ = 0 corresponds to uniform weights. In our synthetic and real data experiments we fix ϕ = 20. W~ is computed analogously. It is important to keep the two penalty terms Ω_W(Θ) and ΩW~(ΘT) on the same scale, else the row or column objective will dominate the solution. We normalize so that the column weights sum to 1/n and the row weights sum to 1/p.

We set the penalty multipliers (λ₁, λ₂, and λ₃) for both the formulations using a CV approach. We randomly split our samples into a training set and a hold-out validation set, fitting the models on the training set and then evaluating the root-mean-squared error (RMSE) on the validation set to choose the best values. In order to reduce the computational complexity, we estimate the multipliers greedily, one or two at a time. From our simulations, we determined that this is a reasonable choice. We recognize that these can be tuned further on a case-by-case basis.

λ₁ is set to the reasonable value as determined in section 4.1 for both formulations, since the clustering results are robust to this choice. For formulation 1, we estimate the best λ₂ by CV using Equation (1). For formulation 2, the tuning process is similar, but we pick a sequence of λ₂ and λ₃. We estimate both Θ^λ2,λ3 and Γ^λ2,λ3, but calculate RMSE with Γ^λ2,λ3, since it is used in the clustering objective. When the path of bi-clusterings is computed, we fix λ₂ to the CV estimate and vary only λ₃.

One can obtain the entire solution paths for the estimated coefficients Θ by varying the penalty multipliers. Here we provide an example using a synthetic dataset generated as follows. We consider the multi-task regression model: Y = XΘ* + E with eij~N(0,σ2). All the entries of design matrix X are generated as iid from N(0, 1). The true regression parameter Θ^* has a bi-cluster (checkerboard) structure. To simulate sparsity, we set the coefficients within many of the blocks in the checkerboard to 0. For the non-zero blocks, we follow the generative model recommended in Chi et al. (2014): the coefficients within each cluster are generated as θ_ij = μ_rc + ϵ_ij with ϵij~N(0,σϵ2) to make them close but not identical, where μ_rc is the mean of the cluster defined by the rth row partition and cth column partition. We set n = 50, p = 20, and k = 15. For the non-zero blocks, we set μ_rc ~ Uniform{−2, −1, 1, 2} and set σ_ϵ = 0.25. We set σ = 1.5. We use relatively small values for p and k since there will be a total of pk solution paths to visualize.

We begin with the solution paths for formulation 1. We first fix a reasonable λ₁ and vary λ₂ to get solution paths for all the coefficients. In our experiment, we chose λ₁ based on cross-validation as described in section 4.1. These paths are shown in Figure 3. We can see that as λ₂ increases, the coefficients begin to merge and eventually for large enough λ₂ they are all equal. The solution paths are smooth in λ₂. Similarly, we fix λ₂ based on the cross-validation scheme described in section 4.2 and vary λ₁ to get solution paths for all the coefficients. This is shown in Figure 4. It is well-known that the solution paths for LASSO are piecewise linear (Rosset and Zhu, 2007), when L is least squares loss. Here, we see that the solution paths are not piecewise linear, but rather a smoothed version of it. This smoothness is imparted by the convex clustering regularization, the third term in Equation (2).

We can obtain the solution path for formulation 2 as functions of two variables. We first fix a reasonable λ₁ and vary λ₂, λ₃ to get solution paths for all the coefficients. These paths are shown in Figure 5. The solution paths are smooth in λ₂ and λ₃. Similarly, we fix a reasonable λ₂ and vary λ₁, λ₃ to get solution paths for all the coefficients. These paths are shown in Figure 6. The solution paths are smooth in λ₁ and λ₃. The reasonable values are obtained using cross-validation.

It is well-known that LASSO tends to select too many variables (Meinshausen and Yu, 2009). Hence ||Θ_·i − Θ_·j||₂ may not be exactly zero in most cases, and we may end up identifying too many clusters as well. In Chi et al. (2014) the authors defined the measure v_ij = ||Θ_·i − Θ_·j||₂ and placed the ith and jth columns in the same group if v_ij ≤ τ for some threshold τ, inspired by Meinshausen and Yu (2009). In our formulation 1, after selecting the best tuning parameters and estimating Θ, we place the ith and jth rows in the same group if ||Θ_i· − Θ_j·||₂ ≤ τ_r. Similarly, if ||Θ_·i − Θ_·j||₂ ≤ τ_c we place the ith and jth columns in the same group. For formulation 2, we repeat the same approach using Γ instead of Θ.

To compute the thresholds τ_r and τ_c, we first calculate [v_col]_ij = ||Θ_·i − Θ_·j||₂ and stack this matrix to vector v_col; similarly we calculate [v_row]_ij = ||Θ_i· − Θ_j·||₂ and stack to vector v_row. In the case of sparse linear regression, τ should be on the order of the noise (Meinshausen and Yu, 2009): τ∝σlog(p)/n, where σ is typically estimated using the standard deviation of residuals. In general we could set τ proportional to the standard deviation of v_row or v_col.

However in our case, we have an additional regression loss term for estimating the parameters and hence there are two sources of randomness, the regression residual and the error in v. Taking these into account, we set τc=12[σlog(p)/n+std(vcol)] and τr=12[σlog(p)/n+std(vrow)]. We set the multiplier to 12, following the usual conservative heuristics.

Although formulations 1 and 2 have been developed for row-column bi-clustering, they can be easily specialized to clustering columns or rows alone, by respectively using only Ω_W(Θ) or ΩW~(ΘT) in Equation (2), or using only Ω_W(Γ) or ΩW~(ΓT) in Equation (3).

Assessing the clustering quality can be hard. In this paper, we use the following three measures to evaluate the quality of clustering: the adjusted Rand index (Hubert and Arabie, 1985) (ARI), the F-1 score (F-1), and the Jaccard index (JI).

Assume B is the true clustering, define TP to be the number of pairs of elements in S that are in the same subset in A and in the same subset in B. This is the true positive and similarly we can define TN, FN, FP as true negative, false negative, and false positive, respectively. Define precision=TPTP+FP and recall=TPTP+FN, the F-1 score is defined as:

Using the same notation as F-1 score, the Jaccard Index is defined as:

For all these three measures, a value of 1 implies the best possible performance, and a value of 0 means that we are doing poorly. In order to compute ARI, F-1, and JI, we choose the value of the multiplier λ₂ in formulation 1, and {λ₂, λ₃} in formulation 2 using the approach described in section 4.2, and obtain the estimated clusterings.

The estimation accuracy is measured by calculating the RMSE on an independent test set, and also the parameter recovery accuracy, ‖Θ^est-Θ*‖/‖Θ*‖ where Θ^est and Θ^* are the estimated and true coefficient matrices.

We focus on multi-task regression: Y = XΘ* + E with eij~N(0,σ2). All the entries of design matrix X are generated as iid from N(0, 1). The true regression parameter Θ^* has a bi-cluster (checkerboard) structure. To simulate sparsity, we set the coefficients within many of the blocks in the checkerboard to 0. For the non-zero blocks, we follow the generative model recommended in Chi et al. (2014): the coefficients within each cluster are generated as θ_ij = μ_rc + ϵ_ij with ϵij~N(0,σϵ2) to make them close but not identical, where μ_rc is the mean of the cluster defined by the r^th row partition and c^th column partition. We set n = 200, p = 500, and k = 250 in our experiment. For the non-zero blocks, we set μ_rc ~ Uniform{−2, −1, 1, 2} and set σ_ϵ = 0.25. We try the low-noise setting (σ = 1.5), where it is relatively easy to estimate the clusters, and the high-noise setting (σ = 3), where it is harder to obtain them.

We compare our formulations 1 and 2 with a 2-step estimate-then-cluster approach: (a) Estimate Θ^ first using LASSO, and (b) perform convex bi-clustering on Θ^. Θ^ is estimated by solving (6) while selecting the best λ₁ as discussed in section 4.1, and the convex bi-clustering step is implemented using COBRA algorithm in Chi et al. (2014). Our baseline clustering performance is the best of: (a) letting each coefficient be its own group, and (b) imposing a single group for all coefficients.

The average clustering quality results on 50 replicates are shown in Table 1 for low and high noise settings. Most performance measures are reported in the format mean±std.dev. In both tables, the first, second, and third blocks correspond to performances of row, column and row-column bi-clusterings, respectively. We optimize only for bi-clusterings, but the row and the column clusterings are obtained as by-products.

From Table 1 we see that both our formulation 1 and 2 give better results on row clustering, column clusterings, and row-column bi-clustering compared to the 2-step procedure. Moreover, the clustering results given by our formulations are more stable, with lesser spread in performance.

It is also instructive to note that the performance obtained for columns is substantially higher compared to those obtained with rows. This could be because of two reasons: (a) the columns of Θ have a one-to-one correspondence to the columns of the task responses Y, and hence any relationship between the tasks is easily inherited, (b) the rows of Θ can be noisier than the columns, since each row contributes to all the tasks.

The performance boost obtained with high noise is much higher compared to that with low noise. This makes sense because when noise level is low, the estimation step in the 2-step approach is more accurate and the error propagated into the clustering step is relatively small. However, at high noise levels, the estimation can be inaccurate. This estimation error propagates into the clustering step and makes the clustering result of 2-step approach unreliable. Since our formulations jointly perform estimation and clustering, they obtain more reliable and stable results.

The RMSEs evaluated on the test set and the parameter recovery accuracy are provided in Table 2. The oracle RMSE (with Θ known) is 1.5 for the low noise setting and 3.0 for the high noise setting in Table 2, and we can see that the proposed methods provide improvements over the others. We also observe improvements in the parameter recovery accuracy. Although the improvement is marginal, it demonstrates that we are not losing estimation accuracy because of the biclustering structure we considered.

The experimental data was obtained from 18 Sorghum varieties planted in 6 replicate plot locations, and we considered the trait of plant height. The 18 variety names are given in the Supplementary Material.

From the RGB and hyperspectral images of each plot, we extract features of length 206. Hence n = 6, p = 206, and the number of tasks k = 18, for each trait considered. The presence of multiple varieties with replicates much smaller in number than predictors poses a major challenge: building separate models for each variety is unrealistic, while a single model does not fit all. This is where our proposed simultaneous estimation and clustering approach provides the flexibility to share information among tasks that leads to learning at the requisite level of robustness. Note that here we use the column-only clustering variant of formulation 1.

The dendrogram for task clusters obtained by sweeping the penalty multiplier λ₂ is given in Figure 7. This provides some interesting insights from a plant science perspective. As highlighted in Figure 7, the predictive models (columns of Θ) for thicker medium dark plants are grouped together. Similar grouping is seen for thinner tall dark plants, and thick tall plants with many light leaves.

To compute RMSE, we perform 6-folds CV where each fold consists of at least one example from each variety. As we only have n = 6 samples per variety (i.e., per task), it is unrealistic to learn separate models for each variety. For each CV split, we first learn a grouping using one of the compared methods, treat all the samples within a group as i.i.d, and estimate their regression coefficients using Lasso. The methods compared with our approach include: (a) single model, which learns a single predictive model using Lasso, treating all the varieties as i.i.d., (b) No group multitask learning, which learns a traditional multitask model using Group Lasso where each variety forms a separate group, and (c) Kang et al. (2011), which uses a mixed integer program to learn shared feature representations among tasks, while simultaneously determining “with whom” each task should share. Results reported in Table 3, indicate the superior quality of our groupings in terms of improved predictive accuracy.

We apply our approach in a multi-response Genome-Wide Association Study (GWAS). While traditional GWAS focuses on associations to single phenotypes, we would like to automatically learn the grouping structure between the phenotypes as well as the features (columns and rows of Θ) using our proposed method. We use the proposed formulations 1 and 2 (bi-clustering variant) in this experiment.

The design matrix X consisted of SNPs of Sorghum varieties. We consider n = 911 varieties and over 80,000 SNPs. We remove duplicate SNPs and also SNPs that do not have significantly high correlation to at least one response variable. Finally, we end up considering p = 2,937 SNPs. The output data Y contains the following 6 response variables (columns) for all the n varieties collected by hand measurements:

For each variety, each trait can be an average of measurements from up to four plants.

The coefficient matrix given by our formulations are visualized in Figure 8. To make the figure easier to interpret, we exclude the rows with all zero coefficients and take the average over the coefficients within each bi-cluster. The light yellow regions are coefficients close to zero; red and blue areas are positive and negative coefficients, respectively. The rows and columns are reordered to best show the checkerboard patterns. We wish to emphasize again that these checkerboard patterns in the coefficient matrices are automatically discovered using our proposed procedures, and are not readily evident, or trivially discoverable from the data.

The two formulations reveal similar bi-clustering patterns up to reordering. For column clusters, the plant height tasks (h1 and h2), the stem diameter tasks (d1 and d2), and the leaf tasks (l1 and l2) group together. Also, the stem diameter and leaf tasks are more related to each other compared to the height tasks. The bi-clustering patterns reveal the groups of SNPs that influence similar phenotypic traits. Coefficients for height features in the GWAS (Figure 9) study show SNPs with strong effects coinciding with locations of Dwarf 3 (Multani et al., 2003) and especially Dwarf 1 (Hilley et al., 2016) genes known to control plant height that are segregating and significant in the population. The lack of any effect at the Dwarf 2 (Hilley et al., 2017) locus supports previous work indicating that this gene is not a strong contributing factor in this population. This demonstrates that we are able to discover existing factors. We also identify potentially new SNPs for further investigation and biological validation, since many coefficients align with loci outside of the previously identified known height genes.

To evaluate predictive accuracy, we split our data set into three parts: 70% training, 15% validation, and 15% test. We estimate the coefficient matrices by optimizing our formulations on the training set, select the tuning parameters based on the validation set (sections 4.2, 4.4), and then calculate the RMSE on the test set. Table 4 shows the RMSE on test set.

We also estimate the RMSE of the proposed formulations and compare it with the RMSE provided by a simple Lasso model and 2-step procedure. This is shown in Table 4. We see that the RMSE of our formulations are slightly less than that of the Lasso and 2-step procedure. Hence, for similar estimation performance, we are able to discover additional interesting structure in the input-output relationship using our proposed methods.

KN and AL were employed by IBM. The remaining authors declare that the research was conducted in the absence of any commercial or financial relationships that could be construed as a potential conflict of interest.