英伟达nGPT重塑Transformer，AI训练速度暴增20倍！文本越长，加速越快

768 0 0

文章摘要

英伟达团队提出了一种新的神经网络架构——归一化Transformer（nGPT），这一架构在超球面上进行表示学习，能够显著提升大型语言模型（LLM）的训练速度，同时保持原有精度。nGPT通过将所有向量归一化为单位范数，使得输入的token在超球面表面上移动，每一层都通过位移来贡献最终的输出预测。实验结果表明，nGPT在不同上下文长度下，训练速度提高了4到20倍，且上下文越长，训练速度越快。

nGPT的主要贡献包括在超球面上优化网络参数，将归一化Transformer作为超球面上的可变度量优化器，以及实现更快的收敛。与传统的Transformer相比，nGPT在训练过程中进行了多项改进，包括移除归一化层、对所有矩阵进行归一化处理、修改隐藏层参数更新方程、调整注意力机制中的softmax缩放因子、对MLP块的中间状态进行重新缩放、对logits进行重新缩放，以及移除权重衰减和学习率预热步骤。

在OpenWebText数据集上的训练和评估中，nGPT显示出了显著的训练加速，以及在下游任务中的性能提升。此外，nGPT的嵌入范数保持固定，而GPT的嵌入范数则表现出明显的变化。nGPT的注意力矩阵和MLP矩阵的条件数较低，表明其在处理信息时更为稳定和有效。这些特性使得nGPT在训练大型语言模型时具有潜在的优势，为实现更高效、更智能的AI模型提供了新的可能性。