Optimal Multitask Linear Regression and Contextual Bandits under Sparse Heterogeneity

Abstract

Large and complex datasets are often collected from several, possibly heterogeneous sources. Multitask learning methods improve efficiency by leveraging commonalities across datasets while accounting for possible differences among them. Here, we study multitask linear regression and contextual bandits under sparse heterogeneity, where the source/task-associated parameters are equal to a global parameter plus a sparse task-specific term. We propose a novel two-stage estimator called MOLAR that leverages this structure by first constructing a covariate-wise weighted median of the task-wise linear regression estimates and then shrinking the task-wise estimates towards the weighted median. Compared to task-wise least squares estimates, MOLAR improves the dependence of the estimation error on the data dimension. Extensions of MOLAR to generalized linear models and constructing confidence intervals are discussed in the paper. We then apply MOLAR to develop methods for sparsely heterogeneous multitask contextual bandits, obtaining improved regret guarantees over single-task bandit methods. We further show that our methods are minimax optimal by providing a number of lower bounds. Finally, we support the efficiency of our methods by performing experiments on both synthetic data and the PISA dataset on student educational outcomes from heterogeneous countries.

Journal of the American Statistical Association, 2024

References in corpus (24)

• Communication-Efficient Learning of Deep Networks from Decentralized Data
• Challenges of Big Data Analysis
• Taking Advantage of Sparsity in Multi-Task Learning
• The multi-armed bandit problem with covariates
• Finite-Time Analysis of Kernelised Contextual Bandits
• Exploiting Shared Representations for Personalized Federated Learning
• Few-Shot Learning via Learning the Representation, Provably
• Provable Meta-Learning of Linear Representations
• Sequential Transfer in Multi-armed Bandit with Finite Set of Models
• Maximin effects in inhomogeneous large-scale data
• Statistical Inference for Online Decision-Making: In a Contextual Bandit Setting
• Sequential Batch Learning in Finite-Action Linear Contextual Bandits
• Multi-Task Learning for Contextual Bandits
• Doubly-Robust Lasso Bandit
• Removing Data Heterogeneity Influence Enhances Network Topology Dependence of Decentralized SGD
• Collaborative causal inference with a distributed data-sharing management
• Multitask Learning and Bandits via Robust Statistics
• Revisiting Optimal Convergence Rate for Smooth and Non-convex Stochastic Decentralized Optimization
• Multi-task Representation Learning with Stochastic Linear Bandits
• Online Sparse Reinforcement Learning
• Learning from Similar Linear Representations: Adaptivity, Minimaxity, and Robustness
• Thompson Sampling for High-Dimensional Sparse Linear Contextual Bandits
• Online Statistical Inference for Contextual Bandits via Stochastic Gradient Descent
• Weighted Distributed Estimation under Heterogeneity