Sharp Analysis for KL-Regularized Contextual Bandits and RLHF

Thanks for your constructive comments!

Q1. The main coverage condition (Definition 2.6) is extremely strong. This will scale at least with the number of contexts: if we take $\theta=\theta’$, and choose $b(x)$ for $x\ne x_0$ and $b(x_0)=B$ where $x_0$ is the minimum probability context under $d_0$, then the expression given in Definition 2.6 will scale as $1/d_0(x_0)\ge M$ for $M$ the number of contexts. Thus, the main sample complexity results of the paper (Theorem 3.3 and Theorem 4.4) really scale with the size of the context space. In general, it is not acceptable to obtain a sample complexity scaling with the size of the context space, as this is typically extremely large, so these results are only meaningful asymptotically as $\epsilon \rightarrow 0$.

A1. Thank you for pointing this out. This is actually a loose in our definition. Without changing the analysis, we can redefine the coverage condition as

Then, the situation you mentioned will not happen.

Q2. Furthermore, the result is also not tight in the regime where $\eta$ is very large.

A2. In the case you are referring to, $\eta$ should be $\Omega(1/ \epsilon)$, which is far from realistic if we are using KL regularization. Thus, a large $\eta$ is not the focus of our framework. Additionally, this problem has been solved by PAC Bayes literature, which suffers from similar issues when analyzing KL-regularization.

Q3. The statement on Line 294 that $D^2\le C_{GL}$ is not correct as a result of this ($C_{GL}$ in general will not scale with context size).

A3. After correcting our definition for data coverage (See A1.), $D$ will not scale with the context size. Thanks for the constructive comment. We will remove the claim that $D^2\le C_{GL}$ in the revision.

Q4. The statement in Theorem 3.1 and Theorem 4.1 that the coverage condition is $O(1)$ is also then incorrect—it should be $N_R(\epsilon)$ I believe.

We will remove the claim that the coverage condition is $O(1)$ in the revision and the coverage coefficient is left as a term in the final bound.

A4. Thanks for your suggestion. We will remove the claim in the revision.

Q5. Is the scaling with $D^2$ really necessary in Theorem 3.3 and Theorem 4.4 or can this be reduced to $C_{GL}$?

A5. Yes, it is necessary. According to our proof, the effectiveness of the intermediate policy highly relies on the data coverage coefficient. If the coverage condition is replaced by $C_{GL}$, the dependence will become multiplicative.

Thanks for your constructive comments!

Q1. The paper might benefit from some polishing. For instance, the definitions of some key terms are not rigorous (see questions below). In addition, the proof may lack rigor. For instance, line 320 – 321 uses Taylor expansion, and thus, if my understanding is correct, the equality (and the following inequality) does not hold.

A1. Thanks for your advice. Are you referring to 391-393 in Proof Sketch? According to our detailed calculation in 1113 - 1137 the equality and inequality do hold.

Q2. In definition 2.7, line 219 – 221, and definition 2.8, line 231 – 233, what does ``x sampled from d_0’’ mean in the sup?

A2. Thanks for pointing it out. They are typos. The subscription should be $x \in \text{supp}(d_0)$.

Q3. Can the authors provide more intuition why the number of samples needed are different in the first and the second stages of the algorithm, as presented in Theorem 3.3?

A3. According to our description in Section 3.2, by iteratively improving the data quality, the first stage is to establish a coarse estimate $\hat \theta_0$ of the reward function over a broad range of contexts. This requires enough diversity in the sampled data to capture global properties of the reward function. The second stage is to collect data that is more aligned with the optimal policy distribution, enabling fine-tuning of the reward function estimate so that the output policy is $\epsilon$-optimal. Therefore, they need different number of samples due to the different purposes.

Thanks for your constructive comments!

Q1. The global policy coverage coefficient can blow up to infinity, and the local policy coverage coefficient can be large (as the KL constraint is in expectation). The current paper only derives additive dependence results for global policy coverage, which can be vacuous if the global coefficient is infinity. On the other hand, the dependence on the local coefficient is still multiplicative (as discussed in Section 3.4), which can be extremely large.

A1. We would like to emphasize that currently there is no result with an additive dependence on the local coverage coefficient. Thus, the additive dependence is an unexpected novel result, but we can only derive the additive relationship for the newly-defined data coverage condition (Definition 2.6). For the local coefficient, it is difficult to further improve the multiplicative relationship and can be left as future work. Note that even for the multiplicative result, we have a faster rate.

Q2. I also have concerns regarding the sample complexity upper bounds. The paper claims to first study the effect of KL regularization in improving the sample complexity for policy optimization from $O(1 / \epsilon^2)$ to $O(1 / \epsilon)$. However, such $O(1 / \epsilon)$ sample complexity result already exists for general strongly convex regularizers [1], which include KL regularization as a special case since the reference policy is assumed to have sufficient coverage. Hence it is likely that the upper bound for CB is already known (given that CB is a special case of MDP). For RLHF, its difference from CB mainly comes from the additional reward learning step, so I expect there could be more explanation on why $O(1 / \epsilon)$ samples are sufficient for reward learning from preference data. However, the current version of the paper seems to lack such comparisons/remarks, which in my opinion are necessary for understanding the mechanism of RLHF (just as strong convexity of KL divergence for CB). For baselines, previous literature suggests that reward learning takes $O(1 / \epsilon^2)$ samples.

