Contextual Bandits for Unbounded Context Distributions

Multi-armed bandit (Robbins, 1952; Lai & Robbins, 1985) is an important sequential decision problem that has been extensively studied (Agrawal, 1995; Auer et al., 2002; Garivier & Cappé, 2011). In many practical applications such as recommender systems and information retrieval in healthcare and finance (Bouneffouf et al., 2020), decision problems are usually modeled as contextual bandits (Woodroofe, 1979), in which the reward depends on some side information, called contexts. At the $t$-th iteration, the decision maker observes the context $\mathbf{X}_{t}$, and then pulls an arm $A_{t}\in\mathcal{A}$ based on $\mathbf{X}_{t}$ and the previous trajectory $(\mathbf{X}_{i},A_{i}),i=1,\ldots,t-1$. Many research assume linear rewards (Abbasi-Yadkori et al., 2011; Bastani & Bayati, 2020; Bastani et al., 2021; Qian et al., 2023; Langford & Zhang, 2007; Dudik et al., 2011; Chu et al., 2011; Li et al., 2010), which is restrictive and may not fit well into practical scenarios. Consequently, in recent years, nonparametric contextual bandits have received significant attention, which does not make any parametric assumption about the reward functions (Perchet & Rigollet, 2013; Guan & Jiang, 2018; Gur et al., 2022; Blanchard et al., 2023; Suk & Kpotufe, 2023; Suk, 2024; Cai et al., 2024).

Despite significant progress on nonparametric contextual bandits, existing studies focus only on the case with bounded contexts, and the probability density functions (pdf) of the contexts are required to be bounded away from zero. However, many practical applications often involve unbounded contexts, such as healthcare (Durand et al., 2018), dynamic pricing (Misra et al., 2019) and recommender systems (Zhou et al., 2017). In particular, the contexts may follow a heavy-tailed distribution (Zangerle & Bauer, 2022), which is significantly different from bounded contexts. Therefore, to bridge the gap between theoretical studies and practical applications of contextual bandits, an in-depth theoretical study of unbounded contexts is crucially needed. Compared with bounded contexts, heavy-tailed context distribution requires the learning method to be adaptive to the pdf of contexts, in order to balance the bias and variance of the estimation of reward functions. On the other hand, compared with existing works on nonparametric classification and regression with identically and independently distributed (i.i.d) data, bandit problems require us to achieve a good balance between exploration and exploitation, thus the learning method needs to be adaptive to the suboptimality gap of reward functions. Therefore, the main challenge of solving nonparametric contextual bandit problems with unbounded contexts is to achieve both bias-variance tradeoff and exploration-exploitation tradeoff using a single algorithm.

In this paper, we solve the nonparametric contextual bandit problem with heavy-tailed contexts. To begin with, we derive a minimax lower bound that characterizes the theoretical limits of contextual bandit learning. We then propose a relatively simple method that uses fixed $k$ combined with upper confidence bound (UCB) exploration and derive the bound of expected regret. Even for bounded contexts, our method improves over an existing nearest neighbor method (Guan & Jiang, 2018) for large margin parameter $\alpha$, since our method uses an improved UCB calculation which is more adaptive to the suboptimality gap of reward functions. Despite such progress, there is still some gap between the regret bound and the minimax lower bound, indicating room for further improvement. To close such a gap, we further propose a new adaptive nearest neighbor approach, which selects $k$ adaptively based on the density of samples and the suboptimality gap of reward functions. Our analysis shows that the regret bound of this new method nearly matches the minimax lower bound up to logarithmic factors, indicating that this method is approximately minimax optimal.

The general guidelines of our adaptive kNN method are summarized as follows. Firstly, with higher context pdf, we use larger $k$, and vice versa. Secondly, given a specific context, if the value of an action is far away from optimal (i.e. large suboptimality gap), then we use smaller $k$, and vice versa. Such a choice of $k$ achieves a good tradeoff between estimation bias and variance. With a lower pdf or larger suboptimality gap, the samples are relatively more sparse. As a result, a large bias may happen due to large kNN distances. Therefore, we use smaller $k$ to control the bias. On the contrary, with a higher pdf or smaller suboptimality gap, the samples are dense and thus we can use larger $k$ to reduce the variance. Note that the pdf and the suboptimality gap are unknown to the learner. Therefore, we design a method, such that the value of $k$ is selected by a data-driven manner, based on the density of existing samples.

