Technical Summary: RanSOM

1. Problem Statement

The paper addresses a fundamental limitation in stochastic optimization for deep learning: curvature-induced bias in standard momentum methods (e.g., Polyak's Heavy Ball) when applied to non-convex landscapes.

In stochastic settings, the momentum vector $m_t$ aggregates gradients from previous iterates. As the optimizer moves from $x_{t-1}$ to $x_t$, the curvature of the loss landscape causes the gradient direction to change. Consequently, $m_{t-1}$ becomes a "stale" estimator of the current gradient $\nabla f(x_t)$. This bias, proportional to the integral of the Hessian along the update path, prevents standard momentum methods from achieving the optimal $O(\epsilon^{-3})$ sample complexity required for non-convex optimization, often limiting them to suboptimal $O(\epsilon^{-4})$ rates.

Existing correction strategies face a "bias-variance bottleneck":

• Difference-based estimators (e.g., STORM): Require individual sample smoothness (often violated by ReLU activations) or suffer from high variance.
• Deterministic Second-Order Momentum (SOM): Requires strong global Lipschitz assumptions on the Hessian (bounded third derivatives).
• Stochastic SOM (SOM-Unif): Relaxes smoothness assumptions but incurs prohibitive computational overhead by requiring auxiliary "look-ahead" queries (3 forward/backward passes per iteration).

2. Methodology: The RanSOM Framework

The authors propose RanSOM (Randomized Second-Order Momentum), a unified framework that eliminates curvature bias without additional assumptions or computational overhead. The core innovation is replacing deterministic step sizes with randomized step sizes drawn from specific distributions.

2.1 Randomized Integration by Parts

Instead of a fixed step $\eta_t$, the step size $s_t$ is treated as a random variable with mean $\mathbb{E}[s_t] = \eta_t$. By leveraging Stein-type identities (integration-by-parts formulas), the method relates the finite difference of gradients (the bias) to a single Hessian-Vector Product (HVP) evaluated at the random endpoint.

The framework utilizes two specific identities:

1. Exponential Identity (Unconstrained): If $s \sim \text{Exp}(\lambda)$, then:
   $$\mathbb{E}[g(s) - g(0)] = \frac{1}{\lambda} \mathbb{E}[g'(s)]$$
   This allows the bias $\nabla f(x_{t+1}) - \nabla f(x_t)$ to be estimated by scaling the HVP at $x_{t+1}$ by $\eta_t$.
2. Beta Identity (Constrained): If $s \sim \text{Beta}(1, K)$, then:
   $$\mathbb{E}[g(s) - g(0)] = \mathbb{E}\left[ \frac{1-s}{K} g'(s) \right]$$
   This accounts for the compact support of the Beta distribution, ensuring feasibility in constrained optimization.

2.2 Algorithmic Instantiations

The framework is instantiated in two algorithms:

• RanSOM-E (Unconstrained): Uses Exponential-distributed steps. It integrates with a Linear Minimization Oracle (LMO) to unify various geometric updates, recovering Normalized SGD ($L_2$-ball), SignSGD ($L_\infty$-ball), and Muon-style spectral updates.
• RanSOM-B (Constrained): Uses Beta-distributed steps on $[0, 1]$. This ensures that the update $x_{t+1} = (1-s_t)x_t + s_t v_t$ (where $v_t$ is the LMO direction) remains strictly within the convex feasible set $C$, avoiding projection operations.

2.3 Computational Efficiency

A critical advantage of RanSOM is its computational cost. The correction term requires an HVP at $x_{t+1}$. In modern automatic differentiation frameworks (e.g., PyTorch, JAX), this HVP can be computed via Pearlmutter's trick in a single combined backward pass alongside the standard gradient computation.

• Cost: $2(F_c + B_c)$ (two forward/backward passes equivalent).
• Comparison: This matches the cost of standard variance-reduced methods like STORM and is significantly cheaper than SOM-Unif, which requires a dedicated auxiliary query ($3(F_c + B_c)$).

3. Key Contributions

1. Unified Framework: A method that eliminates curvature-induced bias using randomized step sizes and Stein identities, applicable to both unconstrained and constrained optimization.
2. Optimal Convergence Rates: Theoretical analysis proves RanSOM achieves the optimal $O(\epsilon^{-3})$ sample complexity under standard bounded noise assumptions.
3. Robustness to Heavy Tails: The method extends to challenging settings with heavy-tailed gradient noise ($p \in (1, 2]$) and generalized $(L_0, L_1)$-smoothness (where the Hessian norm grows with the gradient norm), without requiring the Hessian to be Lipschitz continuous.
4. No Auxiliary Overhead: Unlike prior second-order corrections, RanSOM requires no auxiliary "look-ahead" points, maintaining the standard query complexity of $2(F_c + B_c)$.
5. Clipping-Free: The method achieves optimal rates without explicit gradient clipping, relying instead on the intrinsic normalization provided by the LMO updates and the statistical properties of the randomized steps.

4. Theoretical Results

The paper establishes convergence for non-convex optimization under the following conditions:

• Assumptions: The objective is bounded below and satisfies pointwise $(L_0, L_1)$-smoothness. Stochastic gradients and HVPs are unbiased with affine heavy-tailed noise (moments $p, q \in (1, 2]$).
• Unconstrained (RanSOM-E): Converges to a stationary point with rate:
  $$\frac{1}{T} \sum_{t=0}^{T-1} \mathbb{E}[\|\nabla f(x_t)\|_*] \leq O\left( T^{-\frac{q(p-1)}{2q(p-1)+p}} \right)$$
  For the standard case ($p=q=2$), this recovers the optimal $O(T^{-1/3})$ rate.
• Constrained (RanSOM-B): Converges in terms of the Frank-Wolfe gap $G(x)$ with the same optimal rate $O(T^{-1/3})$ for $p=q=2$.

The analysis highlights an asymmetric scaling of the momentum parameter $\beta$: smaller $\beta$ helps average out high-variance gradients, while larger $\beta$ makes the curvature correction more impactful. The optimal $\beta$ balances these effects.

5. Experimental Results

The authors evaluate RanSOM on three non-convex tasks:

1. Binary Classification (SPLICE): RanSOM-E (Muon variant) achieved the highest test accuracy (84.90%) with significantly lower variance compared to STORM (73.35% ± 10.77%).
2. Sequence Classification (MNIST1D): RanSOM-E (Muon) reached 91.90% accuracy, outperforming STORM (88.83%) and classic SOM, while maintaining high stability (std dev ±0.36 vs. ±1.38 for STORM).
3. Constrained Matrix Completion (MovieLens): RanSOM-B converged to a lower final RMSE than Stochastic Frank-Wolfe baselines (SFW-Polyak, SFW-SOM), demonstrating effectiveness in navigating complex constraint geometries.

6. Significance and Claims

The paper claims that RanSOM resolves the "dilemma" of second-order momentum corrections: it achieves the optimal convergence rate of $O(\epsilon^{-3})$ without the restrictive assumptions of Lipschitz Hessians or individual sample smoothness, and without the computational overhead of auxiliary sampling.

By treating step sizes as random variables, RanSOM enables robust second-order bias correction that is:

• Assumption-light: Handles non-smooth individual losses (e.g., ReLU networks) and heavy-tailed noise.
• Computationally efficient: Matches the cost of standard SGD with momentum.
• Geometrically flexible: Naturally generalizes to normalized optimizers and constrained settings via LMOs.

The authors position RanSOM as a unified solution for modern non-convex optimization, particularly in deep learning architectures where standard smoothness assumptions fail and computational efficiency is paramount.