SeeDNorm: Self-Rescaled Dynamic Normalization

TL;DR

SeeDNorm addresses the limitation of traditional normalization layers that discard input scale information by introducing a self rescaled dynamic normalization that modulates the RMSNorm scaling factor with a data dependent term. By computing a per token dynamic scale through a tanh based function of the input and learnable parameters alpha beta, SeeDNorm preserves input norm information in both forward and backward passes, enabling data dependent gradient adjustments. The authors provide forward and gradient analyses, propose a multi head variant to reduce gradient variance, and validate SeeDNorm across large language models and computer vision tasks, showing faster convergence and improved downstream performance with minimal parameter overhead. The approach yields robust improvements over RMSNorm, LayerNorm and DyT, and demonstrates strong potential for general adoption in transformer based architectures including MoE models and conditional generation scenarios.

Abstract

Normalization layer constitutes an essential component in neural networks. In transformers, the predominantly used RMSNorm constrains vectors to a unit hypersphere, followed by dimension-wise rescaling through a learnable scaling coefficient $γ$ to maintain the representational capacity of the model. However, RMSNorm discards the input norm information in forward pass and a static scaling factor $γ$ may be insufficient to accommodate the wide variability of input data and distributional shifts, thereby limiting further performance improvements, particularly in zero-shot scenarios that large language models routinely encounter. To address this limitation, we propose SeeDNorm, which enhances the representational capability of the model by dynamically adjusting the scaling coefficient based on the current input, thereby preserving the input norm information and enabling data-dependent, self-rescaled dynamic normalization. During backpropagation, SeeDNorm retains the ability of RMSNorm to dynamically adjust gradient according to the input norm. We provide a detailed analysis of the training optimization for SeedNorm and proposed corresponding solutions to address potential instability issues that may arise when applying SeeDNorm. We validate the effectiveness of SeeDNorm across models of varying sizes in large language model pre-training as well as supervised and unsupervised computer vision tasks. By introducing a minimal number of parameters and with neglligible impact on model efficiency, SeeDNorm achieves consistently superior performance compared to previously commonly used normalization layers such as RMSNorm and LayerNorm, as well as element-wise activation alternatives to normalization layers like DyT.

Figures (14)

• Figure 1: Comparisons between SeeDNorm and prior methods across diverse tasks in language modeling and vision. The first two figures respectively depict training loss curve comparisons and average downstream taskaccuracy between the OLMoE-1.3B muennighoff2024olmoeopenmixtureofexpertslanguage baseline (using RMSNorm NEURIPS2019_1e8a1942_RMSNorm) and the SeeDNorm-equipped model, following training on 500B tokens. The latter two figures show training loss comparisons between DyT-based Zhu2025DyT baseline models and SeeDNorm-based models in image generation and MAE He_2022_CVPR_mae pre-training. All loss curves are smoothed with a 0.99 EMA.
• Figure 2: The first two subplots respectively show a comparison of training loss and average downstream task accuracy between OLMoE-7B baseline and the counterparts with SeeDNorm. The last two subplots show a comparison of training loss of OLMo2-550M and OLMo2-1B baseline models and their counterparts with SeeDNorm incorporated. All curves are smoothed using EMA with a coefficient of 0.99.
• Figure 3: The first two subplots show FID comparison at different training steps wich CFG=1.5, and the last two subplots show loss curves of DiT-XL/2 in image generation and ViT-L in MAE. All models are respectively augmented with our proposed SeeDNorm and DyT.
• Figure 4: Two subplots present training loss curves of SeeDNorm with different activation functions, and with various $\bm{\alpha}$ initialization strategies. The loss curves are smoothed using EMA with a coefficient of 0.99.
• Figure 5: Comparisons of the accuracy of all downstream tasks from Table \ref{['validset_and_downsteam_tasks']} in OLMoE-1.3B when using SeeDNorm as the normalization layer versus the default RMSNorm. The figure illustrates the evolution of downstream task accuracy as the total training tokens increase during training, with transparent lines indicating unsmoothed results and solid lines denoting 0.99 EMA-smoothed results.

Theorems & Definitions (1)

• Proposition A.1