Multitask regression with task-specific designs

minimize  ∑ₜ cₜ ‖yₜ - β₀ₜ - Xₜbₜ‖² / (2nₜ) + λ ∑ⱼ ωⱼ ‖(s₁ⱼb₁ⱼ, …, s_qⱼb_qⱼ)‖₂

for q regressions yₜ = β₀ₜ + Xₜbₜ + eₜ on the same p predictors, each with its own observations and sample size nₜ, with multitasklasso, MultitaskLassoPath and CVMultitaskLassoPath. The penalty on the rows of B = [b₁ ⋯ bq] selects a predictor for all the tasks or for none, sₜⱼ is the standard deviation of column j within task t, and the task weights are cₜ = 1 by default. A predictor that matters in every task is found from the observations of all of them, while each task keeps coefficients of its own.

using HighDimensionalModels
using Printf, Random, Statistics

rng = Xoshiro(12)
Random.Xoshiro(0x1ff81e1cb8aef442, 0xcf9dcaabe83d1f77, 0x67e73055571c9066, 0x50e64a4746114b2e, 0xcca177be0f6e32e7)

q = 3 tasks with nₜ = 60, 100 and 150 observations of p = 500 predictors. The first s = 10 predictors matter in every task, with coefficients of their own in each, between 0.5 and 1.5 in absolute value, and the intercepts differ.

q, p, s = 3, 500, 10
ns = [60, 100, 150]
B = zeros(p, q)
B[1:s, :] = (0.5 .+ rand(rng, s, q)) .* rand(rng, (-1, 1), s, q)
β₀ = [1.0, -0.5, 2.0]
Xs = [randn(rng, n, p) for n in ns]
ys = [β₀[t] .+ Xs[t] * B[:, t] + randn(rng, ns[t]) for t in 1:q]
3-element Vector{Vector{Float64}}:
 [-1.0279065933315545, -4.484365839774658, 3.3158458854930717, 1.180662658609221, 4.919713166964762, 4.086082679925056, -1.1506949640006037, -0.0494984531051228, 2.6099678871809457, -0.5585853523151831  …  -3.5535412424825203, -1.4379916484411854, -0.4117071673263353, 3.82909190917478, 0.7619164320449556, -3.0237262044690008, 3.2890873685032513, 1.7026780119545784, 3.504183186726524, -1.1891614011894414]
 [-5.027013441327283, -0.38871839505409655, 2.465615852375333, 0.7643331226282316, -4.041658257512434, -6.084399718855762, -2.355468175858698, -3.6789240416472113, -1.7273990197372744, -1.2311607191368688  …  1.719223827684405, 0.4616326442091216, -0.07821277416278538, -1.106877435192331, 2.3653845364887993, -3.544095640647969, -1.9143577171206867, 5.545081392321425, 0.8015551101833336, -2.7402684313889285]
 [1.5957185693648575, 3.0105101263658414, 4.487359059738563, 1.270606977002148, -2.3087174388769744, 6.058784723344742, -0.8290789875450443, 2.426455701238654, 3.997843725800448, 2.6561221565748303  …  4.727233881551341, 0.29838641833699897, -2.096397600067647, 0.39848419201031327, 1.22079688471774, -3.4643408503331417, -3.043276294208396, 6.5501359193681825, -3.9527736415237165, 5.5411376761528]

1000 new observations of each task; the noise variance, 1, is the smallest prediction error

Xnew = [randn(rng, 1000, p) for _ in 1:q]
ynew = [β₀[t] .+ Xnew[t] * B[:, t] + randn(rng, 1000) for t in 1:q]

truesel(rows) = count(<=(s), rows)
falsesel(rows) = count(>(s), rows)
falsesel (generic function with 1 method)

Cross-validation within each task

Five folds drawn within each task, so that every fold holds out part of every task; each training fold centers, standardizes and fits on its own observations. The default task weights cₜ = 1 count the mean squared error of each task once, whatever its size; task_weights = :observations counts each observation once instead.

cv = CVMultitaskLassoPath(Xs, ys; nfolds = 5, rng = Xoshiro(1))
j = findfirst(==(cv.λ1se), cv.λpath)
B̂ = coef(cv)[2:end, :]                  # the intercepts are in the first row
rows = findall(k -> any(!iszero, view(B̂, k, :)), 1:p)
@printf("λ1se = %.4f, level %d of %d: %d rows selected\n", cv.λ1se, j, length(cv.λpath), length(rows))
println("selected rows ", rows, ", true support 1:$s")
println("held-out error of each task at λ1se (cvm_task) ", round.(cv.cvm_task[:, j]; digits = 2))
λ1se = 0.3291, level 44 of 100: 34 rows selected
selected rows [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 29, 34, 38, 75, 92, 98, 113, 148, 149, 174, 203, 206, 208, 222, 236, 344, 353, 362, 418, 429, 441, 448, 463, 485], true support 1:10
held-out error of each task at λ1se (cvm_task) [4.31, 1.8, 1.63]

Each task on its own

CVLassoPath on each task alone, with five folds, an intercept and standardized columns, against the multitask fit: the selected predictors, the cross-validated error at λ1se, and the error on the new observations, which predict(cv, Xnew) gives as a vector with one prediction vector per task. The multitask fit selects the same rows in every task, and the smaller tasks gain the most from the observations of the others.

ŷ = predict(cv, Xnew)
println("             multitask lasso                  lasso on each task")
println("task   nₜ   true  false  CV error  test error   true  false  CV error  test error")
for t in 1:q
  cvt = CVLassoPath(Xs[t], ys[t]; nfolds = 5, intercept = true, rng = Xoshiro(1))
  jt = findfirst(==(cvt.λ1se), cvt.λpath)
  sel = findall(!iszero, coef(cvt)[2:end])
  @printf("%4d  %4d   %4d  %5d  %8.2f  %10.2f   %4d  %5d  %8.2f  %10.2f\n", t, ns[t],
          truesel(rows), falsesel(rows), cv.cvm_task[t, j], mean(abs2, ynew[t] - ŷ[t]),
          truesel(sel), falsesel(sel), cvt.cvm[jt], mean(abs2, ynew[t] - predict(cvt, Xnew[t])))
end
             multitask lasso                  lasso on each task
task   nₜ   true  false  CV error  test error   true  false  CV error  test error
   1    60     10     24      4.31        3.94      4      6      9.30       11.97
   2   100     10     24      1.80        2.01      8     17      3.05        2.76
   3   150     10     24      1.63        1.69     10      6      2.02        2.20