[Paper Review] Impact of Representation Learning in Linear Bandits

Jiaqi Yang, Wei Hu+2 more|arXiv (Cornell University)|Oct 13, 2020

Advanced Bandit Algorithms Research48 references4 citations

TL;DR

This paper proposes a novel algorithm for multi-task linear bandits that leverages shared low-dimensional representations to improve regret performance. By exploiting a common $k \ll d$-dimensional feature extractor across $T$ tasks, the method achieves $\widetilde{O}(T\sqrt{kN} + \sqrt{dkNT})$ regret, significantly outperforming independent learning ($\widetilde{O}(T\sqrt{dN})$) when $T$ is large, with matching minimax lower bounds confirming optimality up to logarithmic factors.

ABSTRACT

We study how representation learning can improve the efficiency of bandit problems. We study the setting where we play $T$ linear bandits with dimension $d$ concurrently, and these $T$ bandit tasks share a common $k (\ll d)$ dimensional linear representation. For the finite-action setting, we present a new algorithm which achieves $\widetilde{O}(T\sqrt{kN} + \sqrt{dkNT})$ regret, where $N$ is the number of rounds we play for each bandit. When $T$ is sufficiently large, our algorithm significantly outperforms the naive algorithm (playing $T$ bandits independently) that achieves $\widetilde{O}(T\sqrt{d N})$ regret. We also provide an $Ω(T\sqrt{kN} + \sqrt{dkNT})$ regret lower bound, showing that our algorithm is minimax-optimal up to poly-logarithmic factors. Furthermore, we extend our algorithm to the infinite-action setting and obtain a corresponding regret bound which demonstrates the benefit of representation learning in certain regimes. We also present experiments on synthetic and real-world data to illustrate our theoretical findings and demonstrate the effectiveness of our proposed algorithms.

Motivation & Objective

• To theoretically characterize the benefit of representation learning in sequential decision-making, specifically in multi-task linear bandits.
• To design an algorithm that exploits shared low-dimensional representations across related bandit tasks to reduce regret.
• To establish tight regret bounds—both upper and lower—that demonstrate the efficiency gain from representation learning.
• To extend the framework to the infinite-action setting and validate the approach empirically.

Proposed method

• The algorithm jointly learns a shared $k$-dimensional feature extractor $\mathbf{B} \in \mathbb{R}^{d \times k}$ across $T$ tasks, where each task's parameter $\bm{\theta}_t = \mathbf{B} \mathbf{w}_t$.
• It uses a contextual bandit framework where each task's reward is linear in the embedded context $\mathbf{B}^\top \mathbf{x}_{n,t,a}$.
• The method employs a novel exploration strategy that balances learning the shared representation and optimizing individual task policies.
• For the infinite-action setting, it extends the algorithm using a kernelized or continuous-action variant with a regret bound of $\widetilde{O}(d^{1.5}k\sqrt{TN} + kT\sqrt{N})$.
• Theoretical analysis combines concentration inequalities and information-theoretic lower bounds to establish minimax optimality.
• Empirical validation is conducted on synthetic data and the MNIST dataset to demonstrate performance gains under representation learning.

Experimental results

Research questions

• RQ1 Can representation learning provably reduce regret in multi-task linear bandits?
• RQ2 What is the fundamental regret limit when tasks share a common low-dimensional representation?
• RQ3 How does the number of tasks $T$ and the representation dimension $k$ affect the performance gain from representation learning?
• RQ4 Does representation learning remain beneficial in the infinite-action bandit setting?
• RQ5 Under what conditions can negative transfer occur due to insufficient task diversity?

Key findings

• The proposed algorithm achieves $\widetilde{O}(T\sqrt{kN} + \sqrt{dkNT})$ regret in the finite-action setting, which is strictly better than the $\widetilde{O}(T\sqrt{dN})$ regret of independent learning when $T$ is large.
• An $\Omega(T\sqrt{kN} + \sqrt{dkNT})$ regret lower bound is established, proving that the algorithm's regret is minimax-optimal up to poly-logarithmic factors.
• In the infinite-action setting, the algorithm achieves $\widetilde{O}(d^{1.5}k\sqrt{TN} + kT\sqrt{N})$ regret, outperforming the naive $O(Td\sqrt{N})$ baseline when $T = \widetilde{\Omega}(dk^2)$.
• Experiments on synthetic and MNIST data confirm that regret per task decreases with increasing $T$, validating the benefit of representation learning in high-task regimes.
• The algorithm exhibits negative transfer when $T$ is small, demonstrating that insufficient task diversity can hurt performance.
• The hyper-parameter scaling $N_1 = d^{1.5}k\sqrt{N/T}$ is critical—only with $c=1.5$ does the algorithm outperform baselines like PEGE in practice.