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

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.

Multi-task learning aims to train a model on many related source tasks at once, so that knowledge learned from source tasks can help with target tasks. For example, if several prediction problems rely on the same hidden features, it should be possible to learn those shared features jointly instead of treating every task separately. Our paper studies this idea in a clean mathematical setting where all tasks share a low-dimensional linear representation. Although this setting looks simple, the standard likelihood-based training objective is non-convex, which makes it difficult to prove that practical gradient-based algorithms can both run efficiently and achieve the best possible statistical accuracy. We propose a two-phase gradient descent algorithm. The first phase helps the model find a good region of the ground truth space, and the second phase refines the solution using a regularization term that stabilizes the joint learning of shared and task-specific parameters. We prove that this method converges efficiently and achieves near-optimal estimation error, improving previous likelihood-based methods by a factor related to the representation dimension.