In this section, we briefly review the related works about contextual bandits and nearest neighbor methods.

Nonparametric contextual bandits with bounded contexts. (Yang & Zhu, 2002) first introduced the nonparametric contextual bandit problem, proposed an $\epsilon$-greedy approach and proved the consistency. (Rigollet & Zeevi, 2010) derived a minimax lower bound on the regret, and showed that this bound is achievable by a UCB method. (Perchet & Rigollet, 2013) proposed Adaptively Binned Successive Elimination (ABSE), which adapts to the unknown margin parameter. (Qian & Yang, 2016) proposed a kernel estimation method. (Hu et al., 2020) analyzed nonparametric bandit problem under general Hölder smoothness assumption. (Gur et al., 2022) proposed Smoothness-Adaptive Contextual Bandits (SACB), which is adaptive to the smoothness parameter. (Slivkins, 2014; Suk & Kpotufe, 2023; Akhavan et al., 2024; Ghosh et al., 2024; Komiyama et al., 2024; Suk, 2024) analyzed the problem of dynamic regret. Furthermore, (Wanigasekara & Yu, 2019; Locatelli & Carpentier, 2018; Krishnamurthy et al., 2020; Zhu et al., 2022) discussed the case with continuous actions.

Nearest neighbor methods. Nearest neighbor classification has been analyzed in (Chaudhuri & Dasgupta, 2014; Döring et al., 2018) for bounded support of features. (Gadat et al., 2016; Kpotufe, 2011; Cannings et al., 2020; Zhao & Lai, 2021b, a) proposed adaptive nearest neighbor methods for heavy-tailed feature distributions. (Guan & Jiang, 2018) proposed kNN-UCB method for contextual bandits and proved a regret bound $\tilde{O}(T^{\frac{1+d}{2+d}})$.

Compared with existing methods on nonparametric contextual bandits, for unbounded contexts, the methods need to adapt to different density levels of contexts and achieve a better bias and variance tradeoff in the estimation of reward functions. Moreover, existing works on nonparametric classification can not be easily extended here, since the samples are no longer i.i.d and we now need to bound the regret instead of the estimation error. These factors introduce new technical difficulties in theoretical analysis. In this work, to address these challenges, we design new algorithms that are adaptive to both the pdf and the suboptimality of reward functions and provide a corresponding theoretical analysis.

Denote $\mathcal{X}$ as the space of contexts, and $\mathcal{A}$ as the space of actions. Throughout this paper, we discuss the case with infinite $\mathcal{X}$ and finite $\mathcal{A}$. At the $t$-th step, the context $\mathbf{X}_{t}$ is a random variable drawn from a distribution with probability density function (pdf) $f$. Then the agent takes action $A_{t}\in\mathcal{A}$ and receive reward $Y_{t}$:

in which $\eta_{a}(\mathbf{x})$ for $a\in\mathcal{A}$ and $\mathbf{x}\in\mathcal{X}$ is an unknown expected reward function, and $W_{t}$ denotes the noise, with $\mathbb{E}[W_{t}|\mathbf{X}_{1:t},A_{1:t}]=0$.

Throughout this paper, define

as the maximum expected reward of context $\mathbf{x}$. For any suboptimal action $a$, $\eta_{a}(\mathbf{x})<\eta^{*}(\mathbf{x})$. Correspondingly, $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$ is called suboptimality gap.

The performance of an algorithm is evaluated by the expected regret

We then present the assumptions needed for the analysis. To begin with, we state some basic conditions in Assumption 1.

