Gorio Tech Blog search

Transformers without Normalization | Summary

한국어로 읽기

|
Contents

This article explains the key points of Transformers without Normalization.

  • 2025-03-13 (arXiv)
  • Zhu, Jiachen, Chen, Xinlei, He, Kaiming, LeCun, Yann, Liu, Zhuang.
  • FAIR, Meta, New York University, MIT, Princeton University
  • Paper
  • Github
  • Project Page

Read this article in Korean


Summary

  • Transformers without Normalization tests whether Transformers can train effectively without normalization layers, which rescale internal activations using input statistics. Dynamic Tanh (DyT), an element-wise replacement that applies learned scaling and a bounded hyperbolic tangent, achieves 83.6% ImageNet-1K classification accuracy versus 83.1% for the Large Vision Transformer with normalization under the same training hyperparameters.
  • The complete replacement is DyT(x) = γ * tanh(αx) + β, where tanh denotes the hyperbolic tangent, α is one learned scalar per layer, and γ and β scale and shift each channel, or feature dimension. Its design follows observations that normalization often produces approximately linear input-output mappings near zero while compressing extreme activations.
  • Experiments span supervised and self-supervised vision, image generation, language, speech, and genomic sequence modeling, with broadly comparable performance to normalized baselines. Language models require initialization tuning and an additional embedding scalar; compiled implementations show no measured speed advantage, and replacing batch normalization, which computes statistics across examples, reduces accuracy in classic convolutional networks.

1 Introduction

Normalization layers improve optimization and training stability, which has encouraged their retention even when other architectural components change. The paper tests whether learned activation scaling and compression of extreme values can preserve these benefits without computing activation statistics.

The residual block structure remains intact while normalization based on activation statistics is replaced by tanh with learned input scaling. The retained scale-and-shift stage shows that DyT changes the normalization operation while preserving learned output rescaling.

Transformer blocks with Layer Norm and Dynamic Tanh.
Transformer blocks with Layer Norm and Dynamic Tanh.

2 Background: Normalization Layers

For an input of shape (B, T, C), representing batch size, token count, and embedding dimension, normalization generally computes normalization(x) = γ * ((x − μ) / √(σ² + ϵ)) + β. Here μ and σ² are the input mean and variance, ϵ is a small constant, and γ and β are learned channel vectors.

  • Batch normalization (BN) computes statistics across the batch and token dimensions for each channel.
  • Layer normalization (LN) computes mean and variance across channels independently for each token.
  • Root mean square normalization (RMSNorm) omits mean subtraction and rescales using the mean squared activation. The evaluated Transformers use LN except for LLaMA, which uses RMSNorm.

3 What Do Normalization Layers Do?

The authors inspect normalization inputs and outputs before the learned scaling and shifting in trained ViT-B, the Base Vision Transformer for image recognition; wav2vec 2.0 Large, a speech representation model; and DiT-XL, an Extra Large Diffusion Transformer for image generation. Earlier layers show mostly linear mappings, whereas deeper layers often show full or partial S-shaped curves.

  • Approximately 99% of plotted points occupy the central, approximately linear region. Extreme values, such as inputs above 50 or below −50 in the inspected ViT, undergo stronger compression.
  • Coloring activations by token reveals straight lines with different slopes because tokens have different statistics. Coloring by channel reveals that a few channels contribute many extreme values.
  • These observations motivate the hypothesis that nonlinear suppression of extreme values helps training. They do not establish a universal equivalence between LN and tanh.

Selected layers from image recognition, speech, and image generation models show approximately linear central regions and compressed tails. These sampled mappings motivate scaled tanh without establishing that every normalization layer implements the same function.

Normalization input-output mappings across three trained models.
Normalization input-output mappings across three trained models.

Grouping points by token reveals different straight-line slopes, while grouping by channel isolates channels with extreme values. The panels explain how rescaling each token using its own statistics can yield an aggregate S-shaped scatter plot.

Layer normalization mappings colored by token and channel.
Layer normalization mappings colored by token and channel.

4 Dynamic Tanh (DyT)

A complete DyT layer computes DyT(x) = γ * tanh(αx) + β. It replaces normalization in attention blocks, feed-forward network (FFN) blocks, which transform each token’s features, and the final normalization layer, while preserving the architecture’s existing activation functions.

  • Each layer learns one scalar α and per-channel vectors γ and β. The default initialization is α = 0.5, γ = 1, and β = 0, with exceptions for language models.
  • During a forward pass, each element is processed independently without means, variances, or other aggregation. The learned α changes through training rather than being computed from the current input.
  • Project page and code: https://jiachenzhu.github.io/DyT.

5 Experiments

