PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformers without Normalization

Jiachen Zhu, Xinlei Chen, Kaiming He, Yann LeCun, Zhuang Liu

arXiv:2503.10622v2cs.LGcs.AIcs.CLcs.CV

TL;DR

Normalization layers are widely regarded as crucial for training deep networks, yet their necessity in Transformers is questioned by this paper. It replaces them with Dynamic Tanh, an element-wise learnable-scaling operation, and finds comparable or better performance across settings, often without original hyperparameter tuning.

  • Problem

    Normalization layers are widely regarded as crucial or indispensable for effectively training deep networks, leaving their necessity in Transformers insufficiently challenged.

  • Method

    Dynamic Tanh (DyT) replaces Transformer normalization layers with tanh(αx), using a learnable scaling factor and bounded squashing without computing activation statistics.

  • Results

    Across various settings, Transformers with DyT match or exceed normalized counterparts, train stably, and often need no tuning of the original training hyperparameters.

  • Takeaways & Limitations

    The findings challenge the conventional understanding that normalization layers are indispensable for training modern neural networks and provide empirical insights into their properties.

  • Takeaways & Limitations

    Experiments focus on LN and RMSNorm, while preliminary experiments indicate that DyT struggles to directly replace BN and adaptation to other normalization types remains understudied.

Abstract

from arXiv · show

Normalization layers are ubiquitous in modern neural networks and have long been considered essential. This work demonstrates that Transformers without normalization can achieve the same or better performance using a remarkably simple technique. We introduce Dynamic Tanh (DyT), an element-wise operation $DyT($x$) = \tanh(α$x$)$, as a drop-in replacement for normalization layers in Transformers. DyT is inspired by the observation that layer normalization in Transformers often produces tanh-like, $S$-shaped input-output mappings. By incorporating DyT, Transformers without normalization can match or exceed the performance of their normalized counterparts, mostly without hyperparameter tuning. We validate the effectiveness of Transformers with DyT across diverse settings, ranging from recognition to generation, supervised to self-supervised learning, and computer vision to language models. These findings challenge the conventional understanding that normalization layers are indispensable in modern neural networks, and offer new insights into their role in deep networks.

1 Introduction

Normalization layers are widely viewed as essential for training deep networks, but this paper challenges that belief with Dynamic Tanh (DyT), a simple normalization replacement for Transformers.

  • Normalization layers became fundamental because they improve results and accelerate and stabilize convergence, especially in wider and deeper networks.
  • Recent architectures commonly replace attention or convolution layers while retaining normalization layers, reflecting their perceived indispensability.
  • DyT applies tanh(αx) element-wise, using a learnable scale to emulate Layer Normalization's activation scaling and extreme-value squashing without computing activation statistics.
  • DyT directly replaces normalization layers in Transformer architectures, including attention blocks, FFN blocks, and final normalization layers.
  • Transformers with DyT train stably and achieve high final performance across diverse settings, often without tuning the original architecture's training hyperparameters.

2 Background: Normalization Layers

Normalization layers share a formulation based on input statistics and learnable affine parameters, while different methods vary in the axes and ranges used to compute those statistics.

  • Given input x with shape (B, T, C), normalization generally computes an output from the input, its mean and variance, and learnable scaling and shifting parameters.
  • Different normalization methods mainly differ in how they compute and broadcast the input mean and variance.
  • Batch normalization computes statistics across both batch and token dimensions and has primarily been used in ConvNet models.
  • Group normalization and instance normalization retain the common formulation but compute statistics over different axes and ranges for specialized ConvNet tasks.
  • Layer Normalization computes statistics independently for each token in each sample, whereas RMSNorm removes mean-centering and normalizes with zero mean.
  • Layer Normalization mappings can be mostly linear in earlier layers, while deeper layers show more intriguing input-output behavior resembling tanh.

3 What Do Normalization Layers Do?

