Technical Summary: Singularity-aware Optimization via Randomized Geometric Probing

1. Problem Statement

Deep learning optimization theory traditionally relies on the assumption of Lipschitz-continuous gradients, ensuring stable descent and convergence. However, modern neural architectures systematically violate this assumption through non-smooth components such as ReLU activations, quantization operators (in Quantization-Aware Training, QAT), and sparsity-inducing regularizers.

In these non-smooth regimes, the local geometry is characterized not by a single gradient but by the Clarke subdifferential ($\partial_C f(x)$), a convex set of limiting gradients. When the diameter of this set is large, momentum-based adaptive optimizers like Adam suffer from gradient chattering. This phenomenon involves violent oscillations caused by conflicting directional signals within the subdifferential, leading to poor convergence, failure to settle in sharp low-loss basins, and suboptimal generalization.

Existing remedies are insufficient:

• Proximal methods require explicit proximal mappings, which are often intractable for nested non-smooth functions found in Quantized Neural Networks (QNNs).
• Randomized smoothing techniques typically replace the original gradient with a surrogate, introducing estimation bias and high variance that degrade performance in high-precision training.
• Sharpness-Aware Minimization (SAM) relies on first-order Taylor expansions that are theoretically invalid at non-differentiable points.

2. Methodology: Singularity-aware Adam (S-Adam)

The authors propose S-Adam, a novel optimizer that stabilizes training by dynamically modulating step sizes based on local geometric instability, without smoothing the objective function itself.

2.1 Local Geometric Instability (LGI) Metric

The core innovation is the LGI metric ($\rho_t$), a computationally efficient estimator of the Clarke subdifferential diameter. Instead of computing Hessians or explicit subgradients, S-Adam employs randomized geometric probing:

1. At each iteration $t$, $k$ random directions $\{u_1, \dots, u_k\}$ are sampled from the unit sphere $U(S^{d-1})$.
2. Directional derivatives $D_i$ are computed via finite differences: $D_i = \frac{f(x_t + \delta u_i) - f(x_t)}{\delta}$.
3. The LGI score is defined as the normalized variance of these directional derivatives:
   $$\rho_t(x_t) = \frac{\text{Var}_{u_i \sim \{u\}}[D_i]}{\mathbb{E}[D_i^2] + \epsilon}$$
   Theoretically, this metric approximates the relative curvature of a smoothed surrogate and correlates with the diameter of the Clarke subdifferential. A high $\rho_t$ indicates a region of high geometric instability (a singularity).

2.2 Adaptive Geometric Brake

S-Adam modifies the standard Adam update rule by introducing an exponential damping factor (geometric brake) that scales the learning rate in real-time:
$$\hat{\eta}_t = \eta_t \cdot \exp(-\lambda \rho_t)$$

• High Instability ($\rho_t \gg 0$): The exponential term becomes small, aggressively decelerating updates to prevent oscillation across singularities.
• Smooth Regions ($\rho_t \to 0$): The damping factor approaches 1, allowing S-Adam to degenerate into standard Adam for fast convergence.

This mechanism acts as an implicit proximal method with dynamic regularization, preserving the original gradient direction while adaptively controlling step size.

3. Key Contributions

1. LGI Metric: A variance-based estimator of local geometric instability that requires no Hessian computation and avoids the bias of gradient surrogates.
2. Adaptive Damping Mechanism: A geometric brake $\exp(-\lambda \rho_t)$ that stabilizes training in singular regions while maintaining efficiency in smooth basins.
3. Rigorous Convergence Analysis: Using the framework of differential inclusions, the authors prove that S-Adam converges almost surely to $(\delta, \epsilon)$-Clarke stationary points at the optimal stochastic non-smooth rate of $O(1/\sqrt{T})$.
4. Empirical Robustness: Demonstrated superiority in extreme non-smooth settings, specifically low-bit QAT and high-noise small-batch learning.

4. Experimental Results

The paper evaluates S-Adam against AdamW and Prox-SGD across two primary experimental setups:

4.1 Quantization-Aware Training (QAT)

Evaluated on CIFAR-100 (2-bit), TinyImageNet and Imagewoof2-160 (4-bit), and ImageNet (8-bit).

• Accuracy Gains: S-Adam consistently outperformed baselines. On CIFAR-100, it achieved 18.67% accuracy (vs. 15.94% for AdamW). On TinyImageNet, it reached 23.18% (vs. 18.91% for AdamW).
• Stability: While AdamW and Prox-SGD exhibited severe oscillations or divergence (e.g., Prox-SGD collapsing to <10% accuracy on CIFAR-100), S-Adam maintained stable convergence to lower loss floors.
• Efficiency: The low-cost setting ($k=2$ probes) incurred minimal overhead (e.g., 1.4% increase in wall-clock time on ImageNet) while delivering significant accuracy improvements.

4.2 High-Noise Small-Batch Learning

Evaluated on ResNet-18 with batch sizes as low as $N=2$ to maximize stochastic gradient noise.

• Extreme Noise Resilience: At $N=2$, S-Adam ($k=8$) improved accuracy by +24.69% on Imagewoof2-160 compared to AdamW (75.87% vs. 51.18%).
• Batch Size Scaling: The performance gap was largest in the most unstable regimes (small batch sizes) and narrowed as batch sizes increased and the loss landscape smoothed, confirming the adaptive nature of the geometric brake.
• Learning Rate Sensitivity: S-Adam outperformed AdamW across a wide range of learning rates ($10^{-4}$ to $10^{-3}$), demonstrating robustness to hyperparameter choices.

5. Significance and Claims

The paper positions S-Adam as a bridge between non-smooth optimization theory and practical deep learning scalability. Its primary significance lies in:

• Theoretical Grounding: Providing the first adaptive optimizer with provable convergence guarantees to Clarke stationary points in non-smooth deep learning settings, moving beyond heuristics.
• Practical Applicability: Offering a "drop-in" replacement for Adam that handles nested non-smoothness (e.g., quantization) without requiring intractable proximal mappings or altering the optimization dynamics via gradient surrogates.
• Geometric Insight: Demonstrating that stochastic gradient statistics encode actionable geometric information about local instability, enabling optimizers to distinguish between stochastic noise and topological singularities.

The authors conclude that S-Adam effectively mitigates gradient chattering, enabling stable training in critical applications like low-bit quantization and high-noise regimes where standard optimizers fail.