HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization

2025-03-06

AI Summary

In deep Transformer training, the choice of normalization placement—Pre-Norm ensures stability but limits expressivity, while Post-Norm offers superior performance yet suffers from training divergence. To resolve this trade-off, we propose HybridNorm: a heterogeneous intra-block normalization scheme wherein LayerNorm is applied separately to the query, key, and value projections (QKV-Norm) within the attention branch, while Post-Norm is retained in the feed-forward network (FFN) branch. This design achieves the first fine-grained decoupling of normalization paths within a single Transformer block, preserving identity-mapping stability during optimization while enhancing representational capacity. Extensive experiments across multiple large language model (LLM) benchmarks demonstrate that HybridNorm significantly improves training stability, accelerates convergence, and outperforms both Pre-Norm and Post-Norm baselines—achieving state-of-the-art (SOTA) performance. The implementation is publicly available.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Diffusion Models for Vision

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Transformers have become the de facto architecture for a wide range of machine learning tasks, particularly in large language models (LLMs). Despite their remarkable performance, challenges remain in training deep transformer networks, especially regarding the location of layer normalization. While Pre-Norm structures facilitate easier training due to their more prominent identity path, they often yield suboptimal performance compared to Post-Norm. In this paper, we propose $ extbf{HybridNorm}$, a straightforward yet effective hybrid normalization strategy that integrates the advantages of both Pre-Norm and Post-Norm approaches. Specifically, HybridNorm employs QKV normalization within the attention mechanism and Post-Norm in the feed-forward network (FFN) of each transformer block. This design not only stabilizes training but also enhances performance, particularly in the context of LLMs. Comprehensive experiments in both dense and sparse architectures show that HybridNorm consistently outperforms both Pre-Norm and Post-Norm approaches, achieving state-of-the-art results across various benchmarks. These findings highlight the potential of HybridNorm as a more stable and effective technique for improving the training and performance of deep transformer models. %Code will be made publicly available. Code is available at https://github.com/BryceZhuo/HybridNorm.

Problem

• Addresses instability in deep transformer training
• Improves performance of large language models
• Combines Pre-Norm and Post-Norm normalization strategies

Innovation

• HybridNorm integrates Pre-Norm and Post-Norm advantages
• QKV normalization in attention mechanism
• Post-Norm in feed-forward network (FFN)