A2. We would like to clarify that this line of work for KL regularization is completely distinct from ours, because these works focus on the planning where the transition dynamics and reward model are known and focus on the pure policy optimization setting without exploration, while our work considers the statistical sample complexity where the underlying model is unknown and studies the statistical property of the learning problem. Hence, while they can achieve $O(1/t)$ rate in the planning setting, their methods cannot be applied to the learning setting.

[1] Lan, Guanghui. "Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and generalized problem classes." Mathematical programming 198, no. 1 (2023): 1059-1106.

[2] Zhu, Banghua, Michael Jordan, and Jiantao Jiao. "Principled reinforcement learning with human feedback from pairwise or k-wise comparisons." In International Conference on Machine Learning, pp. 43037-43067. PMLR, 2023.

Thanks for your insightful comments and positive feedback.

Q1. Cannot see the tradeoff between $\eta$ and the number of samples needed. Even from the experimental results, the lower $\eta$ the better the performance.

A1. Thanks for your insightful comments. In practice, with stronger KL regularization (i.e. lower $\eta$), the learned policy will be constrained in a small interval close to the reference policy which may limit the model's ability to fit the human feedback signal. Essentially, the reason is that the optimal solution for the regularized objective may not be good enough. However, in our paper, we consider the sample size required to approximately maximize the KL-regularized objective in both theoretical results and experimental results. Hence, it is beyond our scope how the KL-regularization coefficient affects the quality of the optimal solution. As a result, there is no tradeoff on $\eta$ from the perspective of maximizing the regularized objective.

In our theoretical analysis, since the objective we consider is the reward regularized by KL divergence, the considered optimal policy is also the optimal one with the $\eta$-KL ball around $\pi_0$. Hence, if we focus on the suboptimality with respect to such an optimal policy, the smaller $\eta$ becomes, the smaller the policy class we consider, thus the sample complexity is smaller.

Q2. The $O(1)$ coverage assumption is unclear if it is a reasonable one or not.

A2. It is a reasonable assumption since first, it is a standard condition in reinforcement learning literature as discussed in Lines 270-277; second, since the reference policy is obtained from supervised fine-tuning, it is natural to assume that it has a good coverage over the policy class. For details and the example, please refer to the 1st point of official comment.

Q3. In line 294 it says that ‘it is obvious that $D^2 \le C_{GL}$’. Could you please provide a proof for that?

A3. Thanks for the constructive comment. After reviewing the conditions, we realize that $D^2 \le C_{GL}$ does not always hold. We will remove this claim in the revision.

Q4. In line 374 is $s_i \sim \pi_0$ a typo, where does $s_i$ come from?

A4. Thanks for pointing out! It should be $a_i$.

Thanks to the reviewer for their response. While the newly proposed coverage condition is tighter than the previous one, it still seems to not capture the correct dependence. For example, assume we have some $\theta, \theta', \tilde{\theta}, \tilde{\theta}'$ such that for all $\bar{\theta}, \bar{\theta}' \in \{ \theta, \theta', \tilde{\theta}, \tilde{\theta}' \}$, we have $R(\bar{\theta},x,a) = R(\bar{\theta}',x,a)$ for all $(x,a)$ with $d_0(x) > 0$, and that for some $x'$ with $d_0(x') = 0$ (or arbitrarily small), we have that $R(\theta,x',a) - R(\theta',x',a) \neq R(\tilde{\theta},x',a) - R(\tilde{\theta}',x',a)$. Then in this case the denominator for both sets $(\theta,\theta')$ and $(\tilde{\theta}, \tilde{\theta}')$ will be 0, but there does not exist a $b(x')$ that will make the numerator 0 for both, so it follows that the coverage condition will be infinite. However, this should not be the case since $d(x') = 0$, so this $x'$ should not influence the complexity ideally.

For the linear example given here, I am not sure the analysis is correct. By the same argument as what I just gave above, I believe $D$ should depend on the minimum eigenvalue of $\Sigma$, and furthermore the $b(x)$ given here depends on $\theta$, which is not consistent with the given definition.

If the authors could comment on this, that would be helpful.

Thank you very much for your timely response.

For your first question, let $\mathcal{X}\_1 = \\{x \in \mathcal{X}| \exists a\ s.t. R(\tilde\theta, x, a) \neq R(\tilde\theta', x, a)\\}$ be the subset consisting of $x'$ you are referring to. If $d(\mathcal{X}_1) = 0$, then conventionally $\mathcal{X}_1$ should be excluded from $\mathcal{X}$ since we can never encounter the samples in $\mathcal{X}_1$, and thus it is impossible to differentiate $\tilde\theta$ and $\tilde\theta'$. You may argue that we can set $d(\mathcal{X}_1)$ to an extremely small number. In that case, we could only learn $\theta$ based on the samples with $x \in \mathcal{X}_1$. So, it is indeed a hard instance, and it is reasonable that the covering coefficient is large.

For the linear example, we apologize for not mentioning that $b$ could vary under different $\theta, \theta'$, and the formal definition should be as follows.

For any pair of $\theta, \theta'$, there exist $b:\mathcal{X}\rightarrow[-B,B]$ such that

For the G-optimal design technique, which can control the elliptical norm of a given set of feature vectors, please refer to Chapter 21.1 of [1].

[1] T Lattimore, C Szepesvári. Bandit Algorithms.