CoupledNorm: Efficient Normalization via Shared RMS Statistics
Abstract
Normalization layers are treated as computationally inexpensive, yet they still introduce non-negligible latency during autoregressive decoding. We propose CoupledNorm, a simple modification to Pre-Norm Transformer blocks that removes one of the two RMS calculations per layer. CoupledNorm shares the per-token RMS statistics between the Attention and MLP sub-layers while retaining separate learned affine parameters. In GPT-2-scale pretraining, CoupledNorm matches the training loss and yields similar downstream performance. For pretrained 0.6B-8B models, we introduce CoupledNorm post hoc via distillation with small downstream degradation. By fusing the remaining pre-MLP normalization operations into existing kernels, CoupledNorm achieves an end-to-end decoding speedup of up to 2%. These results suggest that shared second-order statistics are sufficient for effective normalization, challenging the need for independent normalization per sub-layer while improving efficiency.