Just One Layer Norm Guarantees Stable Extrapolation

TL;DR

This paper demonstrates that adding just one Layer Norm to neural networks guarantees bounded outputs during extrapolation, providing stability far from training data, supported by theoretical analysis and empirical validation.

Contribution

It proves that a single Layer Norm fundamentally changes the neural tangent kernel, ensuring stability in extrapolation, a novel theoretical insight with practical implications.

Findings

Layer Norm transforms NTK into a bounded-variance kernel.

Networks with Layer Norm produce bounded outputs outside training data.

Empirical results show Layer Norm mitigates instability in finite networks.

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.