Sparse multi-task lasso

Introduction
Setup
Methods
- Idealized simulation
Results

Introduction

UTMOST (Hu et al. 2019) fits predictive models for gene expression across multiple conditions jointly, based on optimizing the penalized objective:

\[ \newcommand\mb{\mathbf{B}} \newcommand\me{\mathbf{E}} \newcommand\mr{\mathbf{R}} \newcommand\mx{\mathbf{X}} \newcommand\my{\mathbf{Y}} \ell = \frac{1}{2} \sum_{k=1}^m \Vert \my_{\cdot k} - \mx \mb_{\cdot k} \Vert_F + \lambda_1 \sum_{k=1}^m \Vert \mb_{\cdot k} \Vert + \lambda_2 \sum_{j=1}^p \Vert \mb_{j \cdot} \Vert_2 \]

This approach was previously proposed by Lee et al. 2010. We reimplemented this method in the package gtexpred. Here, we evaluate our implementation on simulations.

Setup

import functools as ft
import hyperopt as hp
import mtlasso
import numpy as np
import pandas as pd
import sklearn.linear_model as sklm
import sklearn.metrics as skm
import sklearn.model_selection as skms
import timeit
import txpred

import rpy2.robjects.packages
import rpy2.robjects.pandas2ri
rpy2.robjects.pandas2ri.activate()

%matplotlib inline
%config InlineBackend.figure_formats = set(['retina'])

import colorcet
import matplotlib
import matplotlib.pyplot as plt
plt.rcParams['figure.facecolor'] = 'w'
plt.rcParams['font.family'] = 'Nimbus Sans'

Methods

Idealized simulation

def _simulate(X, B, pve):
  n = X.shape[0]
  m = B.shape[1]
  G = X.dot(B)
  E = np.random.normal(size=(n, m), scale=np.sqrt((1 / pve - 1) * G.var(axis=0)))
  Y = G + E
  # Center the data
  X -= X.mean(axis=0)
  Y -= Y.mean(axis=0)
  return X, Y, B

def simulate_identical_effects(n, p, m, n_causal, pve, seed):
  np.random.seed(seed)
  X = np.random.normal(size=(n, p))
  b = np.zeros((p, 1))
  b[:n_causal] = np.random.normal(size=(n_causal, 1))
  B = np.outer(b, np.ones(m))
  return _simulate(X, B, pve)

def simulate_shared_effects(n, p, m, n_causal, pve, seed):
  np.random.seed(seed)
  X = np.random.normal(size=(n, p))
  B = np.zeros((p, m))
  B[:n_causal,:] = np.random.normal(size=(n_causal, m))
  return _simulate(X, B, pve)

Results

Running time

Benchmark running time for a single choice of \(\lambda_1, \lambda_2\) from cold initialization \(\mb = \boldsymbol{0}\).

def evaluate_running_time(n, p, m):
  f = gtexpred.lasso.sparse_multi_task_lasso
  X, Y, B = simulate_identical_effects(n=n, p=p, m=m, n_causal=1, pve=0.1, seed=0)
  res = timeit.repeat(lambda: f(X, Y, lambda1=1, lambda2=1),
                      repeat=10,
                      number=1,
                      globals=locals())
  return np.array(res).mean(), np.array(res).std()

result = []
for n in (100, 200, 400, 800):
  for p in (1000, 2000, 4000):
    for m in (2, 10, 50):
      mean, std = evaluate_running_time(n, p, m)
      result.append((n, p, m, mean, std))
result = pd.DataFrame(result)
result.columns = ['n', 'p', 'm', 'mean', 'std']
result.to_csv('/project2/mstephens/aksarkar/projects/gtex-pred/data/sparse-multi-task-lasso-runtime.txt.gz', sep='\t', compression='gzip')

sbatch --partition=broadwl -n1 -c1 --exclusive --time=60:00
#!/bin/bash
source activate fqtl
python <<EOF
<<imports>>
<<simulation>>
<<evaluate-running-time>>
EOF

Read the results.

result = pd.read_csv('/project2/mstephens/aksarkar/projects/gtex-pred/data/sparse-multi-task-lasso-runtime.txt.gz', sep='\t', index_col=0)

Plot run time as a function of \(n, m, p\). Error bars denote 2 standard deviations over 10 trials.

