[Paper Review] Wasserstein regularization for sparse multi-task regression

Hicham Janati, Marco Cuturi+1 more|arXiv (Cornell University)|May 20, 2018

Sparse and Compressive Sensing Techniques27 references14 citations

TL;DR

This paper proposes Multi-task Wasserstein (MTW), a convex regularizer for sparse multi-task regression that leverages geometric structure via unbalanced optimal transport to model relationships between regression coefficients without requiring overlapping supports. By combining Sinkhorn-regularized optimal transport with proximal coordinate descent, MTW outperforms state-of-the-art methods in recovering true supports across synthetic and neuroimaging datasets, especially in scenarios with non-overlapping or partially overlapping features.

ABSTRACT

We focus in this paper on high-dimensional regression problems where each regressor can be associated to a location in a physical space, or more generally a generic geometric space. Such problems often employ sparse priors, which promote models using a small subset of regressors. To increase statistical power, the so-called multi-task techniques were proposed, which consist in the simultaneous estimation of several related models. Combined with sparsity assumptions, it lead to models enforcing the active regressors to be shared across models, thanks to, for instance L1 / Lq norms. We argue in this paper that these techniques fail to leverage the spatial information associated to regressors. Indeed, while sparse priors enforce that only a small subset of variables is used, the assumption that these regressors overlap across all tasks is overly simplistic given the spatial variability observed in real data. In this paper, we propose a convex regularizer for multi-task regression that encodes a more flexible geometry. Our regularizer is based on unbalanced optimal transport (OT) theory, and can take into account a prior geometric knowledge on the regressor variables, without necessarily requiring overlapping supports. We derive an efficient algorithm based on a regularized formulation of OT, which iterates through applications of Sinkhorn's algorithm along with coordinate descent iterations. The performance of our model is demonstrated on regular grids with both synthetic and real datasets as well as complex triangulated geometries of the cortex with an application in neuroimaging.

Motivation & Objective

• To address the limitation of existing multi-task Lasso methods that assume perfect overlap of active features across tasks, which is often unrealistic in spatially-structured data.
• To incorporate geometric information from regressor locations (e.g., brain regions, climate stations) into multi-task regression without enforcing shared supports.
• To develop a convex, scalable optimization framework that enables efficient learning in high-dimensional, sparse, multi-task regression problems.
• To demonstrate that modeling feature proximity via optimal transport leads to improved support recovery and statistical performance compared to standard sparsity-inducing norms.

Proposed method

• Proposes a novel convex regularizer based on unbalanced optimal transport (OT) to measure similarity between coefficient vectors across tasks, encoding geometric relationships between regressors.
• Uses a regularized formulation of OT with entropic smoothing to enable efficient computation via Sinkhorn’s algorithm, allowing GPU acceleration.
• Combines Sinkhorn iterations with proximal coordinate descent for alternating optimization of the multi-task regression problem.
• Employs a ground metric derived from the physical or geometric structure of the regressors (e.g., cortical mesh geodesic distances) to guide the transport cost.
• Applies the method to both regular grids and complex triangulated geometries, such as the human cortex, to model spatially distributed features.
• Derives a latent barycenter from the coefficient vectors to identify the most representative features across the cohort.

Experimental results

Research questions

• RQ1 Can optimal transport be used to model geometric relationships between regression coefficients in multi-task learning without assuming shared active features?
• RQ2 How does the performance of a Wasserstein-based regularizer compare to traditional ℓ1/ℓq norms in recovering sparse, non-overlapping supports?
• RQ3 To what extent does incorporating geometric structure improve statistical power in high-dimensional regression with limited samples?
• RQ4 Can the proposed method scale efficiently to complex, non-Euclidean geometries such as cortical surfaces?

Key findings

• MTW significantly outperforms independent Lasso and multi-task Lasso in AUC score for support recovery, achieving an AUC of 0.92 on synthetic MEG data with 50% overlap and 6 non-zero features per task.
• In scenarios with non-overlapping supports, MTW achieves an AUC of 0.91, while independent Lasso scores only 0.54, demonstrating its ability to leverage spatial proximity.
• MTW maintains consistent performance gains as the number of tasks increases, with AUC rising from 0.78 (2 tasks) to 0.90 (8 tasks), unlike Dirty models whose performance plateaus.
• The latent barycenter learned by MTW effectively identifies the most representative activation patterns across subjects, highlighting common features in neuroimaging applications.
• On real MEG data with SNR = 4, MTW achieves a mean AUC of 0.88 across 30 trials, outperforming MLL and Dirty models across all overlap levels.
• The method is robust to varying degrees of support overlap, with MTW maintaining superior performance even when overlap is as low as 20%.