# RepUCB: Representation Learning-Based UCB for Heterogeneous Multi-Task Linear Bandits

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 17 references that are in the map.
- First posted on arXiv 1 Apr 2026.

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

- [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)
- [Provable Meta-Learning of Linear Representations](https://mlsanity.com/paper/2002.11684/) (2020)
- [LoRA-GA: Low-Rank Adaptation with Gradient Approximation](https://mlsanity.com/paper/2407.05000/) (2024)
- [Impact of Representation Learning in Linear Bandits](https://mlsanity.com/paper/2010.06531/) (2020)
- [Near-optimal Representation Learning for Linear Bandits and Linear RL](https://mlsanity.com/paper/2102.04132/) (2021)
- [LoRe: Personalizing LLMs via Low-Rank Reward Modeling](https://mlsanity.com/paper/2504.14439/) (2025)
- [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)
- [Active Multi-Task Representation Learning](https://mlsanity.com/paper/2202.00911/) (2022)
- [Non-Stationary Representation Learning in Sequential Linear Bandits](https://mlsanity.com/paper/2201.04805/) (2022)
- [Fast and Sample Efficient Multi-Task Representation Learning in Stochastic Contextual Bandits](https://mlsanity.com/paper/2410.02068/) (2024)
- [Dynamical Low-Rank Compression of Neural Networks with Robustness under Adversarial Attacks](https://mlsanity.com/paper/2505.08022/) (2025)
- [Noisy Low Rank Column-wise Sensing](https://mlsanity.com/paper/2409.08384/) (2024)

## Abstract (arXiv)

Multi-task representation learning exploits the shared structure among related tasks by learning a common latent representation, thereby improving sample efficiency. This paper introduces a novel approach to multi-task representation learning in heterogeneous linear bandits. We consider T concurrent heterogeneous linear bandit tasks, each with feature dimension d, whose reward parameters share a common latent representation of dimension r ll mind, T, capturing the underlying task relatedness. We propose RepUCB, a novel Upper Confidence Bound (UCB) algorithm that leverages shared low-rank representations to enhance learning in a sample-efficient manner. Our algorithm first collects data through an exploration phase, estimates the shared representation, and then performs UCB-based learning on our proposed confidence set. We provide theoretical guarantees for the confidence set and prove that the unknown reward parameters lie within the confidence set with high probability. We derive cumulative regret bound and show that the proposed approach achieves widetildeO(√drNT), a significant improvement over solving the T tasks independently, resulting in a regret of widetildeO(dT√N). We performed numerical simulations to validate the performance of our algorithm for different problem sizes and compared with baseline algorithms.