Near-optimal and Efficient First-Order Algorithm for Multi-Task Learning with Shared Linear Representation

2026-05-01

AI Summary

This work addresses the long-standing challenge in multi-task learning of jointly estimating shared linear representations and task-specific parameters, a problem rendered non-convex by matrix factorization and lacking efficient likelihood-based methods with theoretical guarantees. The paper proposes the first likelihood-based first-order optimization algorithm that converges in only Õ(1) iterations and achieves a near-optimal estimation error bound of Õ(dk/(TN)), where d is the feature dimension, k the representation rank, T the number of tasks, and N the sample size per task. By effectively overcoming the non-convexity barrier, the method reduces the estimation error by a factor of k compared to existing likelihood-based approaches, substantially improving both computational efficiency and statistical accuracy.

Abstract

Multi-task learning (MTL) has emerged as a pivotal paradigm in machine learning by leveraging shared structures across multiple related tasks. Despite its empirical success, the development of likelihood-based efficiently solvable algorithms--even for shared linear representations--remains largely underdeveloped, primarily due to the non-convex structure intrinsic to matrix factorization. This paper introduces a first-order algorithm that jointly learns a shared representation and task-specific parameters, with guaranteed efficiency. Notably, it converges in $\widetilde{\mathcal{O}}(1)$ iterations and attains a \emph{near-optimal} estimation error of $\widetilde{\mathcal{O}}(dk/(TN))$, \emph{improving} over existing likelihood-based methods by a factor of $k$, where $d$, $k$, $T$, $N$ denote input dimension, representation dimension, task count, and samples per task, respectively. Our results justify that likelihood-based first-order methods can efficiently solve the MTL problem.

Problem

• multi-task learning
• shared linear representation
• non-convex optimization
• likelihood-based methods
• matrix factorization

Innovation

• multi-task learning
• shared linear representation
• first-order algorithm
• near-optimal estimation
• non-convex optimization