Small steps no more: Global convergence of stochastic gradient bandits for arbitrary learning rates

We appreciate that the reviewer recognized the contribution of the work. We answer the questions as follows.

include some literature review on deep learning with large learning rates

Thank you for pointing the "edge of stability" paper to us. We will cite this line of work in the related work discussion starting from Line 104 in the paper.

what the exact convergence rate is

Please find a detailed discussion in the common rebuttal, regarding why the techniques in [2] are not applicable and cannot be used here to obtain a rate.

The code is not provided

We will upload a link for running the simulations in the updated version of the paper.

Can we also have similar results for the version with baselines?

Thank you for asking this interesting question.

First, our preliminary calculations show that similar results can be obtained for the version with action-independent baselines, such as $\pi_{\theta_t}^\top r$.

Second, in Algorithm 1, if we replace $R_t(a_t)$ with $R_t(a_t) - \pi_{\theta_t}^\top r$ (or $R_t(a_t) - B_t$ with any other action-independent baseline $B_t \in \mathbb{R}$), then Lemmas 1 and 2 still hold because of $R_t(a_t) - \pi_{\theta_t}^\top r$ is still bounded (by Eq. (1)). Theorem 1 will hold because of the progress is exactly the same as in Eqs. (11) and (13), according to the well known result that action-independent baselines contribute $0$ to policy gradient, meaning that Proposition 1 holds for any action-independent baselines.

What would be the challenges of extending the current analysis to the non-asymptotic rate analysis?

First, the challenge is that monotonic improvement (in expectation) over $\pi_{\theta_t}^\top r$ cannot be shown for large learning rates (non-monotonicity of $\pi_{\theta_t}^\top$ is observed in simulations in Fig. 1 in the main paper, also as pointed in Line 304), unlike Lemma 4.6 of [23] (when learning rate is very small). Therefore, it is not clear how to quantify the progress (how much $r(a^*) - \pi_{\theta_t}^\top r$ is reduced in expectation after one stochastic gradient step). Without quantifying the (expected) progress, it seems difficult to do non-asymptotic rate analysis.

