SimpleGPT: Improving GPT via A Simple Normalization Strategy

TL;DR

This work addresses optimization stability in large Transformer models by adopting a second-order perspective that ties architectural choices to Hessian geometry. It introduces SimpleNorm, a normalization strategy placed immediately after linear mappings, and shows it reduces the Hessian spectral norm, enabling learning rates up to 3x–10x larger than strong baselines. Empirically, SimpleGPT improves stability and reduces training loss across nanoGPT, Llama2, and Llama3 backbones from 1B to 8B parameters, including a notable 0.08 loss improvement at 7B after 60K steps, with minimal runtime overhead. The approach is underpinned by Gauss-Newton dominance and weight-scale invariance, offering a principled path toward scalable, stable training of large language models.

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.

Figures (12)

• Figure 1: SimpleGPT vs. GPT. This figure compares the standard GPT block with the proposed SimpleGPT block, highlighting the structural simplifications introduced by SimpleNorm.
• Figure 2: The largest admissible learning rate for Llama2-B, Llama2-1B with QKNorm, and SimpleGPT-1B.
• Figure 3: The training loss curves of Llama2 7B, Llama2 7B with QKNorm and SimpleGPT 7B under 20K, 40K and 60K training steps.
• Figure 4: The training loss curves of Llama2 1B, Llama2 1B with QKNorm and SimpleGPT 1B under 200K training steps.
• Figure 5: Overall comparison across Llama2 1B, Llama2 1B with QKNorm and SimpleGPT 1B under three different learning rates. Adam-mini uses a $2\times 10^{-4}$ learning rate. In SimpleGPT, we enable a $10\times$ learning rate and obtain better performance.

Theorems & Definitions (5)

• Definition 3.1: $\beta$-smoothness
• Theorem 4.1: Gauss--Newton dominance for SimpleNorm
• Theorem 4.2: Linear curvature scales with $\|\boldsymbol{W}\|_2^2$ while SimpleNorm does not
• proof
• proof