Google Research

Top-k eXtreme Contextual Bandits with Arm Hierarchy

  • Rajat Sen
  • Alexander Rakhlin
  • Lexing Ying
  • Rahul Kidambi
  • Daniel Hill
  • Dean Foster
  • Inderjit Dhillon
Top-k eXtreme Contextual Bandits with Arm Hierarchy, International Conference on Machine Learning (2021) (to appear)


Motivated by modern applications, such as online advertisement and recommender systems, we study the top-$k$ \xtm contextual bandits problem, where the total number of arms can be enormous, and the learner is allowed to select $k$ arms and observe all or some of the rewards for the chosen arms. We first propose an algorithm for the non-\xtm realizable setting, utilizing the Inverse Gap Weighting strategy for selecting multiple arms. We show that our algorithm has a regret guarantee of $O(k\sqrt{(A-k+1)T \log (|\cF|T)})$, where $A$ is the total number of arms and $\cF$ is the class containing the regression function, while only requiring $\tilde{O}(A)$ computation per time step. In the $\xtm$ setting, where the total number of arms can be in the millions, we propose a practically-motivated arm hierarchy model that induces a certain structure in mean rewards to ensure statistical and computational efficiency. The hierarchical structure allows for an exponential reduction in the number of relevant arms for each context, thus resulting in a regret guarantee of $O(k\sqrt{(\log A-k+1)T \log (|\cF|T)})$. Finally, we implement our algorithm using a hierarchical linear function class and show superior performance with respect to well-known benchmarks on simulated bandit feedback experiments using \xtm multi-label classification datasets. On a dataset with three million arms, our reduction scheme has an average inference time of only 7.9 milliseconds, which is a 100x improvement.

Learn more about how we do research

We maintain a portfolio of research projects, providing individuals and teams the freedom to emphasize specific types of work