Now we comment on these assumptions. (a) is the Tsybakov margin condition, which was first introduced in (Audibert & Tsybakov, 2007) for classification problems, and then used in contextual bandit problems (Perchet & Rigollet, 2013). Note that $\text{P}(0<\eta^{*}(\mathbf{X})-\eta_{a}(\mathbf{X})<t)\leq 1$ always hold, thus for any $\eta_{a}$, (a) holds with $C_{\alpha}=1$ and $\alpha=0$. Therefore, this assumption is nontrivial only if it holds with some $\alpha>0$. Moreover, we only consider the case with $\alpha\leq d$ here. If $\alpha>d$, then an arm is either always or never optimal, thus it is easy to achieve logarithmic regret (see (Perchet & Rigollet, 2013), Proposition 3.1). An additional remark is that in (Perchet & Rigollet, 2013; Reeve et al., 2018), the margin assumption is $\text{P}(0<\eta^{*}(\mathbf{X})-\eta_{s}(\mathbf{X})<t)\lesssim t^{\alpha}$, in which $\eta_{s}(\mathbf{x})$ is the second largest one among $\{\eta_{a}(\mathbf{x})|a\in\mathcal{A}\}$. Our assumption (a) is slightly weaker than existing ones since we only impose margin conditions on the suboptimality gap $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$ for each $a$ separately, instead of on the minimum suboptimality gap. In (b), following existing works (Reeve et al., 2018), we assume that the noise has light tails. (c) is a common assumption for various literatures on nonparametric estimation (Mai & Johansson, 2021). It is possible to extend this work to a more general Hölder smoothness assumption by adaptive nearest neighbor weights (Cannings et al., 2020). In this paper, we focus only on Lipschitz continuity for convenience.

Assumption 2 is designed for the case that the contexts have bounded support.

In Assumption 2, the pdf $f$ is required to be bounded away from zero, which is also made in (Perchet & Rigollet, 2013; Guan & Jiang, 2018; Reeve et al., 2018). Note that even for estimation with i.i.d data, this assumption is common (Audibert & Tsybakov, 2007; Döring et al., 2018; Gao et al., 2018).

We then show some assumptions for heavy-tailed distributions.

(a) is a common tail assumption for nonparametric statistics, which has been made in (Gadat et al., 2016; Zhao & Lai, 2021b). $\beta$ describes the tail strength. Smaller $\beta$ indicates that the context distribution has heavy tails, and vice versa. To further illustrate this assumption, we show several examples.

It worths mentioning that although the growth rate of the regret is affected by the value of $\beta$, our proposed algorithms including both fixed and adaptive methods do not require knowing $\beta$.

(b) restricts the suboptimality gap of each action. This is not necessary if the support is bounded. However, with unbounded support, without assumption (b), $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$ can increase appropriately with the decrease of $f(\mathbf{x})$, such that the regret of suboptimal action is large, and the identification of best action is hard.

Finally, we clarify notations as follows. Throughout this paper, $\left\lVert\cdot\right\rVert$ denotes $\ell_{2}$ norm. $a\lesssim b$ denotes $a\leq Cb$ for some constant $C$, which may depend on the constants in Assumption 2. The notation $\gtrsim$ is defined conversely.

In this section, we show the minimax lower bound, which characterizes the theoretical limit of regrets of contextual bandits. Throughout this section, denote $\pi:\mathcal{X}\times\mathcal{X}^{t-1}\times\mathbb{R}^{t-1}\rightarrow\mathcal{A}$ as the policy, such that each action is selected according to policy $\pi$. To be more precise,

which indicates that the action $A_{t}$ at time $t$ depends on the current context and the records of contexts and rewards in previous $t-1$ steps.

The minimax lower bound for the case with bounded support has been shown in Theorem 4.1 in (Rigollet & Zeevi, 2010). For completeness and notation consistency, we state the results below and provide a simplified proof.

We then show the minimax regret bounds for unbounded support, which is a new result that has not been obtained before.

From the results above, with $\beta\rightarrow\infty$, (6) reduces to (5).

