GPAS: Accelerating Convergence of LLM Pretraining via Gradient-Preserving Activation Scaling

TL;DR

GPAS is a technique that scales down activations in Pre-LN Transformers during pretraining, preserving gradients and improving convergence and performance across various large language models.

Contribution

We introduce GPAS, a simple activation scaling method that preserves gradients, enhancing training stability and performance in Pre-LN Transformers and other architectures.

Findings

GPAS consistently improves model performance across sizes from 71M to 1B parameters.

GPAS enhances training stability in Pre-LN, Sandwich-LN, and DeepNorm architectures.

The method is versatile and can be combined with existing training approaches.

Abstract

Modern Large Language Models, such as the LLaMA, Qwen and DeepSeek series, predominantly adopt the Pre-LayerNorm (Pre-LN) Transformer architecture. While being stable during pretraining and scalable to large model sizes, Pre-LN suffers from an exponential growth in activation variance across layers, causing the shortcut to dominate over sub-layer outputs in the residual connection and limiting the learning capacity of deeper layers. To mitigate this issue, we propose Gradient-Preserving Activation Scaling (GPAS), a simple technique that can be used in combination with existing approaches. GPAS works by scaling down the intermediate activations while keeping their gradients unchanged. This leaves information in the activations intact, and avoids the gradient vanishing problem associated with gradient downscaling. Extensive experiments across various model sizes from 71M to 1B show that GPAS achieves consistent performance gains. Beyond enhancing Pre-LN Transformers, GPAS also shows promise in improving alternative architectures such as Sandwich-LN and DeepNorm, demonstrating its versatility and potential for improving training dynamics in a wide range of settings. Our code is available at https://github.com/dandingsky/GPAS.