Paper: Accelerating Distributed Stochastic Optimization via Self-Repellent Random Walks (ICLR-2024)

Reviewer comment (Theory):
A concern in my mind is that the paper assumes $X$ takes values in a finite state space. The decentralized optimization where only a single agent is active each time is a particular example. Are there any other examples where $X$ takes finite value?

Author response:
Our SA-SRRW algorithm is versatile and not limited to decentralized optimization employed with token algorithm only. It is equally applicable to scenarios in stochastic optimization where the entire dataset is fully accessible. This application is particularly prevalent in the machine learning domain, such as in training neural networks using SGD or other accelerated algorithms. In these instances, the dataset forms a *complete graph with self-loops* so that $X$ can be independently and identically sampled from the dataset with marginal distribution $\mu$, implying that these instances can also be transformed into the decentralized optimization in the broader sense. Here, the baseline Markov chain $\mathbf{P}$, as indicated in the SRRW kernel (3), becomes a rank-1 matrix $\mathbf{1}\mathbf{\mu}^T$. 

Meanwhile, $X$ cannot be a continuous random variable. The limitation here is not related to the Poisson method used in our proofs but is inherent to the structure of the SRRW employed in the SA method. Specifically, the concept of self-repellence in SRRW requires counting the visits to each state $X$ to form the empirical measure $\mathbf{x}_n$ at time $n$. This count is crucial for designing the SRRW kernel $\mathbf{K}[\mathbf{x}_n]$. In a continuous state space, such a counting mechanism becomes infeasible since the chance of sampling the same $X$ in the continuous state space is measure zero, posing a significant challenge to the design and implementation of a continuous version of SRRW in our algorithm.

Reviewer comment (Writing):
Under the paragraph below (3), the author tries to explain why SRRW gets its name. More specifically, they said “if node $j$ has been visited more often, so far, the entry $x_j$ becomes larger (than target value $\mu_j$)”. I didn’t understand the logic behind it. If node $j$ has been visited more often, I think $x_j$, as the coordinate of a distribution vector, should be close to $\mu_j$. I can’t tell whether $x_j$ is larger than $\mu_j$ or not. However, I acknowledge that the nonlinear kernel in (3) would encourage the algorithm to visit the states that are less visited, based on its expression.

Author response:
Your interpretation ''$x_j$ should be close to $\mu_j$’ is indeed accurate in the context of long-term behavior, where the empirical distribution $\mathbf{x}$ converges towards the target distribution $\mu$ almost surely. However, the point of confusion seems to arise from our description 'if node $j$ has been visited more often so far'. To clarify, when saying 'node $j$ has been visited more often', we mean that the node $j$ has been visited more often than other nodes (and we have improved this explanation in the revision). In this context, according to the definition of empirical distribution, the entry $[\mathbf{x}_n]_j$ – representing the proportion of visits to node $j$ at time $n$ – can be (temporarily) larger than its target value $\mu_j$. In this case, node $j$ will be chosen with smaller probability at time $n+1$ by the SRRW kernel with $\mathbf{x}_n$ in place of $\mathbf{x}$ in (3). Note that, eventually, from the SRRW convergence result, all these $[\mathbf{x}_n]_j$ will converge to the target $\mu_j$ as n increases, as you pointed out already. In some sense, SRRW tries to exploit such temporal fluctuation (while keeping the stationary behavior the same) to 'force' the random walker to reduce such temporal fluctuation (deviation of $\mathbf{x}$ at time $n$ from the target $\mu$), which translates to smaller sampling variance.

Reviewer comment (Theory):
The newly introduced iterate in (4b) generalizes the original algorithm and allows us to treat the resulting algorithm as a two-time-scale SA. However, I didn’t see much motivation to introduce the new iterate. Could the author elaborate more on the motivation?

