Multi-task Representation Learning for Pure Exploration in Linear Bandits
Yihan Du, Longbo Huang, Wen Sun
Abstract
Despite the recent success of representation learning in sequential decision making, the study of the pure exploration scenario (i.e., identify the best option and minimize the sample complexity) is still limited. In this paper, we study multitask representation learning for best arm identification in linear bandits (RepBAI-LB) and best policy identification in contextual linear bandits (RepBPI-CLB), two popular pure exploration settings with wide applications, e.g., clinical trials and web content optimization. In these two problems, all tasks share a common low-dimensional linear representation, and our goal is to leverage this feature to accelerate the best arm (policy) identification process for all tasks. For these problems, we design computationally and sample efficient algorithms DouExpDes and C-DouExpDes, which perform double experimental designs to plan optimal sample allocations for learning the global representation. We show that by learning the common representation among tasks, our sample complexity is significantly better than that of the native approach which solves tasks independently. To the best of our knowledge, this is the first work to demonstrate the benefits of representation learning for multi-task pure exploration. Recently, an emerging number of works (Yang et al., 2021; 2022; Hu et al., 2021; Cella et al., 2022b) investigate representation learning for sequential decision making, and show that if all tasks share a joint low-rank representation, then by leveraging such a joint representation, it is possible to learn faster than treating each task independently. Despite the accomplishments of these works, they mainly focus on the regret minimization setting, where the performance is measured by the cumulative reward gap between the optimal option and the actually chosen options.
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.
Cited by top-tier papers5
- Multi-task Representation Learning for Pure Exploration in Bilinear BanditsSubhojyoti Mukherjee, Qiaomin Xie, Josiah Hanna, Robert D. NowakNeurIPS 2023 · 10 citations
- Multi-Agent Learning with Heterogeneous Linear Contextual BanditsAnh Do, Thanh Nguyen-Tang, Raman AroraNeurIPS 2023 · 7 citations
- Guarantees for Nonlinear Representation Learning: Non-identical Covariates, Dependent Data, Fewer SamplesThomas T. C. K. Zhang, Bruce D. Lee, Ingvar M. Ziemann, George J. Pappas et al.ICML 2024 · 2 citations
- Provably Efficient Multi-Task Meta Bandit Learning via Shared RepresentationsJiabin Lin, Shana MoothedathNeurIPS 2025 · 2 citations
- Multi-task Linear Regression without Eigenvalue Lower Bounds: Adaptivity, Robustness, and SafetySeok-Jin KimICML 2026
Builds on10
- Optimal Best-arm Identification in Linear BanditsYassir Jedra, Alexandre ProutièreNeurIPS 2020 · 99 citations
- Gamification of Pure Exploration for Linear BanditsRémy Degenne, Pierre Ménard, Xuedong Shang, Michal ValkoICML 2020 · 86 citations
- An Empirical Process Approach to the Union Bound: Practical Algorithms for Combinatorial and Linear BanditsJulian Katz-Samuels, Lalit Jain, Zohar S. Karnin, Kevin JamiesonNeurIPS 2020 · 72 citations
- Near-Optimal Representation Learning for Linear Bandits and Linear RLJiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li et al.ICML 2021 · 60 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
Related papers
- Fast and Sample Efficient Multi-Task Representation Learning in Stochastic Contextual BanditsJiabin Lin, Shana Moothedath, Namrata VaswaniICML 2024 · 9 citations
- Impact of Representation Learning in Linear BanditsJiaqi Yang, Wei Hu, Jason D. Lee, Simon Shaolei DuICLR 2021 · 58 citations
- On the Sample Complexity of Representation Learning in Multi-Task Bandits with Global and Local StructureAlessio Russo, Alexandre ProutièreAAAI 2023 · 5 citations
- Provable General Function Class Representation Learning in Multitask Bandits and MDPRui Lu, Andrew Zhao, Simon S. Du, Gao HuangNeurIPS 2022 · 11 citations
- Beyond task diversity: provable representation transfer for sequential multitask linear banditsThang Duong, Zhi Wang, Chicheng ZhangNeurIPS 2024 · 3 citations
