Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Optimal Multitask Linear Regression and Contextual Bandits under Sparse Heterogeneity

Xinmeng Huang, Kan Xu, Donghwan Lee, Hamed Hassani, Hamsa Bastani, Edgar Dobriban

TL;DR

This work addresses estimation and decision-making across multiple heterogeneous tasks by assuming a global shared parameter plus sparse task-specific adjustments. It introduces MOLAR, a two-stage estimator that robustly aggregates covariate-wise OLS estimates via a weighted median to recover a global parameter and then performs covariate-wise shrinkage toward that shared parameter for each task. The authors prove minimax-optimal rates for both offline multitask linear regression and online contextual bandits under sparse heterogeneity, and extend the framework to generalized linear models and confidence intervals. Empirical results on synthetic data and the PISA dataset demonstrate improved estimation accuracy and regret bounds compared to single-task and prior multitask methods. Overall, the paper provides tight theory and practical algorithms that leverage sparse task differences to achieve significant gains in high-dimensional multitask learning contexts.

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.

Optimal Multitask Linear Regression and Contextual Bandits under Sparse Heterogeneity

TL;DR

This work addresses estimation and decision-making across multiple heterogeneous tasks by assuming a global shared parameter plus sparse task-specific adjustments. It introduces MOLAR, a two-stage estimator that robustly aggregates covariate-wise OLS estimates via a weighted median to recover a global parameter and then performs covariate-wise shrinkage toward that shared parameter for each task. The authors prove minimax-optimal rates for both offline multitask linear regression and online contextual bandits under sparse heterogeneity, and extend the framework to generalized linear models and confidence intervals. Empirical results on synthetic data and the PISA dataset demonstrate improved estimation accuracy and regret bounds compared to single-task and prior multitask methods. Overall, the paper provides tight theory and practical algorithms that leverage sparse task differences to achieve significant gains in high-dimensional multitask learning contexts.

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.
Paper Structure (54 sections, 38 theorems, 244 equations, 10 figures, 3 tables, 5 algorithms)

This paper contains 54 sections, 38 theorems, 244 equations, 10 figures, 3 tables, 5 algorithms.

Table of Contents

  1. Introduction
  2. Contributions
  3. Related Works
  4. Data Heterogeneity.
  5. Multitask Linear Regression & Contextual Bandits.
  6. Transfer Learning.
  7. Notation
  8. Multitask Linear Regression
  9. Algorithm Overview
  10. Analysis of Estimation Error under Gaussian Noise
  11. Lower Bound
  12. Linear Contextual Bandit
  13. Algorithm Overview
  14. Regret Analysis
  15. Lower Bound
  16. ...and 39 more sections

Key Result

Proposition 1

TakingThroughout the paper, we assume $\{\sigma_m\}_{m=1}^M$ are known as they can be easily estimated using the OLS-based formula $\widehat{\sigma}_m^2 = \|Y^{(m)}-{\mathbf{X}}^{(m)}\widehat{\beta}^{(m)}_{\rm ind}\|_2^2/(n_m-d)$. Experiments with estimated variances can be found in Appendix app:rob where $\sigma_w:= W_{[M]}^{-1}\sum_{m=1}^Mw_m(\sigma_m/\sqrt{n_m})$ is the weighted average of stan

Figures (10)

  • Figure 1: The differences in the least squares estimates of a measure of educational attainment in selected countries for the PISA dataset. See details in Appendix \ref{['app:experi']}.
  • Figure 2: Average $\ell_1$ estimation error for multitask linear regression. (Left): Fixing $s = 20, \,M = 30$ and varying $n$. (Middle): Fixing $s = 20,\,n = 5,000$ and varying $M$. (Right): Fixing $M = 30, \,n = 5,000$ and varying $s$. The standard error bars (barely visible) are obtained from ten independent trials.
  • Figure 3: Average $\ell_1$ estimation error for multitask linear regression with $n=10,000$ and $\rho=s/d$ fixed. (Left): $\rho=0.1$. (Right): $\rho=0.2$. The standard error bars are obtained from ten independent trials.
  • Figure 4: Regret $R_T^{(m)}$ of instances with activation probability $0.778$ (Left), $0.466$ (Middle), $0.318$ (Right), respectively, where shaded regions depict the corresponding 95% normal confidence intervals based on standard errors calculated over twenty independent trials.
  • Figure 5: Regret $R_T^{(m)}$ of Canada, UAE, and Denmark of the PISA dataset. The shaded regions depict the corresponding 95% normal confidence intervals based on standard errors calculated over twenty independent trials.
  • ...and 5 more figures

Theorems & Definitions (70)

  • Proposition 1: Error bound for well-aligned coordinates
  • Remark 1: Comparison with xu2021multitask
  • Corollary 1: Error bound for global parameter
  • Theorem 1: Error bound for task-wise parameters
  • Remark 2: Varying heterogeneity levels
  • Remark 3: Extensions of MOLAR
  • Corollary 2: Transfer learning for a data-poor task
  • Theorem 2: Minimax lower bound for linear regression under sparse heterogeneity
  • Remark 4
  • Lemma 1: Parameter estimation bound for heterogeneous bandits
  • ...and 60 more