In trained Transformers, deeper layer-normalization layers produce collectively tanh-like, S-shaped input-output mappings. LN remains linear per token, but differing token statistics and channel ranges collectively yield nonlinear squashing of extreme activations.

  • Analysis setup: The analysis measures tensor elements immediately before and after LN in ViT-B, wav2vec 2.0 Large, and DiT-XL forward passes.Because LN preserves tensor dimensions, each input element can be matched directly to its output for visualization.
  • Tanh-like mappings with layer normalization: Across three trained networks, deeper LN layers show input-output curves resembling full or partial tanh-shaped S-curves.Earlier LN layers are mostly linear, while the tanh-like behavior is more apparent in deeper layers.
  • Normalization by tokens and channels: LN linearly transforms each token’s activations, but differing token means and standard deviations prevent collective linearity across all tensor elements.Tokens with smaller input ranges have smaller variance and therefore steeper normalized straight-line mappings.
  • Tanh-like mappings with layer normalization: Most points (∼99%) lie in the central, approximately linear range, while extreme values such as x > 50 or x < -50 in ViT are squashed toward less extreme values.This nonlinear, disproportionate squashing cannot be approximated by a simple affine transformation.
  • Related observations: Prior work is cited as linking LN’s nonlinearities to representational capacity, while the squashing behavior is compared with biological-neuron saturation for large inputs.These are presented as related observations rather than as the paper’s primary analysis result.
  • Normalization by tokens and channels: Channels span different input ranges, and channels with large extreme values contribute segments that are squashed most strongly by LN.The Figure 4 visualization groups activations by tokens and channels to explain how per-token linearity produces a collective S-shaped curve.

4 Dynamic Tanh (DyT)

Dynamic Tanh is proposed as a simple, element-wise replacement for normalization layers, motivated by the tanh-like mappings observed in LN. It can replace normalization throughout existing Transformer architectures while preserving nonlinear extreme-value squashing and requiring little hyperparameter retuning.

  • Definition and motivation: Dynamic Tanh is defined as an element-wise operation that replaces normalization layers with a tanh transformation and learnable scaling.Its purpose is to emulate LN’s tanh-like input-output behavior.
  • Parameterization: DyT uses a learnable scalar α to scale inputs before tanh, while learnable per-channel γ and β restore output scale and offset.γ and β are the same type of affine parameters used in normalization layers.
  • Integration: One DyT layer directly replaces one normalization layer in attention blocks, FFN blocks, and the final normalization layer.The original activation functions, such as GELU or ReLU, remain unchanged.
  • Integration: DyT generally performs well with little need to tune the original architectures’ hyperparameters.The default α initialization is 0.5 except for LLM training, while γ starts at ones and β at zeros.
  • Implementation: The pseudocode represents DyT inputs with shape [B, T, C], where B is batch size, T is token count, and C is embedding dimension.The implementation initializes α, γ, and β as learnable parameters.
  • Operational distinction: DyT processes each tensor element independently without computing statistics or other aggregation.It preserves nonlinear squashing of extreme values while approximately linearly transforming central inputs.

5 Experiments

The experiments replace LN or RMSNorm with DyT while retaining the original architectures and training protocols. Across vision, generation, language, speech, and DNA tasks, DyT generally matches or improves normalized-model performance.

  • Experimental setup: DyT replaces LN or RMSNorm using the normalized models’ original architectures, protocols, and hyperparameters.The replacement is evaluated across diverse tasks and domains.
  • Vision: DyT performs slightly better than LN for supervised ViT and ConvNeXt classification across architectures and model sizes.The comparison uses ImageNet-1K top-1 classification accuracy.
  • Vision: DyT performs on par with LN across self-supervised ImageNet-1K pretraining methods and model sizes.The evaluated methods are MAE and DINO.
  • Generation: DyT achieves comparable or improved FID over LN across DiT model sizes on ImageNet image generation.FID is evaluated using the standard ImageNet reference batch.
  • Language: DyT and RMSNorm have closely aligned LLaMA pretraining loss curves across model sizes.LLaMA models are pretrained on The Pile and additionally evaluated on zero-shot language-model tasks.
  • Speech and DNA: DyT performs comparably to LN in wav2vec 2.0 speech pretraining and maintains comparable performance to LN for DNA sequence modeling.The speech evaluation reports validation loss, while the DNA evaluation uses GenomicBenchmarks classification accuracy.

