Softsign: Smooth Sign in Your Optimizer For Better Parameter Heterogeneity Handling
Abstract
Sign-based and LMO-inspired optimizers have recently attracted substantial attention in deep learning due to their strong performance and low memory footprint. However, their fixed-magnitude updates can hurt terminal convergence: they decouple update mechanisms from gradient magnitudes and fail to account for parameter heterogeneity, often leading to oscillation rather than convergence. We propose SoftSignum, a smooth relaxation of sign-based optimization that replaces the hard sign map with a temperature-controlled soft-sign transformation, enabling a parameter-wise transition from sign-like updates to magnitude-sensitive SGD-like steps. We complement it with an adaptive quantile-based temperature schedule and extend the same principle to matrix-valued optimizers, obtaining SoftMuon. We also develop a generalized geometry-relaxation framework based on strongly convex regularizers and Fenchel conjugates, proving convergence in stochastic non-convex setting. Experiments on diverse deep learning tasks, including LLM pretraining, show that SoftSignum and SoftMuon consistently improve over their hard sign-based counterparts and standard AdamW.
Lay Summary
Optimizers are tools used to iteratively minimize an artificial intelligence model's learning error during training. Many fast, popular optimizers only look at which direction a learning step should go, forcing the AI model to take a rigid, giant stride every single time. However, because different parts of the AI model learn at different speeds, these uniform giant steps cause the model to overshoot its target. It endlessly paces back and forth, failing to settle on the absolute best final solution. To solve this, we created two new optimizers, SoftSignum and SoftMuon. Think of our tools as giving the model a dynamic stride that naturally shortens for some model parts: early on, it takes large strides to cover ground quickly, but as parts of the AI model approach their destination, their step sizes smoothly shrink to tiny, precise adjustments, while other still have giant steps. Our method combines rapid, early-stage speed with surgical precision. In tests across various AI tasks our optimizers consistently learned faster and reached higher final accuracy than standard industry formulas.