English

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

Machine Learning 2026-05-07 v2 Optimization and Control

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 O~(1)\widetilde{\mathcal{O}}(1) iterations and attains a \emph{near-optimal} estimation error of O~(dk/(TN))\widetilde{\mathcal{O}}(dk/(TN)), \emph{improving} over existing likelihood-based methods by a factor of kk, where dd, kk, TT, NN 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.

Keywords

Cite

@article{arxiv.2605.00473,
  title  = {Near-optimal and Efficient First-Order Algorithm for Multi-Task Learning with Shared Linear Representation},
  author = {Shihong Ding and Fangyu Du and Cong Fang},
  journal= {arXiv preprint arXiv:2605.00473},
  year   = {2026}
}