Training nGPT

LLM PTSAD MoE ML Mamba/SSM Transformer
608
2026年08月02日
归一化Transformer（nGPT）通过将模型参数向量与激活向量约束在单位超球面上，实现超球面表征学习。本文提出了一套适用于nGPT的实用训练方案，并在当前主流的混合型Mamba-2–Transformer专家混合（MoE）模型上对其进行了评估。该方案包含逻辑回归输出梯度预处理（Logit Gradient Preconditioning）、对数形式的学习率衰减（Logarithmic Learning Rate Decay）、门控AdamW优化器（GatedAdamW）、角度更新调控机制，以及可选的探索机制。相较于采用AdamW优化器训练、具有相同混合MoE架构但未施加归一化的基线模型，总参数量为140亿的nGPT模型仅需约一半的训练token数，即可达到同等的验证损失。该训练方案在所考察的各类模型上均展现出良好的可扩展性，涵盖参数量最高达140亿的模型。
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.
许愿