Softsign: Smooth Sign in Your Optimizer For Better Parameter Heterogeneity Handling
This work provides an incremental improvement for deep learning practitioners using sign-based optimizers, addressing their terminal convergence issues.
This paper addresses the issue of fixed-magnitude updates in sign-based optimizers, which can hinder terminal convergence due to parameter heterogeneity. The authors propose SoftSignum, a smooth relaxation that allows for a parameter-wise transition from sign-like to magnitude-sensitive updates, and demonstrate consistent improvements over hard sign-based optimizers and AdamW in various deep learning tasks, including LLM pretraining.
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.