Just One Layer Norm Guarantees Stable Extrapolation

TL;DR

This work addresses the challenge of neural network extrapolation beyond the training distribution by leveraging Neural Tangent Kernel theory. It shows that inserting at least one Layer Norm transforms the NTK into a bounded-variance form, guaranteeing bounded predictions even far from training data, whereas networks without LN can explode under extrapolation. The authors prove both a general unbounded-extrapolation result for LN-free networks and a bound for LN-containing networks, and they corroborate these findings with toy-, protein-, and face-age experiments, demonstrating improved out-of-distribution stability. The results provide a principled architectural cue: including LN can yield safer, more realistic extrapolations in high-stakes tasks, with practical implications for proteomics and biometric age estimation, among others.

Abstract

In spite of their prevalence, the behaviour of Neural Networks when extrapolating far from the training distribution remains poorly understood, with existing results limited to specific cases. In this work, we prove general results -- the first of their kind -- by applying Neural Tangent Kernel (NTK) theory to analyse infinitely-wide neural networks trained until convergence and prove that the inclusion of just one Layer Norm (LN) fundamentally alters the induced NTK, transforming it into a bounded-variance kernel. As a result, the output of an infinitely wide network with at least one LN remains bounded, even on inputs far from the training data. In contrast, we show that a broad class of networks without LN can produce pathologically large outputs for certain inputs. We support these theoretical findings with empirical experiments on finite-width networks, demonstrating that while standard NNs often exhibit uncontrolled growth outside the training domain, a single LN layer effectively mitigates this instability. Finally, we explore real-world implications of this extrapolatory stability, including applications to predicting residue sizes in proteins larger than those seen during training and estimating age from facial images of underrepresented ethnicities absent from the training set.

Figures (7)

• Figure 1: Predictions made by networks with various architectures when trained on synthetic datasets. Red dot show the train set datapoints. The solid lines indicate average values over 5 seeds and shaded areas are 95% confidence intervals of the mean estimator.
• Figure 2: Heatmaps showing the values of empirical NTK values $\Theta(\bm{x}, \bm{x}')$ plotted on domain $\bm{x}, \bm{x}' \in [-25,25]$ with brighter colours indicating higher values. Note that the scales are different for each heatplot, with the values range for the NTK of Standard NN being orders of magnitude higher than others. The displayed values are averages over 5 seeds.
• Figure 3: Histogram of predictions made by each model on the out-of-distribution data for protein experiment (left) and LOWESS cleveland1981lowess trendlines fitted to the relationship between protein surface area and average prediction for each method (right). Both plots are produced from data aggregated from 10 seeds. Shaded areas in the right plot are 95%-confidence intervals.
• Figure 4: LOWESS cleveland1981lowess trendlines fitted to average OOD model prediction on the UTK-Face task as a function of average cosine similarity of ResNet-18 features with training samples. Shaded areas are 95% confidence intervals produced over 250 seeds.
• Figure 5: Predictions made by networks with various activations when trained on synthetic datasets. Plots above consider the case of standard NN without LayerNorm, whereas plots below show the case of varying activations while keeping the LayerNorm in the architecture. Red dots show the train set datapoints. The solid lines indicate average values over 5 seeds and shaded areas are 95% confidence intervals of the mean estimator.

Theorems & Definitions (57)

• Theorem 3.1
• proof : Proof sketch
• Theorem 3.2
• proof : Proof sketch
• Theorem 3.3
• proof : Proof sketch
• Proposition 3.3
• proof : Proof sketch
• Definition C.1
• Remark C.1