Author response:
We thank the reviewer for commenting on the SRRW iterate (4b). This SRRW iterate generalizes the step size $\gamma_n$ from $1/(n+1)$ to a general form $1/(n+1)^a$  for $a\in(0.5,1]$, which is crucial for enhancing the algorithm’s flexibility and performance. This generalization allows for the exploration of the step sizes $\beta_n=1/(n+1)^b =o(\gamma_n)$ (case(i)), which is suggested as the best algorithmic choice among all three cases in simulation. Without this generalization, i.e., if we stick to the original step size of $\gamma_n=1/(n+1)$ as in [Doshi et al. 2023], the case (i) becomes vacuous as it is impossible to choose the step size $\beta_n=o(1/n)$ (such step size becomes summable, fundamentally violating the condition for the step size in any stochastic approximation setting). To reflect this response, we have added a new footnote #7 in the revision to elaborate on the motivation for the generalized step size in the SRRW iterates.

>Doshi, V., Hu, J., & Eun, D. Y. (2023). Self-Repellent Random Walks on General Graphs--Achieving Minimal Sampling Variance via Nonlinear Markov Chains. International Conference on Machine Learning. PMLR, 2023.

Reviewer comment (Theory):
If we set $\alpha=\infty$, from (6), the variance matrix would be zero. Could I say, in this case, $x_n$ actually converges to $\mu$ at a faster rate than $\gamma_n$ so that the CLT degenerates to zero? How to explain the intuition behind it?

Author response:
It is crucial to clarify that our CLT results, as presented in Theorem 3.3, are specifically tailored for any 'given’ $\alpha < \infty$, where $\mathbf{x}_n$ converges to $\mu$ at the exact asymptotic rate $\gamma_n$ and the resulting covariance tends to zero for larger $\alpha$ without actually reaching it. It's important to understand that our analysis does not extend to the case of $\alpha=\infty$ since the SRRW kernel in (3) is undefined. Consequently, the behavior of $\mathbf{x}_n$ in such a setting, including its convergence rate to $\mu$ and the resulting asymptotic covariance, falls outside the scope of our theoretical framework.

Reviewer comment (Experiments):
From the simulation, it seems the larger $\alpha$, the faster convergence. Is there any suggestion for a practical choice of $\alpha$?

Author response:
In practice, we recommend using moderate values of $\alpha$, e.g., $\alpha=5$. Since covariance in case (i) decreases at a rate of $O(1/\alpha^2)$, even moderate values of $\alpha$ are sufficient to achieve small errors, e.g., $\alpha=5$ in Figure 2(c). Moreover, it’s important to consider the potential drawbacks of setting $\alpha$ too large. As illustrated in Figure 2(b), larger $\alpha$ values could slow down convergence in the early stage of the training process. For example, when $\alpha$ is set to $10$ or $20$, performance is slightly inferior to $\alpha=5$ until $n\approx 5×10^4$.

Reviewer comment (Theory):
If the $\alpha$ is also updated, for example, we use $\alpha_t$ at the iterate $t$ that increases at a carefully selected rate in advance (or perhaps could be updated using a third iterate). Can we accelerate the convergence rate of $\theta_n$ so that its mean square error is faster than $\beta_n$?

Author response:
Assume $\alpha_n$ is updated through a third iterate and converges to some finite value $\alpha^*$, we speculate that the asymptotic rate of MSE cannot go faster than $\beta_n$. Our current analysis, as detailed in Theorem 3.3, addresses scenarios with a fixed $\alpha$ value. Here, the influence of the SRRW sequence $\{X_n\}$ is encapsulated within the matrix $\mathbf{U}$, which contributes to the asymptotic covariance $\mathbf{V}_{\theta}$ of $\theta_n$. Introducing a time-varying $\alpha_n$ adds a new dimension to the system, potentially transforming the SA-SRRW algorithm, which we analyze via *two-timescale* SA setting, into a *three-timescale* SA. The CLT analysis of such a complex scenario remains unexplored and presents an interesting avenue for future research. Intuitively speaking, in the asymptotic regime, as $\alpha_n$ already approaches $\alpha^*$, the behavior of sequence $\{X_n\}$ would likely be similar to that driven by SRRW with constant $\alpha^*$. This means that while varying $\alpha_n$ impacts the asymptotic covariance $\mathbf{V}_{\theta}$ through matrix $\mathbf{U}$, the convergence rate of $\theta_n$, determined by $\beta_n^{-1/2}$, would remain unchanged. Hence, even with a dynamic $\alpha_n$, the MSE of $\theta_n-\theta^*$ is still expected to converge at a rate tied to $\beta_n$, according to our current theoretical framework.