For supervised ImageNet-1K classification, DyT and LN use matching training hyperparameters. Top-1 accuracy, the fraction of examples whose highest-scoring prediction is correct, changes from 82.3% to 82.5% for ViT-B, 83.1% to 83.6% for ViT-L, 83.7% to 83.7% for ConvNeXt-B, and 84.3% to 84.4% for ConvNeXt-L; B and L denote Base and Large, and ConvNeXt is the convolutional architecture tested alongside Transformers.

  • Training loss curves for ViT-B and ConvNeXt-B follow similar trajectories over 300 epochs, supporting comparable convergence behavior in these experiments.

Under matching training hyperparameters, DyT changes supervised accuracy by 0.0 to 0.5 percentage points across the four architectures and sizes. ViT-L has the largest reported gain.

Supervised ImageNet-1K top-1 accuracy with LN and DyT.
Supervised ImageNet-1K top-1 accuracy with LN and DyT.

Self-supervised vision experiments pretrain without labels and then fine-tune for classification with labels. Masked autoencoders (MAE), trained to reconstruct masked image content, achieve LN/DyT accuracies of 83.2%/83.2% for ViT-B and 85.5%/85.4% for ViT-L; DINO, a method trained with an objective over learned image representations, achieves 83.2%/83.4% with patch size 16 and 84.1%/84.5% with patch size 8.

  • For image generation, DiT-B, DiT-L, and DiT-XL use patch sizes 4, 4, and 2, respectively. Fréchet Inception Distance (FID), the reported image generation quality measure for which lower is better, changes from 64.9 to 63.9, 45.9 to 45.7, and 19.9 to 20.8.
  • The DiT experiments retain the affine parameters used to condition generation on class labels. Appendix A also specifies that DyT removes the zero initialization retained by the LN baseline, so this comparison includes an initialization change.

DyT lowers FID for DiT-B and DiT-L but raises it from 19.9 to 20.8 for DiT-XL. The largest tested generator therefore performs worse on this metric.

ImageNet generation FID with LN and DyT, where lower is better.
ImageNet generation FID with LN and DyT, where lower is better.

LLaMA large language models (LLMs) with 7B, 13B, 34B, and 70B parameters are pretrained on 200B tokens from The Pile text dataset, where B denotes billion. Across 15 zero-shot tasks, evaluated without task-specific examples using the lm-eval evaluation framework, the reported average scores are identical between RMSNorm and DyT at 0.513, 0.529, 0.536, and 0.549; final RMSNorm/DyT training losses are 1.59/1.60, 1.53/1.54, 1.50/1.50, and 1.45/1.45.

  • DyT language models add a learned scalar after the embedding layer and use tuned α initialization. Their loss curves closely track the RMSNorm baselines.
  • For speech pretraining on the LibriSpeech speech dataset, wav2vec 2.0 Base has validation loss 1.95 with either method, while Large changes from 1.92 to 1.91. The initial group normalization layer, which computes statistics within groups of channels and handles raw input data here, remains in both speech models.
  • For genomic sequence modeling, HyenaDNA and Caduceus pretrain on GRCh38 human reference genome data and are evaluated on GenomicBenchmarks, a collection of genomic sequence classification datasets. LN and DyT both achieve average accuracies of 85.2% for HyenaDNA and 86.9% for Caduceus.

The reported average zero-shot score is unchanged at every model size. Training loss is 0.01 higher for DyT at 7B and 13B and equal at 34B and 70B, with the additional embedding scalar and tuned initialization used for DyT.

LLaMA training loss and average scores across 15 zero-shot tasks.
LLaMA training loss and average scores across 15 zero-shot tasks.

6 Analysis

The analysis tests the contributions of the squashing function and learned scalar, examines how α relates to activation scale, and compares DyT with previous approaches to training without normalization.

6.1 Ablations of tanh and α

Replacing tanh with the identity function, which leaves its input unchanged, causes training to diverge after reaching peak accuracies of 58.5% for ViT-S and 61.0% for ViT-B; S denotes Small. Bounded alternatives train successfully, but tanh gives the highest accuracy: 80.3%/82.5% for ViT-S/ViT-B, compared with 79.9%/82.2% for hardtanh, a clipped linear function, and 79.6%/81.6% for sigmoid, a squashing function bounded between 0 and 1.

  • Removing learned α reduces ViT-B accuracy from 82.5% to 81.1% with tanh, 82.2% to 80.7% with hardtanh, and 81.6% to 80.7% with sigmoid.
  • These results support bounded compression and learned scaling under the tested training recipe. The authors’ explanation for tanh’s advantage, its smoothness and zero-centered output, remains a hypothesis.

Identity replacements diverge in both tested model sizes, whereas all three bounded alternatives complete training. Tanh has the highest accuracy, although the experiment does not isolate which property accounts for its advantage over the other bounded functions.

