Explain the directional first-moment bias and its relationship to learning-rate decay

Determine the cause of the directional first-moment bias observed under NF4 round-to-nearest quantization for the LM-head optimizer state and characterize its possible relationship to learning-rate decay.

Background

The paper finds that deterministic NF4 round-to-nearest quantization of the LM-head first moment can produce a persistent inward radial bias during late-stage GPT-style 2.7B pretraining, accompanied by gradient growth and a sharp validation-loss regression. Switching the LM-head first moment to stochastic rounding alleviates the observed instability, but the mechanism generating the directional bias is not identified. The authors also leave unresolved whether the bias is connected to the learning-rate decay schedule, making the causal relationship between quantization bias and late-training dynamics an explicit open problem.

References

Targeted LM-head SR alleviates the observed late-training regression, but the cause of the directional first-moment bias and its possible relationship to learning-rate decay remain unresolved.

— Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization  (2610.12444 - Li et al., 8 Oct 2026) in Conclusion, subsection “Limitations and future work”