Softsign: Smooth Sign in Your Optimizer For Better Parameter Heterogeneity Handling

2026-05-29

AI Summary

This work addresses the convergence difficulties and inadequate handling of parameter heterogeneity inherent in sign-based optimizers, which stem from their fixed-magnitude updates. To overcome these limitations, the authors propose SoftSignum and SoftMuon optimizers that employ a temperature-controlled soft sign transformation, enabling a smooth transition from hard sign updates to gradient-magnitude-sensitive updates. By integrating adaptive quantile-based temperature scheduling, matrix manifold extensions, and a geometric relaxation framework grounded in strongly convex regularization and Fenchel conjugacy, the study establishes the first convergence guarantees for sign-based optimization under non-convex stochastic settings. Experimental results demonstrate that the proposed methods significantly outperform conventional sign optimizers and AdamW in large language model pretraining, achieving superior convergence stability and final model performance.

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.

Problem

• sign-based optimization
• parameter heterogeneity
• update magnitude
• convergence
• gradient magnitude

Innovation

• SoftSignum
• soft-sign transformation
• parameter heterogeneity
• adaptive temperature scheduling
• geometry-relaxation framework