SimpleGPT: Improving GPT via A Simple Normalization Strategy

TL;DR

This paper introduces SimpleNorm, a normalization strategy that stabilizes activation scales and reduces the Hessian spectral norm, enabling larger learning rates and improved training stability for large GPT models.

Contribution

We propose SimpleNorm, a simple normalization method that improves GPT training by stabilizing activations and allowing larger learning rates, backed by theoretical analysis and extensive experiments.

Findings

SimpleNorm reduces the Hessian spectral norm.

Models trained with SimpleNorm tolerate 3-10x larger learning rates.

SimpleGPT outperforms baselines in training stability and performance.

Abstract

In this work, we revisit Transformer optimization through the lens of second-order geometry and establish a direct connection between architectural design, activation scale, the Hessian matrix, and the maximum tolerable learning rate. We introduce a simple normalization strategy, termed SimpleNorm, which stabilizes intermediate activation scales by construction. Then, by analyzing the Hessian of the loss with respect to network activations, we theoretically show that SimpleNorm significantly reduces the spectral norm of the Hessian, thereby permitting larger stable learning rates. We validate our theoretical findings through extensive experiments on large GPT models at parameter scales 1B, 1.4B, 7B and 8B. Empirically, SimpleGPT, our SimpleNorm-based network, tolerates learning rates 3$\times$-10$\times$ larger than standard convention, consistently demonstrates strong optimization stability, and achieves substantially better performance than well-established baselines. Specifically, when training 7B-scale models for 60K steps, SimpleGPT achieves a training loss that is 0.08 lower than that of LLaMA2 with QKNorm, reducing the loss from 2.290 to 2.208. Our source code will be released at https://github.com/Ocram7/SimpleGPT.