Reviewer comment (Theory):
Could the authors comment a little bit more on the stability of the iterates? This is assumed to always hold, but as the authors point out, practical tricks are required to ensure this, in which case the theoretical analysis might also need tweaks.

Author response:
Answer: We appreciate the reviewer's question regarding the stability of the iterates in our SA-SRRW algorithm, and we recognize the potential impact of practical tricks on our theoretical analysis. The stability of SRRW iterates $\mathbf{x}_n$ in (4b) is always ensured through the truncation method in [1], where the detailed steps have been shown in [2]’s Appendix E with $\gamma_n=1/(n+1)$. This analysis can be extended to generalized step size $\gamma_n = 1/(n+1)^a$ for $a \in (0.5,1]$ as well since it satisfies the step size assumption (A1) in [1]. 

However, the stability of the SA iterates $\theta_n$ presents a more complex challenge. This is due to its dependency on the SRRW iterates $\mathbf{x}_n$. Currently, there is a lack of comprehensive stability analysis for $\theta_n$ within the domain of two-timescale SA with controlled Markov noise. Even the latest result in this area, which provides only the almost sure convergence (no CLT result therein), still requires the stability assumption [3]. This situation mirrors the challenges faced in single-timescale SA literature, where establishing stability conditions often requires significant analytical effort, e.g., Chapter 4 in [4], [1], [5], [6]. 

Therefore, while we ensure the stability of the SRRW iterates using established methods, the stability of $\theta_n$ under the two-timescale framework remains an open question. Addressing this would not only require modifying our theoretical analysis to accommodate practical algorithmic adjustments but also contribute significantly to the broader field of two-timescale SA, which we acknowledge as an important aspect of future work.

[1]. Andrieu, C., Moulines, É., & Priouret, P. Stability of stochastic approximation under verifiable conditions. SIAM Journal on control and optimization, 44(1), 283-312, 2005.

[2]. Doshi, V., Hu, J., & Eun, D. Y. Self-Repellent Random Walks on General Graphs--Achieving Minimal Sampling Variance via Nonlinear Markov Chains. International Conference on Machine Learning. PMLR, 2023.

[3]. Yaji, V. G., & Bhatnagar, S. Stochastic recursive inclusions in two timescales with nonadditive iterate-dependent Markov noise. Mathematics of Operations Research, 45(4), 1405-1444, 2020.

[4]. Borkar, V. Stochastic Approximation: A Dynamical Systems Viewpoint: Second Edition. Texts and Readings in Mathematics. Hindustan Book Agency, 2022.

[5]. Fort, G., Moulines, E., Schreck, A., and Vihola, M. Convergence of markovian stochastic approximation with discontinuous dynamics. SIAM Journal on Control and Optimization, 54(2):866–893, 2016.

[6]. Yaji, V. G., & Bhatnagar, S. Analysis of stochastic approximation schemes with set valued maps in the absence of a stability guarantee and their stabilization. IEEE Transactions on Automatic Control,65(3), 1100–1115, 2020.

Reviewer comment (Experiments):
Are there any tradeoffs in picking $\alpha$? Should $\alpha$ always be as large as possible, or will that affect other aspects of the algorithm, e.g. stability?

Author response:
Answer: While a larger $\alpha$ does indeed favor a smaller asymptotic covariance as per our CLT, there are practical considerations that suggest a more balanced approach to choosing $\alpha$. In practice, a moderate value of $\alpha$ often strikes the right balance between reducing asymptotic covariance and maintaining efficient convergence in the initial phases of training. For instance, as observed in our simulations (Figure 2(b)), we noticed that when $\alpha$ is set to larger values like 10 or 20, the performance initially lags behind that of $\alpha=5$ until about $n\approx 5×10^4$ iterations. In addition, asymptotic covariance in case (i) decreases at a rate of $O(1/\alpha^2)$, implying that even moderate values of $\alpha$ are sufficient to achieve small errors, e.g., $\alpha=5$ in Figure 2(c). 