For bounded context support, regret comes mainly from the region with $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})\lesssim T^{-1/(d+2)}$ (which is the classical rate for nonparametric estimation (Tsybakov, 2009)), in which the identification of best action is not guaranteed to be correct. However, with heavy-tailed contexts, regret may also come from the tail, i.e. the region with small $f(\mathbf{x})$, where the number of samples around $\mathbf{x}$ is not enough to yield a reliable best action identification. For heavy-tailed cases, i.e. $\beta$ is small, the regret caused by the tail region may dominate. This also explains why we need to use different techniques to derive the minimax lower bound for heavy-tailed contexts.

In the remainder of this paper, we claim that a method is nearly minimax optimal if the dependence of expected regret on $T$ matches (5) or (6). Following conventions in existing works (Rigollet & Zeevi, 2010; Perchet & Rigollet, 2013; Hu et al., 2020; Gur et al., 2022), currently, the minimax lower bounds are derived for contextual bandit problems with only two actions, thus we do not consider the minimax optimality of regrets with respect to the number of actions $|\mathcal{A}|$.

To begin with, we propose and analyze a simple nearest neighbor method with fixed $k$. We make the following definitions first.

Denote $n_{a}(t)=|\{i<t|A_{i}=a\}|$ as the number of steps with action $a$ before time step $t$. Let $\mathcal{N}_{t}(\mathbf{x},a)$ be the set of $k$ nearest neighbors among $\{i<t|A_{i}=a\}$. Define

as the $k$ nearest neighbor distance, i.e. the distance from $\mathbf{x}$ to its $k$-th nearest neighbor among all previous steps with action $a$.

With the above notations, we describe the fixed $k$ nearest neighbor method as follows. If $n_{a}(t)\geq k$, then

in which $b$ has a fixed value

If $n_{a}(t)<k$, then

Here we explain our design. If $n_{a}(t)\geq k$, then it is possible to give a UCB estimate of $\eta_{a}(\mathbf{x})$, shown in (9). $L\rho_{a,t}(\mathbf{x})$ bounds the estimation bias, while $b$ is an upper bound of the error caused by random noise that holds with high probability. In Lemma 5 in Appendix E, we show that $\hat{\eta}_{a,t}(\mathbf{x})$ is a valid UCB estimate of $\eta_{a,t}(\mathbf{x})$, i.e. $\hat{\eta}_{a,t}(\mathbf{x})\geq\eta_{a,t}(\mathbf{x})$ holds with high probability. If $n_{a}(t)<k$, then it is impossible to give a UCB estimate. In this case, we just let $\hat{\eta}_{a,t}(\mathbf{x})$ to be infinite.

Finally, the algorithm selects the action $A_{t}$ with the maximum UCB value:

According to (11), as long as an action has not been taken for at least $k$ times, the UCB estimate of $\eta_{a}$ will be infinite. Note that the selection rule (12) ensures that the actions with infinite UCB values will be taken first. Therefore, the first $k|\mathcal{A}|$ steps are used for pure exploration. In this stage, the agent takes each action $a$ for $k$ times. After $k|\mathcal{A}|$ steps, the UCB values for all $\mathbf{x}\in\mathcal{X}$ and $a\in\mathcal{A}$ become finite. Since then, at each step, the action is selected with the maximum UCB value specified in (9).

The procedures above are summarized in Algorithm 1. Compared with (Guan & Jiang, 2018), our method constructs the UCB differently. In (Guan & Jiang, 2018), the UCB is $\frac{1}{k}\underset{i\in\mathcal{N}_{t}(\mathbf{x},a)}{\sum}Y_{i}+\sigma(T_{a}(t-1))$, in which $\sigma(T_{a}(t-1))$ is uniform among all $\mathbf{x}$ with fixed action $a$.

Therefore, the method (Guan & Jiang, 2018) is not adaptive to the suboptimality gap $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$. On the contrary, our method has a term $L\rho_{a,t}(\mathbf{x})$ that varies for different $\mathbf{x}$, and thus adapts better to the suboptimality gap. The bound of regret is shown in Theorem 3.

