Efficient Distributed Optimization under Heavy-Tailed Noise

We thank reviewer y2B4 for their review. As the reviewer noted, TailOPT can be readily combined with a wide range of optimization strategies, which greatly enhances its practical applicability. We propose several novel instantiations of TailOPT (such as $Bi^2Clip$) that achieves the strongest empirical performance without the extra memory/compute/communication overhead compared with baselines. Additionally, the TailOPT framework provides theoretical convergence guarantees under heavy-tailed noise with unbounded gradient variance and local updates, and these results have not appeared in prior works.

[No Decoder-Based Transformers?] We clarify that our evaluations already include encoder-only (RoBERTa), encoder-decoder (T5), and decoder-only (GPT-2, DistillGPT) transformers. Results on decoder-based architectures (T5, GPT) are included in both Section 6 and Appendix H. For clarity, we summarize key empirical results below:

• Section 6.1 (Synthetic Experiments, Full results in Appendix G.1): We carefully inject heavy-tailed noise into gradients via corrupted labels, simulating both a controlled heavy-tailed and non-heavy-tailed regime. Results show that TailOPT significantly outperforms Adam, with the performance gap widening under heavy-tailed conditions. Notably, we show that heavy-tailed noise destabilizes conventional training methods, a setting in which TailOPT instantions consistently demonstrate optimal performance.

• Section 6.2 (GLUE Benchmark): Table 10 (expanded Table 1) extensively evaluates ~150 optimizer-dataset pairs. $Bi^2Clip$, which applies adaptivity to both inner/outer optimizers, consistently achieves the best performance. We find that the action of simply switching the inner optimizer from SGD to $BiClip$ yields approximately 10% gain on GLUE ** average ** with identical memory and communication cost, highlighting the importance of TailOPT and heavy-tailed convergence in training modern transformer models.

• Section 6.3 and Appendix H.2 (Machine Translation/Generation): TailOPT variants outperform $Adam^2$, DiLoCo, and FedAdam across nearly all machine translation tasks, while being more resource-efficient. Table 11 (expanded Table 2) confirms optimal performance of $Bi^2Clip$ on generative tasks via T5, and Appendix H.3 explores robustness under partial participation using GPT-2 and DistillGPT. Across all settings, algorithmic instantiations from TailOPT consistently achieve state-of-the-art performance.

[Other Questions] We use $d$ to denote the coordinate dimension, following optimization literature conventions. The variables $u_t$ and $d_t$ represent “up” and “down” clipping thresholds, respectively. We are happy to revise notation to avoid confusion.

Regarding hyperparameter tuning, TailOPT, particularly $Bi^2Clip$, is easy to tune. We use fixed clipping values across all iterations (lines 290-306) and sweep over only three grid values each for $u_t$ and $d_t$ for $Bi^2Clip$ (Table 3, Appendix G.5). For example, for MNLI, we sweep $d_t$ over {1e-7, 5e-5, 1e-4}, and the best performing validation accuracies are 0.850, 0.849, and 0.850, with barely any difference. Despite this minimal tuning effort, $Bi^2Clip$ exceeds state-of-the-art performance. In contrast, other baselines (such as Adam) requires careful tuning of more hyperparameters, that is, the learning rate, $\beta_1$, $\beta_2$, and $\varepsilon$. Moreover, $BiClip$ requires only 33% of the memory used by Adam to store gradient statistics and even avoids the need to store or compute auxiliary gradient statistics at all, leading to significant improvements in memory, compute efficiency, and communication cost.

[Theoretical Results] Assumptions are listed in Sections 3, 5, lines 145-164, 226-234 (L-smoothness, bounded $\alpha$-moment) for the theorems and corollaries mentioned. We also believe the reviewer may have unintentionally interpreted our setting through the lens of the bounded stochastic gradient literature, where the state-of-the-art convergence rate for non-convex objectives is $\mathcal{O}(T^{-1/2})$. However, these results do not generalize to the heavy-tailed regime with local updates (under the unbounded stochastic gradient variance assumption), which is the focus of our work. TailOPT is, to our knowledge, the first general framework to provably achieve this optimal rate under heavy-tailed noise, a setting that is highly relevant to training modern models such as large language models and large-scale transformers. Due to the unbounded variance in this regime, both the theoretical analysis and algorithmic design are significantly more challenging, which TailOPT addresses theoretically with convergence guarantees which are rigorously validated by empirical results. Our extensive evaluations further support these theoretical advantages, demonstrating consistent improvements over state-of-the-art baselines.

We are happy to answer any further questions.

We thank reviewer QUNM for the feedback and for finding our theoretical and empirical results convincing. We fully agree that addressing heavy-tailed gradient variance is a critical challenge, particularly in the context of training large-scale models such as LLMs. As the reviewer notes, the empirical performance of the novel $BiClip$ method (an instantiation of our proposed TailOPT framework) matches or exceeds that of Adam, while avoiding the substantial memory, communication, and compute costs associated with maintaining extra gradient statistics. Moreover, TailOPT is equipped with theoretical guarantees in the heavy-tailed regime with unbounded variance and local updates. We answer the reviewer’s clarification questions below.

[Clarification of Theory Sections]