As explained in the previous response, the stability of SRRW iterates can always be guaranteed through truncation method in [1] for any $\alpha<\infty$. Thus, although we do not have the stability of $\theta_n$, due to the robust behavior of SRRW iterates, we speculate that large $\alpha$ is unlikely to compromise the stability of $\theta_n$.

Reviewer comment (Theory):
While the analysis is on asymptotic covariance, could the authors comment on the likely finite-time convergence rate behavior of their SRRW-SA algorithm, and whether improvements (due to the SRRW) can be also shown in terms of finite-time convergence rates?

Author response:
Answer: The current non-asymptotic analysis of two-timescale SA with controlled Markov noise cannot be applied to our SA-SRRW algorithm. Specifically, all the state-of-the-art non-asymptotic analysis requires that the mean field functions of both iterates are globally Lipschitz [7-9]. This globally Lipschitzness condition is in place even for finite-time bounds in the singe-timescale SA [10-12]. However, the mean field function $\boldsymbol{\pi}(\mathbf{x})-\mathbf{x}$ of SRRW iterates (4b) exhibits only locally Lipschitz continuity, where $\mathbf{x}$ is in the interior of the probability simplex. This is due to the polynomial form $(x_i/μ_i)^{-\alpha}$ inside $\boldsymbol{\pi}(\mathbf{x})$, which was elaborated on in Appendix D [2] as a necessary condition to ensure scale invariance and local information reliance. This distinction is critical as it places our SA-SRRW algorithm outside the existing non-asymptotic analysis. To address this, we include additional insights in the paragraph following (4) in our revised manuscript and provide a detailed discussion in the new appendix B.

[7]. Doan, T. T. Finite-time convergence rates of nonlinear two-time-scale stochastic approximation under markovian noise. arXiv preprint arXiv:2104.01627, 2021.

[8]. Zeng, S., Doan, T. T., and Romberg, J. A two-time-scale stochastic optimization framework with applications in control and reinforcement learning. arXiv preprint arXiv:2109.14756, 2021.

[9]. Doan, T. T. Nonlinear two-time-scale stochastic approximation convergence and finite-time performance. IEEE Transactions on Automatic Control, 2022.

[10]. Chen, Z., Maguluri, S. T., Shakkottai, S., & Shanmugam, K. Finite-sample analysis of stochastic approximation using smooth convex envelopes. arXiv preprint arXiv:2002.00874, 2020.

[11]. Sun, T., Sun, Y., & Yin, W. On markov chain gradient descent. Advances in neural information processing systems, 31, 2018.

[12]. Even, M. Stochastic gradient descent under markovian sampling schemes. International Conference on Machine Learning, 2023.

Reviewer comment (Theory):
The authors suggest that $\beta_n = o(\gamma_n)$ is the best step-size choice if a $O(1/\alpha^2)$ decay rate in the asymptotic covariance is desired. However, if faster convergence is to be desired, a larger learning rate (i.e. larger $\beta_n$) should always be picked, i.e. we would want to pick the maximal possible step-size required by the assumptions in the paper, i.e. $\beta_n = O(1/n^{0.5+\epsilon})$ for some $\epsilon > 0$ as close to 0 as possible, in which case $\gamma_n$ can at best match the step-size of $\beta_n$, meaning that we can only get $O(1/\alpha)$ in the asymptotic covariance decay again. Could the authors comment on this possible issue?

