Near-Optimal and Efficient First-Order Algorithm for Multi-Task Learning with Shared Linear Representation
Shihong Ding, Fangyu Du, Cong Fang
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 iterations and attains a near-optimal estimation error of , improving over existing likelihood-based methods by a factor of , where , , , 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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext b43c7c09-416d-41f2-96fc-50622ae83419Builds on10
- FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPsAlekh Agarwal, Sham M. Kakade, Akshay Krishnamurthy, Wen SunNeurIPS 2020 · 271 citations
- Provable Meta-Learning of Linear RepresentationsNilesh Tripuraneni, Chi Jin, Michael I. JordanICML 2021 · 218 citations
- Few-Shot Learning via Learning the Representation, ProvablySimon Shaolei Du, Wei Hu, Sham M. Kakade, Jason D. Lee et al.ICLR 2021 · 56 citations
- Last Iterate Risk Bounds of SGD with Decaying Stepsize for Overparameterized Linear RegressionJingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu et al.ICML 2022 · 38 citations
- Provable Benefit of Multitask Representation Learning in Reinforcement LearningYuan Cheng, Songtao Feng, Jing Yang, Hong Zhang et al.NeurIPS 2022 · 33 citations
Related papers
- Near-Optimal Representation Learning for Linear Bandits and Linear RLJiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li et al.ICML 2021 · 60 citations
- Bridging Multi-Task Learning and Meta-Learning: Towards Efficient Training and Effective AdaptationHaoxiang Wang, Han Zhao, Bo LiICML 2021 · 108 citations
- Multi-Task Learning with User Preferences: Gradient Descent with Controlled Ascent in Pareto OptimizationDebabrata Mahapatra, Vaibhav RajanICML 2020 · 182 citations
- Fast and Sample Efficient Multi-Task Representation Learning in Stochastic Contextual BanditsJiabin Lin, Shana Moothedath, Namrata VaswaniICML 2024 · 9 citations
- Provably Efficient Multi-Task Meta Bandit Learning via Shared RepresentationsJiabin Lin, Shana MoothedathNeurIPS 2025 · 2 citations