Second, there exist analyses for non-monotonic improvement over objective functions (such as Nesterov's accelerated gradient), but for convex functions. However, here we have a non-concave maximization (Line 95).

Please also find a detailed discussion in the common rebuttal.

what would be the challenges of extending the current proof techniques to other RL algorithms like natural gradient methods?

Natural gradient methods achieve faster convergence results than standard policy gradient when using exact gradient updates. However, with stochastic on-policy sampling $a_t \sim \pi_{\theta_t}(\cdot)$, if we use constant learning rates $\eta \in \Theta(1)$ for natural policy gradient method, then it can fail, by converging to sub-optimal deterministic policies with positive probability (Theorem 3 of [19] and Proposition 1 of Chung2021).

[Chung2021] Beyond variance reduction: Understanding the true impact of baselines on policy optimization, in ICML 2021.

We appreciate that the reviewer understood and recognized the contribution of the work. We answer the questions as follows.

$\eta = 100$ and $\eta = 1000$ do not converge. Can you discuss this phenomenon in detail?

We ran more iterations for $\eta = 100$ and $\eta = 1000$, and eventually all runs converged. Please see Fig. 1 in the rebuttal pdf for results. A new observation is that those curves converge when the optimal action is sampled, i.e., $a_t = a^*$. This is aligned with theory, since the first part of Theorem 2 (Line 241) proved that the optimal action will be sampled for infinitely many times as $t \to \infty$.

larger $K = 10$

We ran experiments using $K = 10$ as suggested for learning rates $\eta \in \\{1, 10\\}$, extending the example in the main paper to $r = (0.2, 0.05, -0.1, -0.4, -0.5, -0.6, -0.7, -0.8, -0.9, -1.0)^\top$. Please see Fig. 2 in the rebuttal pdf for results.

It would be helpful to discuss the challenge when using the technique in [1] for a constant learning rate.

First, monotonic improvement over $\pi_{\theta_t}^\top r$ cannot be established for large learning rates. In fact, for large learning rates, non-monotonicity of $\pi_{\theta_t}^\top$ is observed in Fig. 1 (metioned in Line 304).

Second, in particular, Lemma 4.6 in [1] requires very small constant learning rates ($\eta \le \frac{\Delta^2}{40 \cdot K^{3/2} \cdot R_{\max}^3 }$), which follows from their Lemma 4.2 (smoothness) and Lemma 4.3 (noise growth condition). Using large learning rates makes those two lemmas not applicable, and monotonic improvement of $\pi_{\theta_t}^\top r$ is no longer assured.

We appreciate that the reviewer understood and recognized the contribution of the work, and we thank the reviewer for carefully reading and checking the results. The main concerns are addressed as follows.

detailed explanations for Lemma 1 and Lemma 2

According to Eq. (1), the sampled reward $R_t(a_t)$ is in a bounded range of $[-R_{\max}, R_{\max}]$ with $R_{\max} < \infty$. By design, we use a constant learning rate $\eta < \infty$. This argument is shown in Eqs. (26) and (27) for Lemma 1 in the appendix.

elaborate on the observed differences in convergence behavior for these learning rates

From the simulations, smaller $\eta$ values such as $1$ and $10$ always enter the final stage of $\pi_{\theta_t}(a^*)$ being close to $1$ more quickly, while larger $\eta$ values ($100$ and $1000$) can reduce the sub-optimality gap $r(a^*) - \pi_{\theta_t}^\top r$ to much smaller values when $\pi_{\theta_t}$ is already in the final stage of $\pi_{\theta_t}(a^*) \approx 1$. However, there is a trade off, such that larger $\eta$ values (on average) spend a longer time entering the final stage.

learning rates outside the tested range

We ran experiments on the same example in the paper using $\eta = 0.01$ and $\eta = 0.1$, which outside the tested range of $\{1, 10, 100, 1000\}$ (since $1000$ is already a super large learning rate, we did not test $\eta > 1000$). Please see Fig. 3 in the rebuttal pdf for results.

How would you recommend practitioners choose an appropriate learning rate...Are there any heuristics or rules of thumb...?

In deep learning, large learning rates have been applied empirically. For example, [Li2020] showed that increasing learning rates can achieve good performance for networks with BatchNorm, which is similar to what we speculated for RL (Line 345).

However, multiple factors affect practical performance. Our heuristic is that when the optimal solution is deterministic (no entropy regularization), then stochastic policy gradient could work with large learning rates. If there is entropy regularization, then eventually the learning rate has to be decayed, since otherwise oscillation can happen (optimal solution is not deterministic, and overshooting can happen).

[Li2020] An Exponential Learning Rate Schedule for Deep Learning, in ICLR 2020.

What challenges do you foresee in extending these results to more general reinforcement learning settings? How might the approach change?

First, currently we do not know if the key properties we established for bandits (Lemmas 1 and 2) will hold in MDPs or not, which is the key challenge of extending to general reinforcement learning settings.

Second, in general MDPs, exploration over the state space is a problem that does not exist in bandits [1], which requires the stochastic gradient methods to be modified accordingly.

Assumption 1 ... How critical is this assumption ... how to relax this assumption in future work?

First, Assumption 1 is for the simplicity of our arguments. It shouldn't be critical, since one can collapse actions with the same rewards without changing the optimal expected reward / value.

Second, regarding how to relax this assumption, please refer to the common rebuttal for a more detailed discussion.

Could you suggest specific directions or methods for this refined analysis

We speculate that the established asymptotic convergence of $\pi_{\theta_t}(a^*) \to 1$ as $t \to \infty$ can be split into two phases: early and final stages, corresponding to when $\pi_{\theta_t}(a^*) < 1 - \epsilon$ or $\pi_{\theta_t}(a^*) \ge 1 - \epsilon$ respectively.

The phase transition happens at the time point $t_0 < \infty$ such that for all $t \ge t_0$, we have $\pi_{\theta_t}(a^*) \ge 1 - \epsilon$. Therefore, a refined analysis would be to characterize how the time $t_0$ depends on learning rates, as well as other problem specific quantities, such as reward gap and reward range.

Could you provide more insights or conjectures about the expected rate of convergence for different learning rates?

First, for all learning rates, from Figs. 1(a) and 1(b) in the main paper, $\log{ (r(a^*) - \pi_{\theta_t}^\top r) } \approx - \log{t} + c$ (from numerical values, slopes are $\approx -1$), which means that $r(a^*) - \pi_{\theta_t}^\top r \approx C/t$, i.e., the rate of convergence is about $O(1/t)$ (as mentioned in Line 315).

Second, for different learning rates, we speculate that the time spent in early stage scales with $\eta$, such as $O(\eta)$, while the rate in final stage scales with $1/\eta$, such as $O(1/(\eta \cdot t))$, aligned with previous observed differences in convergence behavior for different learning rates in Fig. 1 in the main paper.

We thank the reviewer for taking time to review our work. We hope the following can help clarify matters.

heavily relies on prior work [23]

This is simply incorrect. Our analysis is significantly different from [23], which is built on smoothness and growth condition, while our results do not exploit. Please refer to the common rebuttal for a detailed explanation.

In simulations, if the learning rate is initially set to large number and then decreased, does algorithm show fast convergence?

First, we ran experiments using adaptive decaying learning rates as suggested on the same example in the paper. The total number of iterations is $T = 2 \times 10^6$. For the first $10^6$ iterations, we use a linearly decayed learning rate from $\eta = 1000$ (or $\eta = 100$) to $\eta = 1$ at $t = 10^6$, i.e., $\eta_t = \mathbf{1000} \times (1 - t/10^6) + t/10^6$ (or replace $\mathbf{1000}$ with $\mathbf{100}$) for all $t \in [1, 10^6]$. For the second $10^6$ iterations, we decay the learning rate as $O(1/\sqrt{t})$, i.e., $\eta_t = 1/\sqrt{t-10^6}$ for all $t \in [10^6 + 1, 2 \times 10^6]$. Please see Fig. 4 in the rebuttal pdf for results.

Second, we believe that this decaying learning rate scheme will lead to slower convergence speeds than using constant learning rates, as suggested by the analysis carried out in the paper of Lu2024.

[Lu2024] Towards Principled, Practical Policy Gradient for Bandits and Tabular MDPs.

We appreciate that the reviewer understood and recognized the contribution of the work. We answer the questions as follows.

Please refer to the common rebuttal for questions regarding Assumption 1, rate of convergence, and comparison with [23].

...maximizing the cumulative reward (or minimizing the expected regret). Can we apply the stochastic gradient bandit algorithm in that setting?

In fact, if we can obtain a rate of convergence (say $O(1/t)$ as speculated in the common rebuttal), such that $\mathbb{E}{ [ r(a^*) - \pi_{\theta_t}^\top r]} \le C/t$, then this will imply a expected regret upper bound by summing up sub-optimality gap, $\sum_{t=1}^{T}{ \Big( r(a^*) - \mathbb{E}{[\pi_{\theta_t}^\top r]} \Big) } \le \sum_{t=1}^{T}{C/t} \le C \cdot ( \log{T} + 1 ),$ which means that to get such a result, it is sufficient to prove a rate of convergence result.

adaptive learning rates...decreased systematically for consequent batches?

First, we ran experiments using adaptive decaying learning rates on the same example in the paper. The total number of iterations is $T = 2 \times 10^6$. For the first $10^6$ iterations, we use a linearly decayed learning rate from $\eta = 1000$ (or $\eta = 100$) to $\eta = 1$ at $t = 10^6$, i.e., $\eta_t = \mathbf{1000} \times (1 - t/10^6) + t/10^6$ (or replace $\mathbf{1000}$ with $\mathbf{100}$) for all $t \in [1, 10^6]$. For the second $10^6$ iterations, we decay the learning rate as $O(1/\sqrt{t})$, i.e., $\eta_t = 1/\sqrt{t-10^6}$ for all $t \in [10^6 + 1, 2 \times 10^6]$. Please see Fig. 4 in the rebuttal pdf for results.

Second, we believe that this decaying learning rate scheme will lead to slower convergence speeds than using constant learning rates, as suggested by the analysis carried out in the paper of Lu2024.

[Lu2024] Towards Principled, Practical Policy Gradient for Bandits and Tabular MDPs.

Intuitively the right learning rate should depend on the underlying difficulty of the reward generating mechanism for different arms, however I am not sure if that is apparent from the discussion in the paper.

Thank you for pointing this out. We agree with the reviewer on their intuition and insights.

First, we would re-emphasize that the asymptotic convergence result is the very first result for using large learning rates in stochastic gradient bandits, which is already difficult to obtain.

Second, asymptotic convergence result itself without any characterization of rates is not enough to tell the difference in dependence on problem difficulty. As mentioned in Line 335 in the paper, a more refined characterization of convergence behaviors is needed to reflect this dependence.

Can this be extended to the contextual bandits framework?

Thank you for asking, and we are also pursuing this interesting direction. It is doable but not immediate. Extra work needs to be done for analyzing stochastic gradients using linear features (rather than the softmax tabular parameterizaion here), which is of independent of interest.