[Paper Review] Multi-task Representation Learning with Stochastic Linear Bandits

Leonardo Cella, Karim Lounici+2 more|arXiv (Cornell University)|Feb 21, 2022

Advanced Bandit Algorithms Research4 citations

TL;DR

This paper proposes a trace norm regularized greedy policy for multi-task stochastic linear contextual bandits that implicitly learns a low-rank shared representation across tasks without requiring prior knowledge of the rank or invertible arm covariance. The method achieves a regret bound of $ T ilde{O}( ext{poly}(r) ext{poly}(d) ext{poly}(N)) $, which is minimax optimal up to logarithmic factors and significantly improves upon independent task learning and prior methods like MLingreedy, especially when $ T \geq d $.

ABSTRACT

We study the problem of transfer-learning in the setting of stochastic linear bandit tasks. We consider that a low dimensional linear representation is shared across the tasks, and study the benefit of learning this representation in the multi-task learning setting. Following recent results to design stochastic bandit policies, we propose an efficient greedy policy based on trace norm regularization. It implicitly learns a low dimensional representation by encouraging the matrix formed by the task regression vectors to be of low rank. Unlike previous work in the literature, our policy does not need to know the rank of the underlying matrix. We derive an upper bound on the multi-task regret of our policy, which is, up to logarithmic factors, of order $\sqrt{NdT(T+d)r}$, where $T$ is the number of tasks, $r$ the rank, $d$ the number of variables and $N$ the number of rounds per task. We show the benefit of our strategy compared to the baseline $Td\sqrt{N}$ obtained by solving each task independently. We also provide a lower bound to the multi-task regret. Finally, we corroborate our theoretical findings with preliminary experiments on synthetic data.

Motivation & Objective

• To address the challenge of transfer learning in stochastic linear contextual bandits with shared low-dimensional representations across tasks.
• To develop a policy that does not require prior knowledge of the true rank of the shared representation matrix.
• To design a method robust to non-invertible arm covariance matrices and effective even with small time horizons.
• To achieve regret bounds comparable to an oracle policy that knows the true representation in advance.
• To outperform existing baselines like independent task learning and MLingreedy, particularly in high-dimensional or low-horizon settings.

Proposed method

• The method employs a greedy policy based on trace norm regularization to encourage low-rank structure in the matrix of task-specific regression vectors.
• It uses a novel martingale concentration argument to derive oracle inequalities under the restricted strong convexity condition.
• The policy estimates the shared representation via a regularized empirical risk minimization that promotes low-rank solutions through the nuclear norm.
• The algorithm adaptively learns the underlying representation without requiring the rank to be specified in advance.
• It operates under minimal assumptions: no boundedness on arms and no invertibility of the arm covariance matrix.
• The policy is computationally efficient and does not require solving complex matrix factorization problems at each step.

Experimental results

Research questions

• RQ1 Can a multi-task bandit policy achieve regret comparable to an oracle that knows the true low-dimensional representation, without prior knowledge of the rank?
• RQ2 How does trace norm regularization improve learning efficiency in multi-task linear bandits with shared representations?
• RQ3 What is the impact of the number of tasks $ T $, dimension $ d $, and time horizon $ N $ on the regret of a joint learning policy?
• RQ4 How does the proposed policy compare to existing methods like MLingreedy when the rank is misspecified or the arm covariance is singular?
• RQ5 Can the policy achieve minimax optimal regret without requiring invertible arm covariance or bounded arm features?

Key findings

• The proposed policy achieves a regret bound of $ T\sqrt{rN} + \sqrt{rNTd} $, up to logarithmic factors, which is minimax optimal.
• The regret bound is significantly better than independent task learning, which incurs a regret of order $ T\sqrt{dN} $.
• The policy outperforms MLingreedy across all configurations, especially when $ d $ is large or $ N $ is small.
• Even with incorrect rank estimation, MLingreedy performs worse than the proposed method, which is robust to rank misspecification.
• The trace norm bandit consistently outperforms both ITL and MLingreedy, and even surpasses the oracle policy when $ r \approx d $, due to continued regularization benefits.
• Numerical experiments confirm that the method is robust to dimensionality and task count, with consistent performance gains over baselines across all tested settings.