ImageNet-1K accuracy under alternative squashing functions.
ImageNet-1K accuracy under alternative squashing functions.

Learning α improves ViT-B accuracy for every tested squashing function, including a rise from 81.1% to 82.5% with tanh. Bounded compression alone does not recover the full DyT result under this recipe.

ViT-B accuracy with and without learned input scaling.
ViT-B accuracy with and without learned input scaling.

6.2 Values of α

During training, α broadly follows the inverse standard deviation, 1/std, of input activations in the inspected layers. Final α values also correlate with 1/std across layers in ViT-B and ConvNeXt-B, suggesting that learned scaling adapts to activation magnitude.

  • Deeper layers tend to have activations with larger standard deviations. Unlike LN’s separate statistics for each token, α supplies one learned scale for an entire layer; it cannot alone perform the nonlinear suppression supplied by tanh.

Training trajectories and final layer values associate larger α with larger inverse activation standard deviation. The association supports interpreting α as learned adaptation to activation scale without demonstrating exact statistical normalization.

Learned α in relation to inverse activation standard deviation.
Learned α in relation to inverse activation standard deviation.

6.3 Comparison with Other Methods

The comparison includes Fixup and SkipInit, initialization schemes intended to prevent large initial activations and gradients, and σReparam, a method that constrains the spectral norm, or largest singular value, of weight matrices. Because Fixup and SkipInit require lower learning rates to avoid divergence, the authors search learning rates for all compared methods, including DyT.

  • For ViT-B, ViT-L, MAE ViT-B, and MAE ViT-L, DyT achieves 82.8%, 83.6%, 83.7%, and 85.8%. The corresponding σReparam results are 82.5%, 83.0%, 83.2%, and 85.4%; Fixup achieves 77.2%, 78.1%, 73.7%, and 74.1%, and SkipInit achieves 74.1%, 75.6%, 73.1%, and 74.0%.
  • These tuned DyT results use a different comparison condition from the original-hyperparameter results in Section 5.

DyT has the highest reported accuracy across the four configurations after a learning-rate search for all methods. This search condition explains why some DyT values differ from the original-recipe results.

ImageNet-1K comparison with alternative normalization-removal methods.
ImageNet-1K comparison with alternative normalization-removal methods.

7 Initialization of α

The paper denotes the initial value of α by α₀ and examines its sensitivity separately for language and non-language models. Tuning generally yields small gains outside language modeling, while language model performance benefits from initialization chosen by model size and layer location.

7.1 Initialization of α for Non-LLM Models

Most non-language experiments perform consistently over a broad range of α₀, with values between 0.5 and 1.2 generally giving good results. Supervised ViT-L is an exception: under its original recipe, α₀ above 0.6 causes divergence.

  • A stability sweep varies width from 512 to 1280, depth from 8 to 32, and learning rates 0.002, 0.004, and 0.008. Larger models and higher learning rates increase failure risk; smaller α₀ or lower learning rates improve stability.
  • The default α₀ = 0.5 yields a stability pattern similar to LN in this sweep.

Failure regions generally expand with model size, learning rate, and larger DyT initialization. The α₀ = 0.5 setting has a stability pattern similar to LN, while smaller α₀ allows more configurations to complete training.

Training completion across model sizes, learning rates, and α₀ values.
Training completion across model sizes, learning rates, and α₀ values.

7.2 Initialization of α for LLMs

Language model initialization is selected by comparing training losses after 30B tokens. The reported attention/other α₀ pairs are 0.8/0.2 for LLaMA 7B, 0.6/0.15 for 13B, and 0.2/0.05 for both 34B and 70B, where “other” covers FFN locations and the final DyT layer.

  • The 7B heatmap reaches loss 1.820 at 0.8/0.2, and the 13B heatmap reaches 1.767 at 0.6/0.15. Both favor higher initialization in attention blocks than at other locations.
  • A separate sweep over widths 1024, 2048, 4096, and 8192 and depths 8, 16, 32, and 64 selects pairs 1.0/1.0, 1.0/0.5, 0.8/0.2, and 0.2/0.05 by width, unchanged across the tested depths. Width therefore determines the selected initialization within this sweep.
  • Table 10 prints width 8196 for the 34B and 70B models, whereas Table 11 uses 8192 in the controlled width sweep.

The lowest displayed losses occur with different attention and other-location initializations: 0.8/0.2 for 7B and 0.6/0.15 for 13B. These 30B-token sweeps support using distinct initial values by location in the language experiments.

LLaMA loss after 30B tokens across location-specific α₀ settings.
LLaMA loss after 30B tokens across location-specific α₀ settings.