• In response to the reviewer’s suggestions, we will clarify the “maximal rate at least” terminology to align with the formal statement in Theorem 6. Thank you for this feedback.
• As discussed after Corollary 1, we compare our results with existing convergence rates of other algorithms. We will further expand the discussion comparing TailOPT to existing baselines, such as Avg+SGD, Adam+SGD, and Adagrad+SGD, which are currently known to converge only under bounded gradient noise. Furthermore, we note that the convergence of the advanced $Adam^2$ and DiLoCo algorithms under heavy-tailed noise remain unknown in prior works. Due to space constraints, we moved discussions around convergence of L2 Clip to Appendix C and Appendix D, but we will discuss connections in the main text in the next version.
• We will also move the discussion from the Appendix (lines 1516-1522) regarding the generalization of BiClip to standard L2 Clipping into the main text, and expand it for clarity. We note that for proving the framework-wide convergence for TailOPT, we prove the bounds separately for various instantiations, including $Bi^2Clip$, RMSProp-$BiClip$, etc. That is, the $Bi^2Clip$ theoretical result does recover other TailOPT instantiations such as Adagrad-$BiClip$, as they are based on different inner and outer optimizers.

We thank reviewer 6e8Q for their review. We appreciate the reviewer's acknowledgements about the strengths of our paper, particularly the novelty and efficiency of our proposed algorithm.

[Proof Challenges] The challenges in the analysis lie in heavy-tailed noise with unbounded gradient variance and bias introduced by clipping. The core idea is to decompose model updates into descent directions and clipping bias. We carefully design clipping thresholds so that the bias introduced by thresholding diminishes asymptotically. We will incorporate the reviewer's feedback and further clarify key insights.

[Ablation Studies] In our submission, we already conducted ablation studies on all components of TailOPT and explored most combinations (e.g., Avg-L2Clip vs Avg-BiClip to explore clipping strategies, Adam-SGD vs Avg-SGD to explore effects of outer optimizers, etc) in Section 6. For instance, Table 1 (or Table 10, Appendix H.1) shows the performance of many methods differing by inner/outer optimizers. We see that (a) our $BiClip$ can mitigate the negative effects of heavy-tailed noise while achieving Adam-like performance, and (b) $BiClip$ applied to both inner/outer optimizers ($Bi^2Clip$) outperforms all baselines.

Additionally, we include a summary table comparing communication/memory requirements for gradient statistics of different algorithms. The communication cost includes model parameters (dimension d). In particular, $Bi^2Clip$ achieves best performance, without incurring extra memory/communication cost compared with baselines.

[Practical Implementation] Our goal is not to develop distributed optimization methods that handle varying network latency or heterogeneous hardware. Instead, we focus on developing, analyzing, and evaluating a communication- and memory-efficient framework (TailOPT) that addresses heavy-tailed gradient noise. That said, we do evaluate TailOPT with node failures theoretically (Appendix D) and empirically (Appendix H.3). We see that under partial participation (e.g., Table 12), algorithms using BiClip achieve the top accuracies. Additionally, the lightweight and modular nature of TailOPT (e.g., table above) ensures ease of TailOPT deployment at scale in practical distributed systems. TailOPT is a stateless algorithm---no algorithm states need be maintained across communication rounds, which also makes it suitable for large-scale potentially resource-constrained settings.

[Bounded Deterministic Gradients]
We assume only the deterministic gradient $\nabla F(x)$ is bounded, while the stochastic gradient $\nabla F(x) + \xi$ is unbounded due to heavy-tailed $\xi$. The former is a common assumption in distributed optimization (e.g., [1-2]), where even the stochastic gradient variance is assumed bounded, a much stronger assumption. Additionally, our assumptions are much weaker than prior works which assume bounded variance on the heavy-tailed stochastic gradient (e.g., see lines 110-124) whereas we allow it to be unbounded.

[Other Questions & Related Work]

• The additive form $F_i(x,\xi) = F_i(x) + \langle \xi, x \rangle$ follows from integrating $\nabla F_i(x,\xi) = \nabla F(x) + \xi$ where $\xi$ models heavy-tailed noise (lines 134-140).
• Theorem 1 holds for all $\beta_2 \in [0,1)$, thus extending RMSProp-style results as the reviewer noted.
• We assume a fixed small $\varepsilon$ in the analysis. For readability, we presented only the schedule-dependent terms in the theorem statement. In Theorem 6, setting $\tau = 0$ invalidates convergence by exploding bounds (lines 1681-1694), which is consistent with empirical results highlighting the importance of the adaptivity parameter.
• We will include discussions on the suggested references, which covers Adagrad-type algorithms. [3] assumes bounded variance and almost surely bounded gradients, a much stronger assumption, and [4] studies affine variance which is different from heavy-tailed.

[1] Zaheer et al., Adaptive methods for nonconvex optimization, 2018

[2] Reddi et al., Adaptive Federated Optimization, 2021

[3] Kavis et al., High Probability Bounds for a Class of Non-Convex Algorithms with Adagrad Stepsize, 2022

[4] Faw et al., Beyond Uniform Smoothness: A Stopped Analysis of Adaptive SGD Affine variance, 2023