cm = plt.get_cmap('Dark2')
plt.clf()
fig, ax = plt.subplots(1, 3)
fig.set_size_inches(6, 2.5)
for a, (m, g) in zip(ax, result.groupby('m')):
  for i, (p, g2) in enumerate(g.groupby('p')):
    a.errorbar(g2['n'], g2['mean'], yerr=2 * g2['std'], c=cm(i), label=p, lw=1)
    a.set_xticks([100, 400, 800])
    a.set_xlabel('$n$')
    a.set_title(f'$m$={m}')
ax[0].set_ylabel('Running time (sec)')
ax[-1].legend(title='$p$', frameon=False, loc='center left', bbox_to_anchor=(1, .5))
fig.tight_layout()

Lasso path

The coordinate descent update (assuming b_jk ≥ 0) is

\[ b_{jk} := \min\left(\frac{\mr_{\cdot k} \mx_{\cdot j} - \lambda_1}{(\mx' \mx)_{jj} + \lambda_2 / \Vert \mb_{j \cdot} \Vert_2}, 0\right) \]

As expected, we recover the lasso update if \(\lambda_2 = 0\). For the lasso, the update reveals that any \(b_{jk}\) for which \(\mr_{\cdot k} \mx_{\cdot j} < \lambda_1\) will be shrunk to 0. This gives a way to find the maximum \(\lambda_1\) for a given data set:

\[ \lambda_{1,\max} = \max \vert \mx' \my / n \vert \]

def est_sparse_multi_task_lasso_path(X, Y, lambda2, trace=0, **kwargs):
  print(f'estimating path for lambda2 = {lambda2}')
  xty = X.T @ Y
  lam_max = np.abs(xty).max()
  grid = np.geomspace(lam_max, .01, 50)
  Bhat = xty / X.shape[0]
  path = dict()
  for a in grid:
    Bhat, _ = mtlasso.lasso.sparse_multi_task_lasso(X, Y, init=Bhat, lambda1=a, lambda2=lambda2, **kwargs)
    path[a] = Bhat[trace].copy()
  return pd.DataFrame.from_dict(path)

Simulate data for 5 conditions and a single causal variant which has identical effect in all conditions. Look at the paths of the coefficients for the simulated causal variant.

X, Y, B = simulate_identical_effects(n=500, p=1000, m=5, n_causal=1, pve=0.1, seed=1)
paths = {
  lambda2: est_sparse_multi_task_lasso_path(X, Y, lambda2=lambda2, atol=1e-4, max_iter=500)
  for lambda2 in (0, 250, 500)
}

cm = plt.get_cmap('Dark2')
plt.clf()
fig, ax = plt.subplots(1, 3, sharey=True)
fig.set_size_inches(6, 2.5)
for a, (lambda2, path) in zip(ax, paths.items()):
  a.set_xscale('log')
  for k, row in path.iterrows():
    a.plot(1 / path.columns, row, lw=1, c=cm(k))
  a.axhline(y=B[0,0], ls=':', c='k', lw=1)
  a.set_xlabel('$1 / \lambda_1$')
  a.set_title(f'$\lambda_2 = {lambda2}$')
ax[0].set_ylabel('$\hat{b}_{0\cdot}$')
fig.tight_layout()

Simulate multiple causal variants.

X, Y, B = simulate_identical_effects(n=500, p=1000, m=5, n_causal=3, pve=0.1, seed=0)
paths = {
  lambda2: est_sparse_multi_task_lasso_path(X, Y, lambda2=lambda2, atol=1e-4, max_iter=500)
  for lambda2 in (0, 250, 500)
}

cm = plt.get_cmap('Dark2')
plt.clf()
fig, ax = plt.subplots(1, 3, sharey=True)
fig.set_size_inches(6, 2.5)
for a, (lambda2, path) in zip(ax, paths.items()):
  a.set_xscale('log')
  for k, row in path.iterrows():
    a.plot(1 / path.columns, row, lw=1, c=cm(k))
  a.axhline(y=B[0,0], ls=':', c='k', lw=1)
  a.set_xlabel('$1 / \lambda_1$')
  a.set_title(f'$\lambda_2 = {lambda2}$')
ax[0].set_ylabel('$\hat{b}_{0\cdot}$')
fig.tight_layout()

Start from the oracle initialization to understand why the estimated coefficients are not driven to be equal.

X, Y, B = simulate_identical_effects(n=500, p=1000, m=5, n_causal=3, pve=0.1, seed=0)
Bhat, Bhat0 = mtlasso.lasso.sparse_multi_task_lasso(X, Y, init=B.copy(), lambda1=0, lambda2=200, atol=1e-4)
Bhat[:3]

array([[ 1.57120005,  1.42403555,  1.36905244,  1.49823112,  1.02940123],
[-0.1172862 , -0.11772508, -0.14990959, -0.12432973, -0.00778592],
[ 1.51250475,  1.25522477,  0.89798756,  1.04866882,  1.42675512]])

Idealized simulation

Simulate an easy example, and verify the method worked.

X, Y, B = simulate_identical_effects(n=500, p=1000, m=10, n_causal=1, pve=0.1, seed=0)
# Hold out validation data to evaluate the final model
Xt, Xv, Yt, Yv = skms.train_test_split(X, Y, test_size=0.1)
mx = Xt.mean(axis=0)
my = Yt.mean(axis=0)

Fix \(\lambda_1, \lambda_2\).

Bhat = gtexpred.lasso.sparse_multi_task_lasso(Xt - mx, Yt - my, lambda1=150, lambda2=100, atol=1e-8)

Make sure the model worked on the training data.

plt.clf()
fig, ax = plt.subplots(3, 1, sharex=True)
fig.set_size_inches(3, 5)

Yhat = (Xt - mx).dot(Bhat)
ax[0].bar(np.arange(10), skm.r2_score(Yt - my, Yhat, multioutput='raw_values'), color='0.7')
ax[0].axhline(y=0.1, ls=':', lw=1, c='k')
ax[0].set_ylabel('$R^2$')

ax[1].bar(np.arange(10), np.median(abs(Yt - my - Yhat), axis=0), color='0.7')
ax[1].set_ylabel('MAE')

ax[2].bar(np.arange(10), skm.mean_squared_error(Yt - my, Yhat, multioutput='raw_values'), color='0.7')
ax[2].set_ylabel('MSE')

ax[2].set_xlabel('Phenotype')
for a in ax:
  a.set_xticks(np.arange(10))
ax[2].set_xticklabels(np.arange(10))
plt.tight_layout()

Evaluate the model on the validation data.

plt.clf()
fig, ax = plt.subplots(3, 1, sharex=True)
fig.set_size_inches(3, 5)

Yhat = (Xv - mx).dot(Bhat)
ax[0].bar(np.arange(10), skm.r2_score(Yv - my, Yhat, multioutput='raw_values'), color='0.7')
ax[0].axhline(y=0.1, ls=':', lw=1, c='k')
ax[0].set_ylabel('$R^2$')

ax[1].bar(np.arange(10), np.median(abs(Yv - my - Yhat), axis=0), color='0.7')
ax[1].set_ylabel('MAE')

ax[2].bar(np.arange(10), skm.mean_squared_error(Yv - my, Yhat, multioutput='raw_values'), color='0.7')
ax[2].set_ylabel('MSE')

ax[2].set_xlabel('Phenotype')
for a in ax:
  a.set_xticks(np.arange(10))
ax[2].set_xticklabels(np.arange(10))
plt.tight_layout()

Compare the fitted coefficients at the true causal variant to the true coefficients.

plt.clf()
plt.gcf().set_size_inches(4, 3)
plt.scatter(np.arange(10), B[0,:], color='r', label='True', s=4)
plt.scatter(np.arange(10), Bhat[0,:], color='k', label='Estimated', s=4)
plt.legend(frameon=False, loc='center left', bbox_to_anchor=(1, .5), handletextpad=0, markerscale=2)
plt.xticks(np.arange(10), np.arange(10))
plt.xlabel('Phenotype')
plt.ylabel('Effect size')
plt.tight_layout()

Cross-validation

Simulate an easy example.

X, Y, B = simulate_identical_effects(n=500, p=1000, m=5, n_causal=1, pve=0.1, seed=0)
# Hold out validation data to evaluate the final model
Xt, Xv, Yt, Yv = skms.train_test_split(X, Y, test_size=0.1)
mx = Xt.mean(axis=0)
my = Yt.mean(axis=0)

Tune \(\lambda_1, \lambda_2\) by 5-fold cross validation within the training data.

# c.f. https://github.com/scikit-learn/scikit-learn/blob/1495f69242646d239d89a5713982946b8ffcf9d9/sklearn/linear_model/coordinate_descent.py#L112
grid = np.geomspace(.1, 1, 10) * Xt.shape[0]
cv_scores = mtlasso.lasso.sparse_multi_task_lasso_cv(
  Xt,
  Yt,
  cv=skms.KFold(n_splits=5),
  lambda1=grid,
  lambda2=grid,
  max_iter=500,
  verbose=True)
cv_scores = pd.DataFrame(cv_scores)
cv_scores.columns = ['fold', 'lambda1', 'lambda2', 'mse']

cv_scores.to_csv('/scratch/midway2/aksarkar/txpred/full-data-cv-scores.txt.gz', sep='\t')

cv_scores = pd.read_csv('/scratch/midway2/aksarkar/txpred/full-data-cv-scores.txt.gz', sep='\t')

lambda1, lambda2 = cv_scores.groupby(['lambda1', 'lambda2'])['mse'].agg(np.mean).idxmin()
lambda1, lambda2

(45.0, 208.87149751257505)

Look at the best performing penalty weight choices.

cv_scores.groupby(['lambda1', 'lambda2'])['mse'].agg(np.mean).sort_values().head().reset_index()

lambda1     lambda2        mse
0  45.000000  208.871498  18.675936
1  75.064524  161.721615  18.697633
2  58.119735  208.871498  18.705495
3  58.119735  161.721615  18.710209
4  45.000000  269.767913  18.737840

Fit the model on the full training data using the best \(\lambda_1, \lambda_2\).

Bhat, B0hat = mtlasso.lasso.sparse_multi_task_lasso(Xt, Yt, lambda1=lambda1, lambda2=lambda2)

Score the final model on the validation data.

pd.DataFrame({
  'r2': skm.r2_score(Yv - my - B0hat, (Xv - mx).dot(Bhat), multioutput='raw_values'),
  'mse': np.square(Yv - my - B0hat - (Xv - mx).dot(Bhat)).mean(axis=0),
  'mae': np.median(np.abs(Yv - my - B0hat - (Xv - mx).dot(Bhat)), axis=0)
})

r2        mse       mae
0  0.212076  23.835363  3.408341
1  0.074575  22.467354  3.564529
2  0.060348  19.873910  2.961368
3  0.116312  22.517498  3.416149
4  0.089665  19.693309  3.128175

Try fixing \(\lambda_1 = 0\)

cv_scores_lam1_0 = mtlasso.lasso.sparse_multi_task_lasso_cv(
  Xt,
  Yt,
  cv=skms.KFold(n_splits=5),
  lambda1=[0.],
  lambda2=grid,
  max_iter=500,
  verbose=True)
cv_scores_lam1_0 = pd.DataFrame(cv_scores_lam1_0)
cv_scores_lam1_0.columns = ['fold', 'lambda1', 'lambda2', 'mse']

cv_scores_lam1_0.groupby(['lambda1', 'lambda2'])['mse'].agg(np.mean).sort_values().head().reset_index()

lambda1     lambda2        mse
0      0.0  300.150811  18.726872
1      0.0  372.063809  18.738914
2      0.0  242.137255  18.856089
3      0.0  195.336637  19.084469
4      0.0  157.581707  19.512256

lambda1, lambda2 = cv_scores_lam1_0.groupby(['lambda1', 'lambda2'])['mse'].agg(np.mean).idxmin()
Bhat, B0hat = mtlasso.lasso.sparse_multi_task_lasso(Xt, Yt, lambda1=lambda1, lambda2=lambda2)
pd.DataFrame({
  'r2': skm.r2_score(Yv - my - B0hat, (Xv - mx).dot(Bhat), multioutput='raw_values'),
  'mse': np.square(Yv - my - B0hat - (Xv - mx).dot(Bhat)).mean(axis=0),
  'mae': np.median(np.abs(Yv - my - B0hat - (Xv - mx).dot(Bhat)), axis=0)
})

r2        mse       mae
0  0.208358  23.947845  3.363572
1  0.075627  22.441819  3.630488
2  0.063343  19.810545  2.960831
3  0.123200  22.341967  3.392123
4  0.088933  19.709142  3.070840

SMBO vs. CV

The goal of CV is to find the parameters of the training algorithm (here, \(\lambda_1\) and \(\lambda_2\)) that minimizes a loss (here, mean-squared error). Above, we evaluated \(k\)-fold CV, which takes \(k \times n^2\) model fits, where \(n\) is the size of the grid of candidate values for \(\lambda_1\) and \(\lambda_2\) This approach scales poorly for high resolution grids (large \(n\)), and self-evidently spends a lot of time on values of the parameters that are highly unlikely to be optimal.

One could in principle do better using sequential model-based optimization (SMBO; e.g., Bergstra et al. 2011). This approach treats the loss as an unknown function \(f\) (which is expensive to evaluate) to be modeled; it uses each evaluation of \(f\) to update a model of \(f\); and it uses the current model of \(f\) to pick the next point at which to evaluate \(f\). Bergstra et al. 2011 propose the Tree of Parzen Estimators as a model of \(f\). In principle, this approach should be able to more efficiently search the space, by focusing effort on parameter values near the optimum.

Simulate the same example as above.

X, Y, B = simulate_identical_effects(n=500, p=1000, m=5, n_causal=1, pve=0.1, seed=0)
# Hold out validation data to evaluate the final model
Xt, Xv, Yt, Yv = skms.train_test_split(X, Y, test_size=0.1)
mx = Xt.mean(axis=0)
my = Yt.mean(axis=0)

Specify CV loss as a black-box function.

Remark One could in principle not do full \(k\)-fold CV within this function; however, then the SMBO algorithm needs to model noisy function evaluations, which seems to be an open problem.

def loss(theta, X, Y, **kwargs):
  lambda1, lambda2 = theta
  scores = []
  init = np.zeros((X.shape[1], Y.shape[1]))
  for train_idx, test_idx in skms.KFold(n_splits=5).split(X, Y):
    Bhat, B0hat = mtlasso.lasso.sparse_multi_task_lasso(X[train_idx], Y[train_idx], lambda1, lambda2, init=init, **kwargs)
    score = mtlasso.lasso._mse(X[test_idx], Y[test_idx] - B0hat, Bhat)
    init = Bhat
    scores.append(score)
  return {'loss': np.array(scores).mean(), 'status': hp.STATUS_OK}

Run SMBO.

np.random.seed(1)
trials = hp.Trials()
opt = hp.fmin(
  ft.partial(loss, X=Xt, Y=Yt, max_iter=500),
  space=[hp.hp.loguniform('lambda1', np.log(.1 * Xt.shape[0]), np.log(Xt.shape[0])),
         hp.hp.loguniform('lambda2', np.log(.1 * Xt.shape[0]), np.log(Xt.shape[0]))],
  algo=hp.tpe.suggest,
  trials=trials,
  max_evals=20)

Look at the contours of the cross-validation loss as a function of \(\lambda_1, \lambda_2\), based on the grid of evaluations in full CV, and compare against the candidate points evaluated in SMBO.

cm = colorcet.cm['bmy']
norm = matplotlib.colors.Normalize(vmin=18., vmax=22.)
plt.clf()
fig, ax = plt.subplots(1, 2, sharey=True)
fig.set_size_inches(5.5, 3)
for a in ax:
  a.set_aspect('equal', adjustable='datalim')
query = cv_scores.pivot_table(index='lambda2', columns='lambda1', values='mse').values
ax[0].contour(grid, grid, query, cmap=cm, norm=norm, levels=15)
ax[0].scatter(lambda1, lambda2, s=32, marker='x', c='k')
ax[0].set_xlabel('$\lambda_1$')
ax[0].set_ylabel('$\lambda_2$')
ax[0].set_title('CV')

query = pd.DataFrame([{k: v[0] for k, v in x['misc']['vals'].items()} for x in trials.trials])
query['loss'] = [x['result']['loss'] for x in trials.trials]
ax[1].scatter(query['lambda1'], query['lambda2'], c=cm(norm(query['loss'])), s=2)
ax[1].scatter(opt['lambda1'], opt['lambda2'], s=32, marker='x', c='k')
ax[1].set_xlabel('$\lambda_1$')
ax[1].set_title('TPE')

fig.tight_layout()

Report the optimal \(\lambda_1, \lambda_2\) estimated by SMBO.

opt

{'lambda1': 46.93776994095178, 'lambda2': 106.27197065892976}

Fit sparse multi-task lasso with that choice of \(\lambda_1, \lambda_2\), and evaluate the performance on the validation data.

Bhat, B0hat = mtlasso.lasso.sparse_multi_task_lasso(Xt, Yt, lambda1=opt['lambda1'], lambda2=opt['lambda2'])
pd.DataFrame({
  'r2': skm.r2_score(Yv - my - B0hat, (Xv - mx).dot(Bhat), multioutput='raw_values'),
  'mse': np.square(Yv - my - B0hat - (Xv - mx).dot(Bhat)).mean(axis=0),
  'mae': np.median(np.abs(Yv - my - B0hat - (Xv - mx).dot(Bhat)), axis=0)
})

r2        mse       mae
0  0.192476  24.428296  3.626432
1  0.024523  23.682503  3.755742
2  0.076187  19.538899  3.059620
3  0.039539  24.473773  3.719337
4  0.019015  21.221685  2.753583

This methodology does slightly worse than full CV, but runs in a fraction of the wall clock time.

Missing data

Simulate an easy example.

X, Y, B = simulate_identical_effects(n=500, p=1000, m=5, n_causal=1, pve=0.1, seed=0)
# Hold out validation data to evaluate the final model
Xt, Xv, Yt, Yv = skms.train_test_split(X, Y, test_size=0.1)
mx = Xt.mean(axis=0)
my = Yt.mean(axis=0)

Suppose not every phenotype is observed in the training data. As an example, mask 50% of the training phenotype matrix entries.

Z = np.random.uniform(size=Yt.shape) < 0.5
Ytm = np.ma.masked_array(Yt, mask=Z)

Fit the model on the observed training data.

grid = np.geomspace(.1, 1, 10) * Xt.shape[0]
cv_scores = mtlasso.lasso.sparse_multi_task_lasso_cv(
  Xt,
  Ytm,
  cv=skms.KFold(n_splits=5),
  lambda1=grid,
  lambda2=grid,
  max_iter=200,
  verbose=True)
cv_scores = pd.DataFrame(cv_scores)
cv_scores.columns = ['fold', 'lambda1', 'lambda2', 'mse']
lambda1, lambda2 = cv_scores.groupby(['lambda1', 'lambda2'])['mse'].agg(np.mean).idxmin()
Bhat0 = mtlasso.lasso.sparse_multi_task_lasso(Xt, Ytm, lambda1, lambda2)

cv_scores.to_csv('/scratch/midway2/aksarkar/txpred/masked-data-cv-scores.txt.gz', sep='\t')

For comparison, fit the model on the full training data.

cv_scores = mtlasso.lasso.sparse_multi_task_lasso_cv(
  Xt,
  Yt,
  cv=skms.KFold(n_splits=5),
  lambda1=grid,
  lambda2=grid,
  max_iter=200)
cv_scores = pd.DataFrame(cv_scores)
cv_scores.columns = ['fold', 'lambda1', 'lambda2', 'mse']
lambda1, lambda2 = cv_scores.groupby(['lambda1', 'lambda2'])['mse'].agg(np.mean).idxmin()
Bhat1 = mtlasso.lasso.sparse_multi_task_lasso(Xt, Yt, lambda1, lambda2)

Compare the fitted values of the two models.

plt.clf()
fig, ax = plt.subplots(1, 5, sharey=True)
fig.set_size_inches(7.5, 2.5)
for i, a in enumerate(ax):
  a.set_aspect('equal', adjustable='datalim')
  a.scatter(Xv @ Bhat0[0][:,i] + Bhat0[1][:,i], Xv @ Bhat1[0][:,i] + Bhat1[1][:,i], s=1, c='k')
  a.set_title(f'Phenotype {i}')
ax[0].set_ylabel('Full data out-of-sample\npredicted value')
a = fig.add_subplot(111, xticks=[], yticks=[], frameon=False)
a.set_xlabel('Masked data out-of-sample predicted value', labelpad=20)
fig.tight_layout()

Evaluate the fits on the validation data.

plt.clf()
fig, ax = plt.subplots(3, 2)
fig.set_size_inches(4, 5)
for a, b, t in zip(ax.T, (Bhat0, Bhat1), ('Masked data', 'Full data')):
  Yhat = Xv @ b[0] + b[1]
  a[0].bar(np.arange(5), skm.r2_score(Yv - my, Yhat, multioutput='raw_values'), color='0.7')
  a[0].axhline(y=0.1, ls=':', lw=1, c='k')
  a[0].set_ylabel('$R^2$')
  a[1].bar(np.arange(5), np.median(abs(Yv - my - Yhat), axis=0), color='0.7')
  a[1].set_ylabel('MAE')
  a[2].bar(np.arange(5), skm.mean_squared_error(Yv - my, Yhat, multioutput='raw_values'), color='0.7')
  a[2].set_ylabel('MSE')
  a[0].set_title(t)
for a in ax.ravel():
  a.set_xticks(np.arange(5))
fig.tight_layout()

Estimate the difference in the estimated effect sizes.

np.linalg.norm(Bhat0[0] - Bhat1[0])

1.8253611488999593

Comparison to univariate elastic-net

Simulate an easy case.

rng = np.random.default_rng(1)
n = 1000
p = 5000
m = 10
X = rng.normal(size=(n, p))
B = np.zeros((p, m))
B[0,:3] = rng.normal(size=(1, 3))
B[1,3:] = rng.normal(size=(1, 7))
Y = X @ B + rng.normal(size=(n, m))

Split the simulated data into training and validation.

Xt, Xv, Yt, Yv = skms.train_test_split(X, Y, test_size=0.1, random_state=1)

Fit univariate elastic-net models for each simulated response, using 4-fold CV on the penalty weights.

enets = []
for k in range(Y.shape[1]):
  print(f'fitting {k}')
  fit = sklm.ElasticNetCV(cv=4, n_jobs=8, normalize=True, fit_intercept=True).fit(Xt, Yt[:,k])
  enets.append(fit)

Report the training and validation MSE for each response.

pd.DataFrame({
  'train': {k: skm.mean_squared_error(Xt @ enets[k].coef_ + enets[k].intercept_, Yt[:,k]) for k in range(10)},
  'val': {k: skm.mean_squared_error(Xv @ enets[k].coef_ + enets[k].intercept_, Yv[:,k]) for k in range(10)},
})

train       val
0  1.162114  0.736277
1  1.098416  1.092427
2  0.790721  2.129253
3  1.010724  1.366760
4  1.303171  1.805048
5  1.023614  1.151061
6  1.075434  0.959652
7  0.764977  3.814086
8  1.025173  1.657425
9  1.096054  1.210242

Fit sparse multi-task lasso, using 4-fold CV to select the optimal lambda2, and taking the optimal lambda1 on the path for that choice.

grid = np.geomspace(.1, 1, 10) * Xt.shape[0]
cv_scores = mtlasso.lasso.sparse_multi_task_lasso_cv(
  Xt,
  Yt,
  cv=skms.KFold(n_splits=4),
  lambda1=grid,
  lambda2=grid,
  max_iter=500,
  verbose=True)
cv_scores = pd.DataFrame(cv_scores)
cv_scores.columns = ['fold', 'lambda1', 'lambda2', 'mse']
lambda1, lambda2 = cv_scores.groupby(['lambda1', 'lambda2'])['mse'].agg(np.mean).idxmin()
Bhat, B0hat = mtlasso.lasso.sparse_multi_task_lasso(Xt, Yt, lambda1=lambda1, lambda2=lambda2)

Report the training and validation MSE for each response.

pd.DataFrame({
  'train': {k: skm.mean_squared_error((Xt @ Bhat + B0hat)[:,k], Yt[:,k]) for k in range(10)},
  'val': {k: skm.mean_squared_error((Xv @ Bhat + B0hat)[:,k], Yv[:,k]) for k in range(10)},
})

train       val
0  1.089945  0.711197
1  1.007758  0.990568
2  1.032572  0.795654
3  0.998397  1.069531
4  1.009947  1.078685
5  1.013060  1.142606
6  1.054464  0.933449
7  1.057387  0.965998
8  1.021355  1.004041
9  1.011801  1.048657

Plot the simulated effect sizes and MSE for the fitted models.

plt.clf()
fig, ax = plt.subplots(3, 1, sharex=True)
fig.set_size_inches(4, 5.5)
i = ax[0].imshow(B[:4], cmap=colorcet.cm['cwr'], vmin=-3, vmax=3)
ax[0].set_aspect('auto')
fig.colorbar(i, label='Ground truth effect size', ax=ax[0], shrink=0.5, orientation='horizontal')

ax[1].bar(range(10), [skm.mean_squared_error(Xt @ enets[k].coef_ + enets[k].intercept_, Yt[:,k]) for k in range(10)],
          color='k', alpha=0.5, label='Training')
ax[1].bar(range(10), [skm.mean_squared_error(Xv @ enets[k].coef_ + enets[k].intercept_, Yv[:,k]) for k in range(10)],
          color='r', alpha=0.5, label='Validation')

ax[2].bar(range(10), [skm.mean_squared_error((Xt @ Bhat + B0hat)[:,k], Yt[:,k]) for k in range(10)],
          color='k', alpha=0.5, label='Training')
ax[2].bar(range(10), [skm.mean_squared_error((Xv @ Bhat + B0hat)[:,k], Yv[:,k]) for k in range(10)],
          color='r', alpha=0.5, label='Validation')

ax[0].set_ylabel('Predictor')
ax[1].set_ylabel('Elastic-net MSE')
ax[2].set_ylabel('mtlasso MSE')
ax[-1].set_xlabel('Response')
for a in ax:
  a.set_xticks(range(10))
for a in ax[1:]:
  a.set_ylim(0, 4)
  a.legend(frameon=False, loc='center left', bbox_to_anchor=(1, .5))
fig.tight_layout()

Look at Fabio's example.

dat = rpy2.robjects.r['readRDS']('/project2/mstephens/fmorgante/test.rds')

X, Y, B, V, _, _, B0, _, idx, _ = [np.array(t) for t in dat]

n = 900
Xt = X[:n]
Xv = X[n:]
Yt = Y[:n]
Yv = Y[n:]

Fit univariate elastic nets.

enets = []
for k in range(Y.shape[1]):
  print(f'fitting {k}')
  fit = sklm.ElasticNetCV(cv=4, n_jobs=8, normalize=True, fit_intercept=True).fit(Xt, Yt[:,k])
  enets.append(fit)

Report the training and validation MSE for each response.

pd.DataFrame({
  'train': {k: skm.mean_squared_error(Xt @ enets[k].coef_ + enets[k].intercept_, Yt[:,k]) for k in range(10)},
  'val': {k: skm.mean_squared_error(Xv @ enets[k].coef_ + enets[k].intercept_, Yv[:,k]) for k in range(10)},
})

train        val
0  25.362915  31.492174
1  19.354158  21.605593
2  34.449236  35.804181
3   4.740143   7.053288
4   1.282722   1.391711
5   0.662963   0.735404
6   3.736664   4.286861
7   2.128571   2.280445
8   0.939318   1.009634
9   5.289383   6.038868

Fit sparse multi-task lasso, using 4-fold CV to select the optimal lambda2, and taking the optimal lambda1 on the path for that choice.

grid = np.geomspace(.1, 1, 10) * Xt.shape[0]
cv_scores = mtlasso.lasso.sparse_multi_task_lasso_cv(
  Xt,
  Yt,
  cv=skms.KFold(n_splits=4),
  lambda1=grid,
  lambda2=grid,
  max_iter=500,
  verbose=True)
lambda1, lambda2 = cv_scores.groupby(['lambda1', 'lambda2'])['mse'].agg(np.mean).idxmin()

Bhat, B0hat = mtlasso.lasso.sparse_multi_task_lasso(Xt, Yt, lambda1=lambda1, lambda2=lambda2)

np.savez('/scratch/midway2/aksarkar/txpred/fabio-ex.npz', Bhat=Bhat, B0hat=B0hat)

Report the training and validation MSE for each response.

pd.DataFrame({
  'train': {k: skm.mean_squared_error((Xt @ Bhat + B0hat)[:,k], Yt[:,k]) for k in range(10)},
  'val': {k: skm.mean_squared_error((Xv @ Bhat + B0hat)[:,k], Yv[:,k]) for k in range(10)},
})

train        val
0  21.119216  27.357582
1  15.093317  18.467278
2  22.991694  31.168670
3   5.986364   6.094789
4   1.161488   1.240754
5   0.640629   0.675358
6   4.054401   3.756979
7   2.015473   2.042536
8   0.869960   0.919764
9   5.217864   5.155873

plt.clf()
fig, ax = plt.subplots(3, 1, sharex=True)
fig.set_size_inches(4, 5.5)
i = ax[0].imshow(B[idx - 1], cmap=colorcet.cm['cwr'], vmin=-3, vmax=3)
ax[0].set_aspect('auto')
fig.colorbar(i, label='Ground truth effect size', ax=ax[0], shrink=0.5, orientation='horizontal')

ax[1].bar(range(10), [skm.mean_squared_error(Xt @ enets[k].coef_ + enets[k].intercept_, Yt[:,k]) for k in range(10)],
          color='k', alpha=0.5, label='Training')
ax[1].bar(range(10), [skm.mean_squared_error(Xv @ enets[k].coef_ + enets[k].intercept_, Yv[:,k]) for k in range(10)],
          color='r', alpha=0.5, label='Validation')

ax[2].bar(range(10), [skm.mean_squared_error((Xt @ Bhat + B0hat)[:,k], Yt[:,k]) for k in range(10)],
          color='k', alpha=0.5, label='Training')
ax[2].bar(range(10), [skm.mean_squared_error((Xv @ Bhat + B0hat)[:,k], Yv[:,k]) for k in range(10)],
          color='r', alpha=0.5, label='Validation')

ax[0].set_ylabel('Predictor')
ax[1].set_ylabel('Elastic-net MSE')
ax[2].set_ylabel('mtlasso MSE')
ax[-1].set_xlabel('Response')
for a in ax:
  a.set_xticks(range(10))
for a in ax[1:]:
  a.legend(frameon=False, loc='center left', bbox_to_anchor=(1, .5))
fig.tight_layout()

Sparse multi-task lasso

Table of Contents

Introduction

Setup

Methods

Idealized simulation

Results

Running time

Lasso path

Idealized simulation

Cross-validation

SMBO vs. CV

Missing data

Comparison to univariate elastic-net