Robust and Fast Training via Per-Sample Clipping

Technical Summary: Robust and Fast Training via Per-Sample Clipping

Problem Statement
The paper addresses the stochastic optimization problem $\min_{x \in \mathbb{R}^d} F(x)$, where $F(x) = \mathbb{E}_{\xi \sim \mathcal{D}}[f(x, \xi)]$. While Stochastic Gradient Descent (SGD) is the standard solver, its convergence guarantees typically rely on the assumption that gradient noise possesses bounded variance (light-tailed distributions). However, recent empirical and theoretical work indicates that in many machine learning applications, gradient noise may be heavy-tailed, possessing finite $p$-th moments for $p \in (1, 2]$ but lacking bounded variance. Under such conditions, standard SGD can perform poorly. Existing robust alternatives, such as standard gradient clipping (Clip-SGD) and Normalized SGD, either fail to achieve optimal sample complexity lower bounds or require specific, often impractical, parameter tuning (e.g., dependence on the initialization gap $\Delta_1$ or specific logarithmic factors of the failure probability $\delta$).

Methodology
The authors propose Per-Sample Clipped SGD (PS-Clip-SGD), a robust gradient estimator that applies clipping to individual samples within a mini-batch before aggregation, rather than clipping the aggregated batch gradient.

For a batch of size $n$ with samples $\xi^{(1)}_t, \dots, \xi^{(n)}_t$, the gradient estimator $G(x_t, \alpha, \beta)$ is defined as:
$$G(x_t, \alpha, \beta) = \frac{1}{n} \sum_{k=1}^n \gamma^{(k)}_t \nabla f(x_t, \xi^{(k)}_t)$$
where the per-sample clipping factor is $\gamma^{(k)}_t := \min\left\{1, \frac{\alpha k^{1/\beta}}{|\nabla f(x_t, \xi^{(k)}_t)|}\right\}$. The parameters $\alpha$ and $\beta$ control the clipping threshold, with $\alpha$ related to the moment bound $\sigma$ and $\beta$ related to the tail index $p$. The update rule follows standard SGD: $x_{t+1} = x_t - \eta_t G(x_t, \alpha, \beta)$.

Key Contributions

1. Optimal In-Expectation Convergence: The authors prove that PS-Clip-SGD achieves optimal sample complexity for non-convex optimization under heavy-tailed noise. Specifically, the method achieves a convergence rate of $O(1/T)$ for the average squared gradient norm, matching the lower bounds established by Zhang et al. (2020). This improves upon existing results for Clip-SGD and Normalized SGD, particularly because the bound holds for the sum of squared gradient norms (a stronger criterion than the sum of norms used in some prior works) without requiring knowledge of the initialization gap $\Delta_1$ for the step size.
2. High-Probability Convergence: The paper establishes high-probability convergence guarantees that match the in-expectation rates up to polylogarithmic factors in the failure probability $\delta$ and the number of steps $T$. Notably, the dependence on $\log(1/\delta)$ is improved compared to state-of-the-art Clip-SGD bounds, and the method does not require step sizes or clipping thresholds to be explicitly tuned as functions of $\delta$.
3. Gradient Accumulation Insight: The authors challenge the common practice of applying gradient clipping only after all accumulation steps in gradient accumulation. They propose and empirically validate that applying clipping at the mini-batch level (Mini-Batch Clipping) during accumulation can improve training performance with virtually no additional computational cost, contrasting with the standard consensus that such intermediate clipping introduces bias.

Results

• Theoretical: Theoretical analysis demonstrates that PS-Clip-SGD provides strong convergence guarantees under Assumptions of $L$-smoothness and finite $p$-th moments. The derived bounds are shown to be tighter than those for Normalized SGD and Clip-SGD in terms of dependence on $\delta$ and the strength of the convergence measure (squared norms vs. norms).
• Quadratic Optimization: In experiments minimizing a quadratic function with Pareto-distributed noise, PS-Clip-SGD significantly outperformed both Clip-SGD and Normalized SGD across various noise regimes ($p \in \{1.2, 1.5, 1.8\}$), even with untuned hyperparameters. High-probability experiments confirmed the theoretical sublinear dependence on $\log(1/\delta)$.
• Image Classification (AlexNet on CIFAR-100): When training AlexNet with momentum, PS-Clip-SGD achieved higher validation accuracy (64.84%) compared to vanilla SGD with momentum (58.90%) and Clip-SGD with momentum (58.94%). This improvement persisted even after accounting for the $\approx 1.3\times$ computational overhead introduced by per-sample gradient norm calculations. The authors attribute this to the effective control of a small number of outlier gradients with very large norms, which standard batch-level clipping fails to detect.
• Language Modeling (GPT-2): In a gradient accumulation setting for training a 124M parameter GPT-2 model, Mini-Batch Clipping (clipping per accumulation step) outperformed standard Clip-SGD (clipping only after full accumulation) across multiple learning rates, despite the prevailing view that intermediate clipping is detrimental.

Significance and Claims
The paper claims to be the first work to propose per-sample clipping specifically as a method for improving training performance in the presence of heavy-tailed noise, distinct from its established use in differential privacy. The significance lies in providing a method that:

1. Achieves optimal theoretical convergence rates for non-convex heavy-tailed problems.
2. Offers robustness to hyperparameter choices, performing well even with unrefined settings.
3. Demonstrates that per-sample clipping can yield practical performance gains in deep learning tasks (like AlexNet) that outweigh its computational costs.
4. Suggests a paradigm shift in gradient accumulation strategies, showing that intermediate clipping can be beneficial rather than harmful.

The authors acknowledge limitations, noting that their theoretical analysis requires knowledge of the noise moment bound $\sigma$ and tail index $p$ to define optimal parameters, and that the assumption of bounded $p$-th moments is stronger than bounded centered $p$-th moments. They also note that while per-sample clipping is effective for smaller models like AlexNet, the computational overhead may render it impractical for very large architectures without further optimization, though mini-batch clipping offers a cost-effective alternative in accumulation scenarios.