Research map / Learning theory and algorithms / Kernels and regression theory
Just One Layer Norm Guarantees Stable Extrapolation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Uses NTK theory on infinitely wide networks to show that a single Layer Norm makes the NTK bounded, so trained network outputs stay bounded out of distribution, while networks without LN can explode.
Problem
Extrapolation behaviour of neural networks far from the training distribution is poorly understood; existing results are limited to special cases, and out-of-distribution outputs can be pathologically large.
Approach
Applies NTK theory to infinitely wide networks trained to convergence and shows that including at least one Layer Norm transforms the NTK into a bounded-variance kernel, so posterior-mean outputs are bounded for any input, whereas networks with positive n-homogeneous activations and no LN can explode.
Key results
- On the UCI protein task, LN-equipped networks reach OOD R^2 of about 0.50 versus 0.31 for a standard NN and 0.27 for XGBoost.
- On UTKFace, LN-first network achieves OOD R^2 0.47 versus 0.27 for a standard NN, with ID R^2 comparable across models (about 0.66-0.68).
- Toy experiments show standard ReLU MLPs extrapolate linearly (explode) while a single LN layer makes outputs saturate.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add at least one LayerNorm instead of using no LayerNorm | Two-layer MLP on the UCI Protein dataset · UCI Protein OOD | R²: 0.31 ± 0.02 -> 0.50 ± 0.02 | Helped |
| Add at least one LayerNorm instead of using no LayerNorm | Two-layer MLP on the UCI Protein dataset · UCI Protein ID | R²: 0.58 ± 0.01 -> 0.60 ± 0.01 | No clear effect |
| Change LayerNorm placement among every layer, first layer only, and second layer only (ablation) | Two-layer MLP on the UCI Protein dataset · UCI Protein OOD | R²: 0.50 ± 0.02; 0.50 ± 0.02; 0.50 ± 0.01 | No clear effect |
| Add at least one LayerNorm instead of using no LayerNorm | ResNet-18 feature-based age predictor · UTKFace age prediction OOD | R²: 0.27 ± 0.03 -> 0.47 ± 0.03 | Helped |
| Use LayerNorm before the second hidden layer only rather than before the first hidden layer only (ablation) | ResNet-18 feature-based age predictor · UTKFace age prediction OOD | R²: 0.47 ± 0.03 -> 0.36 ± 0.04 | Hurt |
| Add LayerNorm instead of using no LayerNorm | ResNet-18 feature-based age predictor · UTKFace age prediction ID | R²: 0.66 ± 0.01 -> 0.68 ± 0.01 | No clear effect |
| Include one LayerNorm anywhere in the network rather than no LayerNorm | Infinitely wide fully connected network trained to convergence · Out-of-distribution inputs | Boundedness of expected output | Helped |
| Train without LayerNorm rather than including LayerNorm | Infinitely wide fully connected network trained to convergence · Out-of-distribution inputs | Boundedness of expected output | Hurt |
| Vary LayerNorm position or number in the toy MLP (ablation) | Toy one-dimensional extrapolation problems · Sinusoid, linear, and quadratic datasets | Extrapolation stability | No clear effect |
| Replace LayerNorm with BatchNorm (ablation) | Toy neural network · Out-of-distribution extrapolation | Explosion prevention | Hurt |
Surprising findings
- BatchNorm does not prevent the explosion, unlike LayerNorm, consistent with the theory since BatchNorm does not alter the NTK.
- Varying the position of the LN layer has an insignificant effect on predictions in toy experiments and on OOD R^2 in the protein task.
Limitations
- Theoretical guarantees hold for infinite width and fully connected networks with positive n-homogeneous activations (and extensions); finite-width behaviour is only verified empirically.
- The bound is worst-case and may not be tight; bounded extrapolation is not always correct when a linear trend is desired.
Experimental setup
Datasets: UCI Physicochemical Properties of Protein Tertiary Structure UTKFace synthetic 1-D datasets (sinusoid, linear, quadratic)
Benchmarks: UCI Protein dataset (OOD and ID R^2) UTKFace age prediction (OOD and ID R^2)
Metrics: R^2 empirical NTK values prediction magnitude on out-of-distribution inputs
Baselines: Standard NN (no LN) XGBoost
Builds on: Neural Tangent Kernel Tensor Programs Layer Normalization Xu et al. NTK extrapolation analysis
Closest papers by meaning
- Extrapolation and Spectral Bias of Neural Nets with Hadamard Product: a Polynomial Net Study (NeurIPS 2022 · similarity 0.88)
- Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $\mu$ Parametrization (ICML 2025 · similarity 0.87)
- On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling (NeurIPS 2025 · similarity 0.86)
- An Infinite-Width Analysis on the Jacobian-Regularised Training of a Neural Network (ICML 2024 · similarity 0.86)
- The Surprising Effectiveness of Infinite-Width NTKs for Characterizing and Improving Model Training (AAAI 2025 · similarity 0.86)
- Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.