Deep neural networks often use a technique called LayerNorm to keep numbers stable during training, and where you place it matters. Pre-LayerNorm, placed before a sublayer, is popular but wastes parameters. Post-LayerNorm, placed after, could be more efficient but training tends to fail as the network gets deeper. This paper introduces KiteNorm, a post-LayerNorm method that adds a regularisation term in the training objective to prevent the hidden-state values from growing too large, fighting the root cause of training instability. In experiments with decoder-only Transformers up to 1.3B parameters, KiteNorm trained stably even with large learning rates and outperformed several leading baselines. This matters because it shows post-norm designs can be practical at scale, potentially enabling more parameter-efficient deep networks without the compromises of pre-norm.
Language models can train more stably by limiting internal variance
Trochelmann and colleagues report that controlling hidden-state variance lets an alternative normalisation design outperform strong baselines in decoder-only models up to 1.3 billion parameters.
Industry
Leon A. Trochelmann · Sajad Movahedi · Shiwei Liu · Antonio Orvieto
ELLIS Institute Tübingen · MPI for Intelligent Systems · Tübingen AI Center
Research Digest··2 min read
In one line
KiteNorm stabilises post-norm Transformer training by constraining hidden-state variance, outperforming pre-norm methods.
Read it as
Trochelmann and colleagues revisit a language-model design largely displaced because it becomes difficult to train as depth increases.
The authors studied where to place layer normalisation, an operation that rescales a model’s internal representations.
How it works
stage by stageThe Daily Edition
The front page in your inbox each weekday morning: the lead, the picks, and why each one earned its place.
What it released
Code