Table of Contents
Fetching ...

Learning with Incomplete Context: Linear Contextual Bandits with Pretrained Imputation

Hao Yan, Heyan Zhang, Yongyi Guo

TL;DR

This work proposes PULSE-UCB, an algorithm that leverages pretrained models trained on the auxiliary data to impute missing features during online decision-making and establishes regret guarantees that decompose into a standard bandit term plus an additional component reflecting pretrained model quality.

Abstract

The rise of large-scale pretrained models has made it feasible to generate predictive or synthetic features at low cost, raising the question of how to incorporate such surrogate predictions into downstream decision-making. We study this problem in the setting of online linear contextual bandits, where contexts may be complex, nonstationary, and only partially observed. In addition to bandit data, we assume access to an auxiliary dataset containing fully observed contexts--common in practice since such data are collected without adaptive interventions. We propose PULSE-UCB, an algorithm that leverages pretrained models trained on the auxiliary data to impute missing features during online decision-making. We establish regret guarantees that decompose into a standard bandit term plus an additional component reflecting pretrained model quality. In the i.i.d. context case with Hölder-smooth missing features, PULSE-UCB achieves near-optimal performance, supported by matching lower bounds. Our results quantify how uncertainty in predicted contexts affects decision quality and how much historical data is needed to improve downstream learning.

Learning with Incomplete Context: Linear Contextual Bandits with Pretrained Imputation

TL;DR

This work proposes PULSE-UCB, an algorithm that leverages pretrained models trained on the auxiliary data to impute missing features during online decision-making and establishes regret guarantees that decompose into a standard bandit term plus an additional component reflecting pretrained model quality.

Abstract

The rise of large-scale pretrained models has made it feasible to generate predictive or synthetic features at low cost, raising the question of how to incorporate such surrogate predictions into downstream decision-making. We study this problem in the setting of online linear contextual bandits, where contexts may be complex, nonstationary, and only partially observed. In addition to bandit data, we assume access to an auxiliary dataset containing fully observed contexts--common in practice since such data are collected without adaptive interventions. We propose PULSE-UCB, an algorithm that leverages pretrained models trained on the auxiliary data to impute missing features during online decision-making. We establish regret guarantees that decompose into a standard bandit term plus an additional component reflecting pretrained model quality. In the i.i.d. context case with Hölder-smooth missing features, PULSE-UCB achieves near-optimal performance, supported by matching lower bounds. Our results quantify how uncertainty in predicted contexts affects decision quality and how much historical data is needed to improve downstream learning.

Paper Structure

This paper contains 43 sections, 17 theorems, 210 equations, 3 figures, 1 algorithm.

Key Result

Lemma 4.1

For any time step $t \in [T]$ and any measurable scalar function $g: {{\mathbb{R}}}^{d_Y} \to {{\mathbb{R}}}$ with $\|g\|_\infty \leq 1$, Here $\mathbb{E}[\cdot]$ denotes the expectations with respect to $\mathbb{P}$.

Figures (3)

  • Figure 1: Comparison of algorithms in synthetic experiments. Left: cumulative regret. Right: 100-step moving average reward. Top: linear missing-feature setting (a). Bottom: nonlinear setting (b). Shaded areas denote $\pm$ one standard error over 30 trials.
  • Figure 2: Algorithm comparison on the Taobao dataset. Shaded areas denote $\pm$ one standard error over 5 runs.
  • Figure 3: Comparison of PULSE-UCB agent learning results under different linearity settings

Theorems & Definitions (23)

  • Lemma 4.1
  • Remark 4.1
  • Lemma 4.2
  • Remark 4.2
  • Theorem 4.1
  • Remark 4.3
  • Proposition 4.1
  • Proposition 4.2
  • Theorem 5.1: Informal Lower Bound
  • Remark 5.1
  • ...and 13 more