Related work connects normalization to gradient stability, initialization sensitivity, smoother optimization, and generalization. The paper also reviews normalization placement in Transformers, specialized initialization, constraints on weights, and gradual removal of normalization through fine-tuning.

  • DyT replaces activation normalization with an element-wise transformation while retaining most of the architecture and training recipe. Its empirical comparisons with alternative removal methods cover the tested vision settings.

9 Limitations

The demonstrated replacement primarily concerns LN and RMSNorm. Directly replacing BN reduces ImageNet-1K accuracy from 76.2% to 68.9% in ResNet-50 and from 72.7% to 71.0% in VGG19, two classic convolutional networks, under matching hyperparameters.

  • The efficiency benchmark uses LLaMA 7B on one Nvidia H100 graphics processing unit with BF16, a numerical precision format, and measures totals for 100 forward passes or 100 forward-backward passes on one 4096-token sequence. Without extra optimizations, full-model inference time falls from 14.1s to 13.0s and training time from 42.6s to 39.1s.
  • After compiling only the RMSNorm or DyT layers with torch.compile, PyTorch’s compilation facility, both models take 12.3s for inference and 38.9s for training. These measurements show no speed advantage in the compiled setting.
  • The result tables report no uncertainty estimates for the small performance differences, limiting assessment of their statistical reliability.

After layer compilation, RMSNorm and DyT have the same reported full-model totals of 12.3s for inference and 38.9s for training. The unoptimized timing advantage therefore does not persist in this compiled benchmark.

LLaMA 7B timing with compiled RMSNorm or DyT layers.
LLaMA 7B timing with compiled RMSNorm or DyT layers.

Direct replacement of BN costs 7.3 percentage points in ResNet-50 and 1.7 percentage points in VGG19 under matching recipes. These results limit the demonstrated applicability of DyT beyond LN and RMSNorm.

ImageNet-1K accuracy after replacing batch normalization with DyT.
ImageNet-1K accuracy after replacing batch normalization with DyT.

10 Conclusion

The experiments show that learned scaling and bounded compression can replace activation-statistic normalization while preserving competitive performance in the evaluated Transformer settings. Applicability to BN-based networks and efficiency across hardware and deployment settings remain unresolved.

Appendix

  • A Experimental Settings: Supervised vision follows ConvNeXt recipes, with beta parameters (0.9, 0.95) for the AdamW optimizer in ViT-L and stochastic depth, which randomly drops residual branches during training, at rates 0.1 for ViT-B and 0.4 for ViT-L. DiT learning rates are selected on LN models and transferred to DyT, while zero initialization is removed only for DyT. LLaMA adds an embedding scalar initialized to √d, where d is the embedding dimension, and uses a batch size of 4M tokens, learning rates 3e-4 for 7B/13B and 1.5e-4 for 34B/70B, and 200B training tokens; M denotes million.
  • A Experimental Settings: The 15 lm-eval identifiers are anli_r1, anli_r2, anli_r3, arc_challenge, arc_easy, boolq, hellaswag, openbookqa, piqa, record, rte, truthfulqa_mc1, truthfulqa_mc2, wic, and winogrande. Speech experiments preserve the first group normalization layer and otherwise use the official hyperparameters. MAE, DINO, HyenaDNA, and Caduceus use publicly released implementations without hyperparameter tuning in the main comparisons.
  • B Hyperparameters: Learning-rate tuning improves DyT ViT-B from 82.5% at 4e-3 to 82.8% at 6e-3, MAE ViT-B from 83.2% at 2.4e-3 to 83.7% at 3.2e-3, and MAE ViT-L from 85.4% at 2.4e-3 to 85.8% at 3.2e-3. Independently tuning α₀ to 1.0 raises ViT-B to 82.6% and MAE ViT-B to 83.4%, and reduces wav2vec 2.0 Large validation loss from 1.91 to 1.90. The tables mark cases without improvement using “-”.
  • C Efficiency of DyT: Unoptimized measurements report full-model time reductions of 7.8% for inference and 8.2% for training, but compiling the replacement layers eliminates the measured advantage. D Replacing Batch Normalization with DyT: Both ResNet-50 and VGG19 lose accuracy under the original torchvision training recipes. The authors hypothesize that the greater frequency of BN layers in these architectures contributes to the difficulty, but do not test that explanation directly.

Brief Thoughts

DyT provides a simple architectural alternative supported by experiments across several domains. The ablations strengthen the evidence: removing bounded compression causes divergence in the tested vision recipe, while removing learned scaling reduces accuracy.

Its practical benefits depend on the setting. Language models need additional scaling and initialization tuning, DiT-XL has worse FID, and compiled timing shows parity; these results support evaluating DyT as an architectural option without establishing general accuracy or runtime gains.