MoLS: Calibrating Adam via SNR to Solve the Gradient Noise Imbalance in LLMs

TL;DR

Large Language Models (LLMs) are modular, yet we often train them with a "one-size-fits-all" learning rate. MoLS (Module-wise Learning Rate Scaling via SNR) reveals that modules like Embeddings and Output Heads suffer from a "noise-damping deficit" in Adam, where high noise suppresses useful updates. By automatically scaling learning rates based on Signal-to-Noise Ratio (SNR) during warmup, MoLS achieves up to 1.4x faster convergence and significantly lower perplexity without manual tuning.

The "Noise-Damping Deficit": Why Adam Isn't Enough

While Adam provides element-wise adaptivity, it doesn't understand the structural hierarchy of an LLM. Within a Transformer, different modules exhibit vastly different gradient statistics.

The authors identify a critical phenomenon: Effective Signal Step ($D_m$) imbalance. In the Low SNR Regime (e.g., Embedding layers), the stochastic noise $\sigma$ overwhelms the true signal $\mu$. Because Adam's denominator contains the second moment (signal + noise), high noise results in a massive denominator that "damps" or suppresses the update. Consequently, while Attention modules are learning rapidly, the Embedding and Head layers are "left behind."

Methodology: Balancing the Signal

The core insight of MoLS is using the Relative Signal-wise Equilibrium principle. If a module has 1/10th the SNR of a "healthy" reference module (like the Value/Output projection), its learning rate should be boosted to compensate.

1. One-shot Calibration

Instead of expensive dynamic tracking, MoLS performs a one-shot SNR estimation during the first 1% of training (roughly 32 iterations). It groups parameters into functional blocks:

• Emb (Embedding)
• QK (Query/Key)
• VO (Value/Output)
• MLP (Feed-forward)
• Head (Output Head)

2. The Rescaling Factor

The scaling factor ${\alpha }_m$ follows an inverse square root relationship with SNR: ${\alpha }_m=\sqrt{\frac{S_{base}}{S_m}}$ This ensures that noise-dominated modules get a square-root-proportional boost to their effective learning rate.

Figure 1: Observed SNR disparities across functional modules during pre-training. Note the significantly lower SNR in Embedding and Head layers.

Experimental Results: SOTA Efficiency

MoLS was tested on Llama models ranging from 60M to 7B parameters.

• Perplexity Gains: On Llama-7B, MoLS reduced final PPL by 1.69 compared to standard 8-bit Adam.
• Speed & Overhead: The SNR estimation adds less than 2% to total training time but results in a 1.2x-1.4x wall-clock speedup due to faster convergence.
• Compatibility: It acts as a plug-and-play wrapper. When combined with Adam-mini, it achieved a 1.53 PPL reduction on a 1.3B model.

Table 1: Comparison of MoLS against standard and memory-efficient optimizers. MoLS consistently hits lower perplexity targets.

Critical Analysis: Is Static Scaling Enough?

The authors address a key concern: Does SNR change too much during training? Their empirical analysis shows that while absolute SNR values drop as the model converges (moving into flatter regions of the loss landscape), the relative ordering of modules remains remarkably stable. This justifies the "one-shot" approach, which keeps the method lightweight and avoids the instability of frequent learning rate jumps.

Limitations:

• The method relies on predefined module boundaries, which might need adjustment for novel architectures like Mixture-of-Experts (MoE).
• It still requires a reference module choice (though VO is shown to be a robust default).

Conclusion

MoLS provides a statistically grounded answer to a problem usually solved by "dark magic" (manual hyperparameter tuning). By revealing the gradient noise imbalance between LLM modules, it allows researchers to treat LLM optimization as a heterogeneous rather than a homogeneous task. For any practitioner training LLMs from scratch or performing extensive SFT, MoLS offers a "free lunch" in terms of convergence speed and final model quality.