6 Analysis

Ablations show that both tanh and the learnable α are important to DyT’s behavior and performance. The learned α tracks activation scale, while DyT outperforms the tested alternative normalization-removal methods.

  • Ablations of tanh and α: Removing tanh and replacing it with the identity function causes unstable training and divergence, whereas squashing functions enable stable training.Among tanh, hardtanh, and sigmoid, tanh achieves the best performance.
  • Ablations of tanh and α: Removing the learnable α degrades performance across all tested squashing functions.This result identifies α as important for overall model performance.
  • Values of α: During training, α tracks the inverse standard deviation of activations, supporting its role in maintaining activations within a suitable range.Both α and 1/std fluctuate consistently during training.
  • Values of α: After training, larger inverse activation standard deviations generally correspond to larger learned α values, while deeper layers tend to have larger activation standard deviations.The analysis suggests α partially functions as a normalization mechanism by approximating 1/std.
  • Comparison with Other Methods: DyT consistently outperforms Fixup and SkipInit across the tested ViT configurations after learning-rate searches.Fixup and SkipInit require significantly lower learning rates to prevent divergence.

7 Initialization of α

The default α initialization is broadly robust for non-LLM models but requires task- and architecture-dependent tuning for LLMs. In LLaMA, model width and block location are more influential than depth.

  • Initialization overview: Tuning α0 rarely produces significant performance improvements, except in LLM training where careful tuning yields noticeable gains.This establishes different initialization sensitivity for LLM and non-LLM models.
  • Initialization of α for Non-LLM Models: Non-LLM performance remains stable for α0 values between 0.5 and 1.2, with supervised ViT-L diverging when α0 exceeds 0.6.The ViT-L instability can be restored by reducing the learning rate.
  • Initialization of α for Non-LLM Models: Increasing model size or learning rate requires lowering α0 for stable training, while smaller α0 generally improves stability.Lowering the learning rate also mitigates instability at higher α0.
  • Initialization of α for Non-LLM Models: α0 = 0.5 is used as the default for non-LLM models because it provides stability comparable to LN while maintaining strong performance.The default applies across the paper’s non-LLM experiments.
  • Initialization of α for LLMs: Larger LLaMA models require smaller α0 values, and higher α0 in attention blocks with lower values elsewhere improves performance.The attention-versus-other initialization distinction is reported for LLaMA training.
  • Initialization of α for LLMs: Model width significantly affects optimal α0, with wider networks benefiting from smaller values, whereas model depth has negligible influence.Wider networks also require more uneven attention-versus-other initialization.

8 Related Work

Prior work explains normalization through optimization, initialization, spectral, and representational mechanisms, while other studies seek to remove or reduce normalization dependence. DyT differs by using a simple drop-in replacement with minimal architectural and recipe changes.

  • Mechanisms of Normalization layers: Prior studies associate normalization with stabilized gradient flow, reduced initialization sensitivity, moderated outlier eigenvalues, auto-tuned learning rates, and smoother loss landscapes.Later work also connects normalization with sharpness reduction and generalization.
  • Normalization in Transformers: Layer normalization has been especially effective for sequential data, and recent work links its nonlinearity to increased representational capacity.RMSNorm is also widely used in contemporary language models.
  • Removing normalization: Normalization-removal methods include self-normalizing activations, specialized initialization schemes, and architectural modifications to reduce reliance on normalization layers.These approaches address training without normalization through different mechanisms.
  • Removing normalization: DyT requires minimal modifications to the architecture and training recipe compared with previous approaches to removing normalization.It replaces normalization layers directly rather than requiring broader architectural changes.

9 Limitations

