# Fast and Sample Efficient Multi-Task Representation Learning in Stochastic Contextual Bandits

arXiv: https://arxiv.org/abs/2410.02068 · HTML: https://mlsanity.com/paper/2410.02068/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 50 references that are in the map.
- First posted on arXiv 2 Oct 2024.

## Builds on (most cited first, 24 of 24)

- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [Few-Shot Learning via Learning the Representation, Provably](https://mlsanity.com/paper/2002.09434/) (2020)
- [The Benefit of Multitask Representation Learning](https://mlsanity.com/paper/1505.06279/) (2015)
- [Actor-Mimic: Deep Multitask and Transfer Reinforcement Learning](https://mlsanity.com/paper/1511.06342/) (2015)
- [Provable Meta-Learning of Linear Representations](https://mlsanity.com/paper/2002.11684/) (2020)
- [Spectral Methods for Data Science: A Statistical Perspective](https://mlsanity.com/paper/2012.08496/) (2020)
- [Batched Multi-armed Bandits Problem](https://mlsanity.com/paper/1904.01763/) (2019)
- [Sequential Batch Learning in Finite-Action Linear Contextual Bandits](https://mlsanity.com/paper/2004.06321/) (2020)
- [High-Dimensional Sparse Linear Bandits](https://mlsanity.com/paper/2011.04020/) (2020)
- [Bilinear Bandits with Low-rank Structure](https://mlsanity.com/paper/1901.02470/) (2019)
- [Impact of Representation Learning in Linear Bandits](https://mlsanity.com/paper/2010.06531/) (2020)
- [Low-Rank Generalized Linear Bandit Problems](https://mlsanity.com/paper/2006.02948/) (2020)
- [Multi-Task Learning for Contextual Bandits](https://mlsanity.com/paper/1705.08618/) (2017)
- [Near-optimal Representation Learning for Linear Bandits and Linear RL](https://mlsanity.com/paper/2102.04132/) (2021)
- [Sharing Knowledge in Multi-Task Deep Reinforcement Learning](https://mlsanity.com/paper/2401.09561/) (2024)
- [Provable Representation Learning for Imitation Learning via Bi-level Optimization](https://mlsanity.com/paper/2002.10544/) (2020)
- [Fast Algorithms for Robust PCA via Gradient Descent](https://mlsanity.com/paper/1605.07784/) (2016)
- [Stochastic Low-Rank Bandits](https://mlsanity.com/paper/1712.04644/) (2017)
- [Sample Efficient Linear Meta-Learning by Alternating Minimization](https://mlsanity.com/paper/2105.08306/) (2021)
- [Multi-task Representation Learning with Stochastic Linear Bandits](https://mlsanity.com/paper/2202.10066/) (2022)
- [Information Directed Sampling for Sparse Linear Bandits](https://mlsanity.com/paper/2105.14267/) (2021)
- [Stochastic Linear Bandits with Hidden Low Rank Structure](https://mlsanity.com/paper/1901.09490/) (2019)
- [Distributed Multitask Learning](https://mlsanity.com/paper/1510.00633/) (2015)
- [Distributed Multi-Task Learning for Stochastic Bandits with Context Distribution and Stage-wise Constraints](https://mlsanity.com/paper/2401.11563/) (2024)

## Built on by (most cited first, 6 of 6)

- [Byzantine Resilient Federated Multi-Task Representation Learning](https://mlsanity.com/paper/2503.19209/) (2025)
- [AltGDmin: Alternating GD and Minimization for Partly-Decoupled (Federated) Optimization](https://mlsanity.com/paper/2504.14741/) (2025)
- [Beyond Centralization: Provable Communication Efficient Decentralized Multi-Task Learning](https://mlsanity.com/paper/2512.22675/) (2025)
- [Diffusion-based Decentralized Federated Multi-Task Representation Learning](https://mlsanity.com/paper/2512.23161/) (2025)
- [RepUCB: Representation Learning-Based UCB for Heterogeneous Multi-Task Linear Bandits](https://mlsanity.com/paper/2604.00531/) (2026)
- [Multi-Task Representation Learning for Conservative Linear Bandits](https://mlsanity.com/paper/2605.12176/) (2026)

## Abstract (arXiv)

We study how representation learning can improve the learning efficiency of contextual bandit problems. We study the setting where we play T contextual linear bandits with dimension d simultaneously, and these T bandit tasks collectively share a common linear representation with a dimensionality of r much smaller than d. We present a new algorithm based on alternating projected gradient descent (GD) and minimization estimator to recover a low-rank feature matrix. Using the proposed estimator, we present a multi-task learning algorithm for linear contextual bandits and prove the regret bound of our algorithm. We presented experiments and compared the performance of our algorithm against benchmark algorithms.