We compare our result with (Guan & Jiang, 2018), which proposes a similar nearest neighbor method. The analysis in (Guan & Jiang, 2018) does not make Tsybakov margin assumption (Assumption 1(a)), and the regret bound is $\tilde{O}(T^{\frac{d+1}{d+2}})$. Without any restriction on $\eta$, Assumption 1(a) holds with $C_{\alpha}=1$ and $\alpha=0$, under which (13) reduces to $\tilde{O}(T^{\frac{d+1}{d+2}})$. Therefore, our result matches (Guan & Jiang, 2018) with $\alpha=0$. If $\alpha>0$, which indicates that a small optimality gap only happens with small probability, then the regret of the method in (Guan & Jiang, 2018) is still $\tilde{O}(T^{\frac{d+1}{d+2}})$, while our result improves it to $\tilde{O}(T^{1-\frac{\alpha+1}{d+2}})$. As discussed earlier, compared with (Guan & Jiang, 2018), our method improves the UCB calculation in (17), and is thus more adaptive to the suboptimalilty gap $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$. With $\alpha>0$, our method achieves smaller regret due to a better tradeoff between exploration and exploitation.

Compared with the minimax lower bound shown in Theorem 1, it can be found that the kNN method with fixed $k$ is not completely optimal. With $d>\alpha+1$, the upper bound matches the lower bound derived in Theorem 1. However, with $d\leq\alpha+1$, the regret is significantly higher than the minimax lower bound, indicating that there is room for further improvement.

We then analyze the performance for heavy-tailed context distribution. The result is shown in the following theorem.

From (15), there are two phase transitions. The first one is at $d=\alpha+1$, while the second one is at $d\beta=1$. Intuitively, the phase transition occurs because the regret is dominated by different regions depending on the settings $\alpha$ and $\beta$. Compared with the minimax lower bound shown in Theorem 2, it can be found that the kNN method with fixed $k$ achieves nearly minimax optimal regret up to logarithm factors if $d>\alpha+1$ and $\beta>1/d$, otherwise the regret bound is suboptimal. Here we provide an intuition of the reason why the kNN method with fixed $k$ achieves suboptimal regrets. In the region where the context pdf $f(\mathbf{x})$ is low, or the suboptimality gap $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$ is large, the samples with action $a$ are relatively sparse. In this case, with fixed $k$, the nearest neighbor distances are too large, resulting in a large estimation bias. On the contrary, if $f(\mathbf{x})$ is high or $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$ is small, then samples with action $a$ are relatively dense, thus the bias is small, and we can increase $k$ to achieve a better bias and variance tradeoff. Therefore, if $k$ is fixed throughout the support set, then the algorithm estimates the reward function $\eta_{a}(\mathbf{x})$ in an inefficient way, resulting in suboptimal regrets. Apart from suboptimal regret, another drawback is that with $d\leq\alpha+1$, the optimal selection of $k$ depends on the margin parameter $\alpha$, which is usually unknown in practice. In the next section, we propose an adaptive nearest neighbor method to address these issues mentioned above.

In the previous section, we have shown that the standard kNN method with fixed $k$ is suboptimal with $d\leq\alpha+1$ or $\beta\leq 1/d$. The intuition is that the standard nearest neighbor method does not adjust $k$ based on the pdf and the suboptimality gap. In this section, we propose an adaptive nearest neighbor approach. To achieve a good exploration-exploitation tradeoff and bias-variance tradeoff, $k$ needs to be smaller for small pdf $f(\mathbf{x})$ or large suboptimality gap $\eta^{*}(x)-\eta_{a}(x)$, and vice versa. However, as both $f(\mathbf{x})$ and $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$ are unknown to the learner, we need to decide $k$ based entirely on existing samples. The guideline of our design is that given a context $\mathbf{X}_{t}$ at time $t$, we use large $k$ if previous samples are relatively dense around $\mathbf{X}_{t}$, and vice versa. To be more precise, for all $\mathbf{x}\in\mathcal{X}$, let

