SimpleGPT: Improving GPT via A Simple Normalization Strategy

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 reduces the training loss from 2.290 to 2.208 compared with Llama2 with QKNorm. Our code is available at \url{https://github.com/Ocram7/SimpleGPT}.

We present SimpleNorm, a simple but effective change to the Transformer architecture that makes training more reliable. By keeping internal signals within a healthy range throughout the network, SimpleNorm helps models learn smoothly and reduces the sensitivity to training settings. As a result, models can be trained using much larger update steps, which can improve training efficiency and reduce the need for extensive tuning. We explain why this behavior occurs through mathematical analysis and verify it in experiments on language models ranging from 1 billion to 8 billion parameters. Across all tested scales, models using SimpleNorm train more stably and achieve better performance than strong existing baselines. For example, in large-scale 7-billion-parameter training runs, our approach consistently reaches lower training loss than widely used Transformer designs. Our findings show that a small architectural change can make large language model training simpler, more robust, and potentially more cost-effective, helping reduce the resources required to develop advanced AI systems.