Author response:
Answer: In our paper, $\beta_n=o(\gamma_n)$ is shown to be the best step-size choice in terms of achieving smaller asymptotic error (captured via asymptotic covariance). These errors are smallest for choices of $\beta_n=O(1/n)$, and not $O(1/n^{0.5+\epsilon})$ for some close-to-zero $\epsilon$, since the step sizes are also the scaling factors under which the finite asymptotic variance is characterized. Thus, larger learning rates actually lead to larger asymptotic variances as well; even though the initial convergence (to a level set around the solution) could be quicker, larger learning rates result in larger variances (larger level sets) for large enough time. This is also demonstrated by the extreme case of constant step size, where, without additional averaging, the SGD iterates may converge quickly, but only to a level set (neighborhood) around the solution, and not to the solution itself, and it is well known that the size of such neighborhood around the solution critically depends on the (asymptotic) variance of the iterates. Therefore, selecting the largest possible learning rate is not necessarily the most desirable.

Reviewer comment (Theory):
It would be nice if the authors provided some intuition as to why the $\beta_n = o(\gamma_n)$ step-size ratio can yield the $O(1/\alpha^2)$ rate. Moreover, it would be nice if the authors spelled out precisely the dependence on the difference between the step-size exponents in their asymptotic rates (i.e. if  $\beta_n = n^{-b}$ and $\alpha_n = n^{-a}$, how the rate actually depends on the difference between $a$ and $b$). If this difference does not affect the asymptotic rate, it would be nice to provide insight into how this difference might play out in practice.

Author response:
Answer: For $\beta_n=o(\gamma_n)$ in case (i), the impact of the SRRW iterates on the SA iterates is primarily reflected in the correlation terms $\mathbf{J}\_{12}(\alpha)$ and $\mathbf{J}\_{22}(\alpha)$, as detailed in equation (55) and Lemma D.3 in Appendix D.2.2. These terms play a pivotal role in shaping the matrix $\mathbf{U}\_{\theta}(\alpha)$ as follows:

$$\mathbf{U}\_{\theta}(\alpha)=\mathbf{U}\_{22}-\mathbf{U}_{21} (\mathbf{J}\_{12}(\alpha) \mathbf{J}\_{22}(\alpha)^{-1})^T-\mathbf{J}\_{12}(\alpha)\mathbf{J}\_{22}(\alpha)^{-1}\mathbf{U}\_{12}+\mathbf{J}\_{12}(\alpha) \mathbf{J}\_{22}(\alpha)^{-1}\mathbf{U}\_{11}(\mathbf{J}\_{12}(\alpha) \mathbf{J}\_{22}(\alpha)^{-1})^T,$$

where $\mathbf{U}\_{ij}$ is defined in (9) for $i,j\in\{1,2\}$. Through algebraic computations, we derive the $O(1/\alpha^2)$ rate for matrix $\mathbf{U}\_{\theta}(\alpha)$ and in turn for the asymptotic covariance $\mathbf{V}\_{\theta}^{(1)}(\alpha)$.

Regarding the difference in step-size exponents ($a$ and $b$ in $\beta_n=n^{-b}$ and $\gamma_n=n^{-a}$), our main CLT results in Theorem 3.3 indicate that for all three cases, the asymptotic covariances remain the same irrespective of the variations in $a$ and $b$. The key change lies in the asymptotic rates $\beta_n^{1/2}$ and $\gamma_n^{1/2}$ of the iterates. In practice, we indeed observe that for a fixed value $b$, the convergence speed in the initial training phase is affected by the different value of $a$. We take $b=0.9$ as an example. Choosing $a$ close to $0.5$ results in the slowest convergence since the SRRW iterates $\mathbf{x}_n$ drastically vary in the initial period, potentially deviating from the target distribution $\boldsymbol{\mu}$ and introducing huge bias to the SA iterates $\theta_n$. As $a$ increases to $0.8$, we note the fastest convergence in case (i) in our simulation setup. However, further increases in $a$ (beyond $0.8$) gradually slow down the convergence since such choice of $a$ would get closer to cases (ii) and (iii). This empirical observation suggests that in practice, as $a$ increases in the range $(0.5,b]$, the convergence becomes faster but then decelerates beyond a certain point close to $b$ in the initial period. This necessitates careful fine-tuning of the step size in case (i) to achieve optimal performance.