Technical Summary: SignMuon: Communication-Efficient Distributed Muon Optimization

1. Problem Statement

The distributed training of large neural networks, particularly Large Language Models (LLMs), faces a critical bottleneck in communication bandwidth and latency. In synchronous data-parallel training, workers must constantly synchronize large volumes of data (gradients, updates, or optimizer states). Two primary limitations exist in current approaches:

1. Full-Precision Communication: Standard optimizers transmit full-precision (e.g., float32) gradients or updates, which is expensive in terms of bandwidth.
2. Coordinatewise Ignorance of Matrix Structure: Many popular communication-efficient methods, such as signSGD, treat weight tensors as one-dimensional vectors. They ignore the inherent matrix structure of parameters (e.g., MLP weights, attention projections), potentially missing optimization benefits derived from matrix geometry.

Existing solutions often address only one of these issues: sign-based methods achieve extreme compression but lack matrix awareness, while matrix-aware optimizers like Muon improve convergence geometry but retain the high communication cost of full-precision synchronization.

2. Methodology: Sign-Muon

The authors propose Sign-Muon, a hybrid optimizer that integrates the communication efficiency of sign-based methods with the matrix-aware geometry of Muon.

Core Mechanism

Sign-Muon operates by combining two distinct components:

1. Muon-style Polar Directions: Instead of using raw gradients, each worker computes a momentum matrix and derives an update direction via polar decomposition. This is approximated locally using the Newton–Schulz (NS) iteration, which orthogonalizes the momentum matrix. This step ensures the update direction respects the spectral-norm geometry of the weight tensors.
2. 1-Bit Sign Aggregation: Rather than transmitting the full-precision polar direction, each worker transmits only the entrywise signs ($\{-1, +1\}$) of the computed direction.
3. Majority Vote Aggregation: In the distributed setting, workers aggregate these sign matrices via a majority vote. This is implemented efficiently as an integer SUM all-reduce (summing the signs across workers) followed by a local thresholding operation ($\text{sign}(\sum S^{(m)})$).

Algorithmic Flow (Distributed)

1. Local Computation: Each worker $m$ computes the stochastic gradient $G^{(m)}_t$, updates local momentum $M^{(m)}_{t+1}$, and computes the polar factor $U^{(m)}_t$ via Newton–Schulz iteration.
2. Sign Extraction: The worker extracts the entrywise sign $S^{(m)}_t = \text{sign}(U^{(m)}_t)$.
3. Communication: Workers perform a single collective communication step:
   • All-Reduce Variant: An integer SUM all-reduce aggregates the signs.
   • All-Gather Variant (1-bit): Signs are packed into bits (1 bit per entry), and an All-Gather disseminates the packed buffers.
4. Local Update: Each worker computes the majority vote $\bar{S}_t$ locally and updates parameters: $W_{t+1} = W_t - \eta_t \bar{S}_t$.
5. Optional Local Orthogonalization: Workers may optionally apply a local polar step to the aggregated sign matrix to further enforce orthogonality without incurring extra communication.

Communication Efficiency

• Payload Size: The method reduces communication to 1 bit per parameter (or 8 bits if using int8 for implementation simplicity).
• Reduction Factor: Compared to float32 all-reduce, Sign-Muon achieves a 32× bandwidth reduction (with bit-packing) or 4× reduction (with int8).
• Local Computation: All spectral-norm normalization and orthogonalization (Newton–Schulz) are performed locally, incurring zero additional network communication cost.

3. Theoretical Contributions

The paper provides a convergence analysis under spectral-norm smoothness and bounded-variance stochastic gradients.

• Convergence Rate: For an $\ell_1$-based stationarity measure, Sign-Muon achieves an $O(1/\sqrt{T})$ nonconvex convergence rate.
• Noise Reduction: Under the assumption of unimodal symmetric noise, the majority vote mechanism across $M$ workers reduces the stochastic noise term by a factor of $1/\sqrt{M}$, matching the theoretical benefits of distributed signSGD.
• Geometry: The analysis utilizes a spectral-norm descent lemma, distinguishing it from standard signSGD analyses that rely on coordinatewise smoothness. The spectral normalization ensures the update direction has a spectral norm of at most 1.

4. Empirical Results

The authors evaluated Sign-Muon on two workloads: CIFAR-10 (image classification with ResNet architectures) and nanoGPT (language modeling).

CIFAR-10 (ResNet-50)

• Accuracy: Across 330 hyperparameter configurations, Sign-Muon achieved the best validation accuracy of 92.15%.
• Distributed Efficiency: A 4-GPU majority-vote variant reached 92.02% accuracy (matching the single-GPU performance) while reducing training time by 37% compared to matched effective batch sizes.
• Scaling: Sign-Muon demonstrated superior weak-scaling behavior compared to SignAdam and SignSGD. On 16 GPUs, it incurred significantly lower training-time overheads (+77.8% vs. +150–163% for baselines in All-Gather settings).

nanoGPT

• Performance: Sign-Muon achieved lower perplexity and better "anytime" performance (running minimum perplexity over wall-clock time) compared to other sign-based baselines.
• Scaling: Favorable weak-scaling was observed up to 16 GPUs.
• Hyperparameters: The method showed robustness to normalization schemes (Frobenius vs. Spectral scaling) and required minimal tuning (e.g., 1 Newton–Schulz iteration often sufficed).

Memory and Overhead

• Memory: Sign-Muon maintains a persistent memory footprint similar to SGD (one momentum buffer), significantly lower than Adam/AdamW which maintain two moment vectors.
• Computation: The local Newton–Schulz iterations add computational overhead per step but are offset by the massive reduction in communication time, particularly in bandwidth-constrained environments.

5. Significance and Claims

The paper claims that Sign-Muon successfully bridges the gap between two previously distinct lines of research:

1. Communication Efficiency: It retains the extreme compression (1-bit) and robustness (majority vote) of sign-based methods.
2. Matrix Awareness: It incorporates the geometric benefits of Muon (polar decomposition) without sacrificing communication efficiency.

Key Claims:

• Bridging the Gap: Sign-Muon is the first optimizer to embed majority-vote sign aggregation directly within Muon's polar-step framework.
• Scalability: It offers a practical solution for large-scale distributed training where bandwidth is the primary bottleneck, reducing communication volume by up to 32× compared to float32.
• Theoretical Soundness: The method provides rigorous convergence guarantees under spectral-norm smoothness, a more appropriate metric for matrix-structured updates than coordinatewise smoothness.
• Practicality: Empirical results on standard benchmarks (CIFAR-10, nanoGPT) demonstrate that Sign-Muon not only matches but often exceeds the performance of full-precision and other sign-based baselines while significantly reducing training time in distributed settings.

The authors emphasize that the method is particularly effective when network bandwidth is the bottleneck, as the local computation of orthogonalization is "free" in terms of communication cost.