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

2024-10-02

International Conference on Machine Learning

AI Summary

In multi-task linear contextual bandits, high-dimensional contexts incur substantial sample and computational overhead. Method: This paper investigates efficient learning under a shared low-dimensional linear representation. We propose a unified framework combining alternating projected gradient descent with a minimization estimator, enabling the first theoretically guaranteed recovery of low-rank feature matrices under stochastic context assumptions. Contribution/Results: Our analysis establishes rigorous regret convergence guarantees for multi-task learning, with an upper bound that strictly improves upon single-task baselines. Experiments demonstrate 3–5× higher sample efficiency and significantly accelerated convergence across multiple tasks. The core innovation lies in tightly coupling low-rank structural priors with online decision dynamics—achieving both statistical efficiency and computational tractability without compromising theoretical soundness.

Technology Category

• Machine Learning: Online Learning & Bandits
• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

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.

Problem

• Multi-Armed Bandits
• Contextual Bandits
• Efficient Learning

Innovation

• Multi-task Learning
• Contextual Bandits
• Efficient Feature Learning