[Paper Review] Top-$k$ eXtreme Contextual Bandits with Arm Hierarchy

Rajat Sen, Alexander Rakhlin+5 more|arXiv (Cornell University)|Feb 15, 2021

Advanced Bandit Algorithms Research45 references4 citations

TL;DR

This paper proposes a hierarchical top-$k$ eXtreme contextual bandit algorithm that leverages arm hierarchy to achieve efficient exploration and regret minimization in settings with millions of arms. By structuring arms into a tree based on label embeddings and using inverse gap weighting with linear regression, the method reduces computation and regret dependence on total arms, achieving 7.9ms average inference time on a 3M-arm dataset—100x faster than prior methods.

ABSTRACT

Motivated by modern applications, such as online advertisement and recommender systems, we study the top-$k$ extreme 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-extreme 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 (|\mathcal{F}|T)})$, where $A$ is the total number of arms and $\mathcal{F}$ is the class containing the regression function, while only requiring $ ilde{O}(A)$ computation per time step. In the extreme 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 (|\mathcal{F}|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 extreme 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.

Motivation & Objective

• Address the scalability challenge in top-$k$ contextual bandits with extremely large arm counts (e.g., millions of arms).
• Overcome the computational and regret inefficiency of standard algorithms that scale poorly with $A$, the total number of arms.
• Introduce a practical arm hierarchy model to exploit structural correlations in reward functions for improved statistical and computational efficiency.
• Achieve low regret with sublinear dependence on $A$ while maintaining $\tilde{O}(A)$ per-step computation.
• Demonstrate empirical superiority on real-world eXtreme multi-label classification datasets with massive arm counts.

Proposed method

• Propose a hierarchical tree structure over arms using recursive 2-means clustering on label embeddings (via PIFA) to group semantically similar arms.
• Define routing functions at each internal node as one-vs-all linear classifiers trained on a small held-out dataset to guide arm selection.
• Use inverse gap weighting (IGW) for selecting $k$ arms at each time step, adapted to the multi-arm setting with hierarchical exploration.
• Train linear regression models ($\tilde{f}(x,\tilde{a}) = \nu_{\tilde{a}}^T[x;1]$) at each node and leaf for reward estimation using online feedback.
• Implement beam search over the tree to identify effective arms, with random sampling from subtrees when non-leaf nodes are selected.
• Use C++ with Eigen for efficient inference, enabling low-latency deployment on large-scale datasets.

Experimental results

Research questions

• RQ1 Can a hierarchical structure in the arm space reduce the effective number of arms considered per context, improving computational and statistical efficiency?
• RQ2 How can the inverse gap weighting (IGW) strategy be extended to support top-$k$ arm selection in extreme contextual bandit problems?
• RQ3 What is the theoretical regret bound of the proposed algorithm in the eXtreme setting, and how does it scale with $A$, $k$, and $T$?
• RQ4 Can the hierarchical model achieve practical speedups (e.g., inference time) without sacrificing regret performance on large-scale datasets?
• RQ5 How does the algorithm compare to standard baselines (e.g., $\epsilon$-greedy, Boltzmann, IGW) in terms of regret and inference efficiency on real-world multi-label datasets?

Key findings

• The proposed algorithm achieves a regret bound of $O(k\sqrt{(\log A - k + 1)T\log(|\mathcal{F}|T)})$ in the eXtreme setting, with logarithmic dependence on $A$ instead of linear.
• On a 3-million-arm dataset, the inference time was reduced to 7.9 milliseconds on average—representing a 100x speedup over baseline methods.
• The hierarchical structure enabled a significant reduction in the number of relevant arms per context, improving both efficiency and scalability.
• Hyperparameter tuning on the Eurlex-4k dataset yielded optimal parameters: $C=1.0$, $\beta=1.0$, and $\epsilon=0.167$ for the respective baselines.
• The algorithm outperformed standard benchmarks in simulated bandit feedback experiments, demonstrating superior regret and inference performance on large-scale multi-label datasets.
• The method maintains theoretical regret guarantees while achieving practical efficiency, bridging the gap between theory and real-world deployment in extreme contextual bandits.