The experiments establish scope boundaries for DyT: its adaptation to normalization types beyond LN and RMSNorm remains unresolved, and its efficiency benefits are not guaranteed after optimization.

  • DyT is evaluated primarily as a replacement for LN and RMSNorm; adaptation to other normalization layers remains insufficiently studied.Preliminary experiments indicate difficulty replacing BN directly in classic networks such as ResNets.
  • DyT provides no speedup over properly compiled or optimized normalization-layer models, leaving hardware- and deployment-specific benefits uncertain.The paper specifically reports no speedup after proper compilation/optimization.

10 Conclusion

The paper shows that Transformers can be trained without normalization layers by replacing them with Dynamic Tanh. Across varied settings, DyT models match or exceed normalized counterparts, challenging the presumed necessity of normalization.

  • DyT replaces traditional normalization by scaling activations with learnable α and squashing extremes through a bounded, S-shaped tanh function.The operation is designed to capture normalization-layer behavior with a simpler element-wise function.
  • Across various settings, models with DyT match or exceed the performance of their normalized counterparts.The reported settings span modern neural networks, particularly Transformers.
  • These findings challenge the conventional understanding that normalization layers are necessary for training modern neural networks.The study also presents its results as insights into normalization-layer mechanisms.

A Experimental Settings

The experimental settings largely reuse established implementations and hyperparameters, while selected tasks receive targeted adjustments. Additional studies test whether DyT benefits from tuning learning rates or α initialization.

  • Supervised image classification: ImageNet-1K supervised classification follows ConvNeXt recipes, with original hyperparameters retained for ConvNeXt-B and ConvNeXt-L.ViT models use ConvNeXt-based settings with specified changes for ViT-L and stochastic depth.
  • Large Language Models: LLaMA DyT models add a learnable scalar after embeddings because otherwise activations are too small and training struggles to progress.The scalar is initialized to the square root of the model embedding dimension and resembles the original Transformer's fixed scaling.
  • Large Language Models: 15 zero-shot commonsense reasoning tasks are evaluated for pretrained LLaMA models, with average performance reported across tasks.The tasks are drawn from lm-eval and include ANLI, ARC, BoolQ, HellaSwag, and Winogrande.
  • Self-supervised learning in speech: Wav2vec 2.0 retains its first group normalization layer for input-data normalization and uses the official implementation without hyperparameter changes.The Base and Large models are evaluated using final validation loss.
  • Other tasks: For MAE, DINO, HyenaDNA, and Caduceus, publicly released code is used without hyperparameter tuning for both LN and DyT models.
  • Hyperparameter tuning: Learning-rate tuning yields only modest DyT improvements, suggesting LN-optimized original hyperparameters are already well suited to DyT.The comparison is summarized in Table 12.
  • Hyperparameter tuning: Tuning α0 produces only minor gains for selected models, while the default α0 = 0.5 generally achieves near-optimal performance.The comparison is summarized in Table 13.

C Efficiency of DyT

DyT is computationally simpler because it is element-wise and avoids internal reductions, yielding lower unoptimized layer computation but nearly identical latency to RMSNorm after compilation. Its replacement of BN also reduces accuracy in tested classic ConvNets.

  • Uncompiled efficiency: DyT reduces computation time relative to RMSNorm layers on an H100 GPU in unoptimized LLaMA 7B inference and training measurements.The benchmark uses BF16 precision and 100 forward or forward-backward passes on 4096-token sequences.
  • Compiled efficiency: After compilation, RMSNorm and DyT layer latency becomes nearly identical.This result is reported for compiled LLaMA 7B latency measurements.
  • Efficiency mechanism: DyT is element-wise and avoids the reduction operation required within normalization layers, offering a potential hardware advantage when reductions bottleneck execution.The paper also identifies possible fusion with a preceding matrix multiplication as an optimization opportunity.
  • Replacing batch normalization: Replacing BN with DyT causes a noticeable ImageNet-1K classification-accuracy drop in ResNet-50 and VGG19.The paper suggests the greater frequency of BN in these ConvNets may relate to the difficulty of substitution.
Loading 2503.10622v2…