The Role of Contextual Information in Best Arm Identification

Masahiro Kato and Kaito Ariu. 2026. Journal of Machine Learning Research 27(51), 1–61. JMLR page, arXiv:2106.14077. 日本語版

Best-arm identification looks for the option with the largest mean reward using as few trials as it can. This study treats the case where a covariate is observed at each round, and examines the problem of identifying the arm whose mean reward is largest once the covariate has been marginalized out. It gives a lower bound on the number of samples required and an algorithm that reaches that bound asymptotically.

The problem

The quantity of interest is the mean reward marginalized over the covariate distribution, not the conditional mean reward at each covariate value. In the language of causal inference, this is the position that looks at the average treatment effect rather than at effects for individuals. For best-arm identification without covariates, an instance-specific lower bound on the sample complexity was already known, along with an algorithm matching it asymptotically. How that bound changes once covariates can be used was not.

The proposed method

We derive an instance-specific lower bound on the sample complexity, both for continuous covariates and for covariates that take finitely many values. The bound fixes the proportions in which the arms have to be pulled to attain it, so the algorithm is a version of Track-and-Stop with covariates that selects arms so as to track that optimal allocation. A stopping rule and a decision rule are designed alongside it, and they are shown to meet the specified error probability.

Main results and conditions

The lower bound derived here is smaller than the covariate-free bound of Garivier and Kaufmann (2016). Covariates therefore leave room to identify the best arm faster, and the sample complexity upper bound of the proposed algorithm agrees with this lower bound asymptotically. Two settings carry the results. One is a two-armed Gaussian bandit in which the rewards and covariates follow a multivariate normal distribution. The other is the case where the reward distribution belongs to a single-parameter exponential family and the covariate takes finitely many values. Numerical experiments confirm that the algorithm reaches the best arm faster when it uses covariates.

Relation to earlier work

For best-arm identification without covariates, Kaufmann, Cappé, and Garivier gave an instance-specific lower bound on the sample complexity (2016, On the Complexity of Best-Arm Identification in Multi-Armed Bandit Models), and Garivier and Kaufmann proposed Track-and-Stop, which matches that lower bound asymptotically (2016, Optimal Best Arm Identification with Fixed Confidence). Degenne, Koolen, and Ménard show how to solve the problem of finding this optimal allocation as a game (2019, Non-Asymptotic Pure Exploration by Solving Games). Using covariates to estimate a marginalized quantity efficiently is an idea close to the adaptive experimental design of Hahn, Hirano, and Karlan (2011, Adaptive Experimental Design Using the Propensity Score), where assigning treatment according to the covariates is likewise shown to lower the variance of the estimator. Our work takes that effect, discussed there as a gain in estimation efficiency, and puts a number on it in the language of sample complexity for best-arm identification.

Where it applies

The setting we have in mind is an experiment that compares several candidates in order to settle on the best one quickly. When information such as user attributes or time of day is available, feeding it into the assignment cuts the number of trials needed to reach the same level of confidence. The designs covered on the Adaptive Experimental Design page form a pair with this one, separated by whether the goal is to estimate an effect precisely or to choose the best candidate.

Paper and materials

BibTeX

@article{role-of-contextual-information-bai,
  author       = {Masahiro Kato and Kaito Ariu},
  title        = {The Role of Contextual Information in Best Arm Identification},
  journal      = {Journal of Machine Learning Research},
  volume       = {27},
  number       = {51},
  pages        = {1--61},
  year         = {2026},
  url          = {https://jmlr.org/papers/v27/22-0358.html},
}