Training nGPT

2026-08-02

AI Summary

This work investigates efficient training of large-scale Mixture-of-Experts (MoE) models with hyperspherical representation learning. By constraining both parameters and activation vectors onto the unit hypersphere, it systematically integrates hyperspherical representation learning into a hybrid Mamba-2 and Transformer architecture for the first time. The study introduces several training optimization strategies, including logit gradient preconditioning, logarithmic learning rate decay, a GatedAdamW optimizer, angular update control, and an optional exploration mechanism. Evaluated on a 14B-parameter model, the proposed approach achieves validation loss comparable to standard AdamW training using only about half the number of training tokens, substantially reducing computational cost while demonstrating superior scalability and training stability.

Abstract

The normalized Transformer (nGPT) realizes hyperspherical representation learning by constraining model parameter vectors and activation vectors to the unit hypersphere. In this paper, we describe a practical training recipe for nGPT and evaluate it on modern hybrid Mamba-2--Transformer Mixture-of-Experts (MoE) models. The recipe introduces Logit Gradient Preconditioning, Logarithmic Learning Rate Decay, GatedAdamW, angular update control, and optional exploration mechanisms. Compared with an unnormalized model of the same hybrid MoE architecture trained with AdamW, the 14B-total-parameter nGPT model reaches the same validation loss using approximately half as many training tokens. The recipe scales across the models considered, which contain up to 14B total parameters.

Problem

• normalized Transformer
• hyperspherical representation learning
• Mixture-of-Experts
• large-scale model training
• unit hypersphere constraint

Innovation

• normalized Transformer
• hyperspherical representation learning
• Logit Gradient Preconditioning
• GatedAdamW
• Mixture-of-Experts