Reproduce “Learning with Incomplete Context: Linear Contextual Bandits with Pretrained Imputation” on OpenResearch

About this paper

Hao Yan, Heyan Zhang, Yongyi Guo

An algorithm called PULSE-UCB integrates pretrained imputation models with linear contextual bandits to address partially observed online contexts, achieving regret guarantees that decompose into standard linear bandit terms and an imputation-dependent component. This approach demonstrated near-optimal performance and improved decision-making on synthetic and real-world datasets, even under covariate shifts.

Key insights

• The cumulative regret in contextual bandits with imputed contexts can be theoretically decomposed into a standard linear bandit regret term and an additional term directly dependent on imputation error.
• The quality of context imputation, quantified by the KL divergence ($D_t$) between true and pretrained conditional distributions, directly determines the magnitude of the additional regret.
• Leveraging sufficient auxiliary data for pretraining allows the algorithm to achieve near-optimal regret rates, almost matching the performance of an oracle with full context access, thereby mitigating information loss from missing contexts.

Summary by alphaXiv.