in which $\rho_{a,t,j}(\mathbf{x})$ is the distance from $x$ to its $j$-th nearest neighbors among existing samples with action $a$, i.e. $\{i<t|A_{i}=a\}$. Such selection of $k$ makes the bias term $L\rho_{a,t,j}(\mathbf{x})$ matches the variance term $\sqrt{\ln T/j}$, thus (16) achieves a good tradeoff between bias and variance. The exploration-exploitation tradeoff is also desirable as $\rho_{a,t,j}(\mathbf{x})$ is large with large $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$, which yields smaller $k$. Note that (16) can be calculated only if $L\rho_{a,t,1}\leq\sqrt{\ln T}$, which means that the $1$-nearest neighbor distance can not be too large. At some time step $t$, for some action $a$, if there is no existing samples, or $\mathbf{X}_{t}$ is more than $\sqrt{\ln T}/L$ far away from any existing samples $\mathbf{X}_{1},\ldots,\mathbf{X}_{t-1}$, then we can just let the UCB estimate to be infinite, i.e. $\hat{\eta}_{a,t}(\mathbf{x})=\infty$. Otherwise, we calculate the upper confidence bound as follows:

in which $\mathcal{N}_{t}(\mathbf{x},a)$ is the set of $k_{a,t}(\mathbf{x})$ neighbors of $x$ among $\{i<t|A_{i}=a\}$, $\rho_{a,t}(\mathbf{x})$ is the corresponding $k_{a,t}(\mathbf{x})$ neighbor distance of $\mathbf{x}$, i.e. $\rho_{a,t}(\mathbf{x})=\rho_{a,t,k_{a,t}(\mathbf{x})}(\mathbf{x})$, and

Similar to the fixed nearest neighbor method, the last two terms in (17) cover the uncertainty of reward function estimation. The term $b_{a,t}(\mathbf{x})$ gives a high probability bound of random error, and $L\rho_{a,t}(\mathbf{x})$ bounds the bias. With the UCB calculation in (17), the $\hat{\eta}_{a,t}(\mathbf{x})$ is an upper bound of $\eta(\mathbf{x})$ that holds with high probability, so that the exploration and exploitation can be balanced well. The complete description of the newly proposed adaptive nearest neighbor method is shown in Algorithm 2.

We then analyze the regret of the adaptive method for both bounded and unbounded supports of contexts.

By comparing Theorem 5 with Theorem 3, it can be found that for the case with bounded support, the adaptive method improves over the fixed $k$ nearest neighbor method. From the minimax bound in Theorem 1, the fixed $k$ method is only optimal for $d\geq\alpha+1$, while the adaptive method is also optimal for $d<\alpha+1$, up to logarithm factors. An intuitive explanation is that with large $\alpha$, the suboptimality gap $\eta^{*}(\mathbf{x})-\eta_{a}(\mathbf{x})$ is small only in a small region, and the exploration-exploitation tradeoff becomes harder, thus the advantage of the adaptive method over the fixed one becomes more obvious.

We then analyze the performance of the adaptive nearest neighbor method for heavy-tailed distribution. The result is shown in the following theorem.

Compared with the minimax lower bound shown in Theorem 2, it can be found that our method achieves nearly minimax optimal regret up to a logarithm factor. Regarding this result, we have some additional remarks.

To begin with, to validate our theoretical analysis, we run experiments using some synthesized data. We then move on to experiments with the MNIST dataset (LeCun, 1998).

This paper analyzes the contextual bandit problem that allows the context distribution to be heavy-tailed. To begin with, we have derived the minimax lower bound of the expected cumulative regret. We then show that the expected cumulative regret of the fixed $k$ nearest neighbor method is suboptimal compared with the minimax lower bound. To close the gap, we have proposed an adaptive nearest neighbor approach, which significantly improves the performance, and the bound of expected regret matches the minimax lower bound up to logarithm factors. Finally, we have conducted numerical experiments to validate our results.

In the future, this work can be extended in the following ways. Firstly, following existing analysis in (Gur et al., 2022), it may be meaningful to design a smoothness adaptive method that can handle any Hölder smoothness parameters. Secondly, it is worth extending current work to handle dynamic regret functions. Finally, the theories and methods developed in this paper can be extended to more complicated tasks, such as reinforcement learning (Zhao & Lai, 2024).