Just One Layer Norm Guarantees Stable Extrapolation

We would like to thank the reviewer for reading our submission and providing valuable feedback. We are glad to know the reviewer finds our results interesting. We address all questions/concerns below:

Layernorm position

Intuitively, if after the LayerNorm operation, the distribution of activation is ensured to have a mean of 0 and variance of 1, regardless of the input. As such no matter what the input, the activations after LayerNorm cannot get too “extreme”. As the final output of the network is just the output of the remaining layers of the network, and activations serving as an input to the rest of the network cannot be too “extreme”, the output also cannot get too “extreme” and, as such remains bounded.

Missing related work

We thank the reviewer for pointing out related work. We will include these papers in the related work section. We would like to highlight that while these papers conduct a form of normalisation to improve OOD performance, they appear to be quite different from layer normalisation, and explicit stability bounds are not derived.

MSE vs CE Loss

Our results regarding the resulting NTK of networks with LayerNorm hold regardless of the loss function being used, but the question of how the NTK affects the trained network’s prediction when a loss different than MSE is used is generally an open problem in the literature. However, if an analytical expression was derived connecting the NTK to the trained network’s prediction for a loss other than MSE, we believe one could easily plug our NTK bound (which is a property only of the kernel and is independent of the choice of loss function) into it and analyse the resulting expression.

Generalizability beyond the kernel regime

While this is an interesting and desirable extension to the theorems presented in this work, the machinery for proving such statements beyond the kernel regime is simply not sufficiently developed yet. As the field evolves, this is certainly a question worth addressing, but for now we point the reviewer towards the empirical results demonstrating the practical applicability of our theory: finitely-wide networks also exhibit characteristics of a feature-learning regime, so the supportive results here also lend credibility to such an extension.

Other stabilizing components

For the case of batch norm these results would not hold, as batch norm does not fundamentally alter the NTK and only manifests as an affine transformation. Other stabilising components such as the softmax associated with cross-entropy loss, or even manually clipping network outputs, also of course achieve the stability property, but these are explicit design choices; we focus on LayerNorm, as it is a popular component already present in many deep learning architectures, which interestingly causes this bounded output behaviour organically.

To empirical validate the claim about BatchNorm, we ran more experiments on the sine function with a fully-connected neural network utilising batchnorm instead of LayerNorm, we present new results below (due to new policy we cannot share figures via links, so have to do it via a table):

Where values after ± are 95% CIs over 5 seeds. We can thus see that while LayerNorm quickly saturates for extreme inputs with both negative and positive signs, the same is not true for BatchNorm which exhibits the explosion, similarly as the Standard NN.

RKHS statement

The reviewer is correct in this assertion. However, in the particular case of a network with ReLU activation functions it is indeed true that the RKHS only contains functions exhibiting this property, from which the theorem intuitively follows, rather than an assumption being made on the RKHS as a result of the theorem. This can be seen explicitly in Corollary D.2.

The limitation that NTP is used for the theoretical results

This is again a limitation of the entire field making use of NTK theory. We would like to point out that for finite-width networks, the difference between standard (LeCun) parametrization (SP) and NTP is of no importance when considering proper scaling of the learning rate (instead manifesting only as a difference in the empirical NTK, which is never calculated), and as such all empirical evidence we provide also support applicability of the theory in practical applications using SP. Theoretical applicability to more modern parametrizations such as MuP would require extensions to the theory as mentioned above, although we note that, at initialisation, our assumptions on distribution of parameters remain true here, and thus it would be highly unintuitive that the learned function would sufficiently deviate from our theory.

We would like to thank the reviewer for spending time to read through our submission and write a comprehensive review. We are glad to hear the reviewer finds our theoretical results strong. We address all questions/concerns below. As per reviewer request, we ran a number of additional experiments and include the results below in the relevant sections.

Limitation to Infinite-Width Theory

While it is true our theoretical guarantees hold for infinitely wide networks, this is true for most of the theoretical results in the field of theoretical analysis of neural networks. As such, although it is a limitation, it is a limitation of the entire field at the moment rather than our paper specifically.

Beyond fully‑connected layers: Can Theorem B.2’s extension via Tensor Programs be instantiated for a modern ResNet/Transformer?

Yes, Theorem B.2 can be instantiated with any architecture that can be expressed via Tensor Programs, which includes Transformer and ResNet (and convolutional networks in general).

We conduct a simple experiment with an end-to-end training using a transformer architecture, where during training the model is given sequences of form $y_i = a*i + b$, where i is the index that is restricted to [0,20] during training and $a$ and $b$ are sampled from standard normal. As such the model is trained to infer the pattern in the sequence and continue it. During testing, we use $a=1$, $b=1$ and use $i$ from range [0, 100]. We use a decoder-only transformer architecture with embeddings being linear projections from input space to latent space of size 64, we use 4 attention heads, and we use a total of 2 attention layers. We try two versions of this architecture, one with a single LayerNorm right after the input embeddings and the other one with no LayerNorm anywhere in the entire architecture. We provide the table of results over 5 seeds below:

As such, it would appear the presence of a single LayerNorm makes the output quickly saturate, whereas lack of any LayerNorm causes an explosion to infinity. We plan on adding this experiment to the final version of the manuscript.

Activation assumptions:

We point the reviewer towards Theorem B.2 and Theorem B.3, which extend our boundedness results to a much more general class of activation functions, including all popular modern activation functions such as GeLU and Swish. The weaker statement was included in the main body for consistency with Theorem 3.1, which relies on analytical computation of the NTK, which is not tractable for such activation functions even if empirically true.

To verify this claim experimentally, we rerun the sine experiment with different activation functions. We present new results below (due to new policy we cannot share figure via links, so have to do it via a table):

Where values after ± are 95% CIs over 5 seeds. We can thus see that regardless of the activation function, network with LayerNorm quickly saturates, whereas network without LN explodes.

Boundedness ≠ Correctness:

We agree with the reviewer on this point, as we wrote in the conclusion: “we do not claim that LN ‘solves’ the problem of extrapolation. Instead, we aimed to highlight the differences in extrapolatory behaviour of networks with and without LN, so that practitioners can make more informed architectural choices”. Indeed, making guarantees about OOD predictions would be likely impossible without further assumptions.

Dense Theoretical Presentation

We will use the additional page in camera-ready version to improve presentation and provide a more intuitive explanation of the derived results and proofs.

LN Sufficient but Not Proven Necessary

The reviewer is correct that one can guarantee bounded input in various other ways, simplest one being just manually clipping the output at some user-specified threshold. We focus on LayerNorm, as it is a popular component already present in many deep learning architectures, which interestingly causes this bounded output behaviour organically.

Finite‑width behaviour

As per request of the reviewer, we trained neural networks with different sizes on the sine synthetic dataset and recorded max value in the evaluation domain. We present the results below:

NTK-GP denotes the theoretically calculated max value of the infinite neural network (Gaussian Process with Neural Tangent Kernel). It would thus appear that for widths larger than 512, one can already get the max absolute prediction to converge within 10% of the max value of the theoretical infinite network.

Bound tightness:

As the bound grows with dataset size, we expect it to be relatively tight on small dataset, but it might become loose as the number of datapoints grows. As per reviewer request, below we plot the theoretical bound together with the observed empirical maximum for a network with LN after the first layer and a standard NN on the sine synthetic dataset, where we measure the maximum over the domain [-10K, 10K].

The bound is of the same order of magnitude for small datasets, but starts to become loose for larger ones. However, the output of standard NN can easily exceed even that loose bound, so the bound is by no means vacuous.

We would like to thank the reviewer for taking time to read our submission and ask insightful questions. We are happy to know the reviewer finds our results interesting. We address the questions/concerns below:

Why keeping the output bounded is the desirable goal for dealing with out-of-distribution data

As we wrote in the Conclusion section: “As we go outside of data distribution, what exactly constitutes a “good” extrapolation is unknown to us unless further assumptions can be made. A stable (i.e. bounded) extrapolation, while desirable in many cases outlined in this paper, does not need to be always correct.”. Within the paper we merely wanted to point out that the introduction of LayerNorm makes the network output’s bounded, but whether that property is desirable on the problem at hand is a different question. We highlighted real-world examples, where this is the case, but ultimately, as we also wrote in Conclusions: “we do not claim that LN ‘solves’ the problem of extrapolation. Instead, we aimed to highlight the differences in extrapolatory behaviour of networks with and without LN, so that practitioners can make more informed architectural choices”.

Is it possible layer norm will shrink the eigenvalue

In certain situations, such as the case of very close training points very far from the origin, this is indeed possible. However, standard normalisation of the training set and appropriate architecture design limits the practical impact, and regularisation such as that suggested by “Observation Noise and Initialization in Wide Neural Networks” Calvo-Ordoñez et al. 2025 or even early stopping can arbitrarily lower-bound the eigenvalues altogether.

We take the opportunity to stress here that it is not the value of the bound itself, which is derived as an absolute worst-case from Gaussian Process theory, but its mere existence which is the remarkable finding of this work, along with the implications for the RKHS of such networks. As an illustration of this, consider the case of a one-dimensional Gaussian Process with 0 prior mean, a kernel defined as $k(x,x’)=1$, and with infinitesimal observation noise. Such a model will clearly have a posterior mean function equal to a constant at the mean of the training outputs, despite having infinitesimal kernel matrix eigenvalues and being subject to the same bound as the one provided in the paper.

Controlling the shrinking of the eigenvalues of the proposed model

To address this concern, we now provide mathematical analysis on the effect of including the layer norm. For the sake of brevity, we restrict this to the case of an $L$-layer network equipped with ReLU activations, biases initialised with variance $1$, and a single layer norm positioned at some layer $h_{LN}$, although this analysis can easily be extended to any situation considered in the paper. From the proof of Theorem 3.2, we see that the NTK of this network, $\Theta_{LN}(x,x’)$ differs from that of the equivalent network without a layer norm operation, $\Theta(x,x’)$ by a factor of 1 over the geometric mean of the variances of the pre-activations at layer $h_{LN}$, or $\frac{1}{\sqrt{(h_{LN}+\lVert x\rVert)(h_{LN}+\lVert x’\rVert)}}$. Thus, the Gram matrix of the NTK for the LN network, $K_{LN}$, relates to the Gram matrix of the no-LN network, $K$, by $K=DK_{LN}D$, where $D_{ij}=\frac{1}{\sqrt{h_{LN}+\lVert X_i\rVert}}\delta_{ij}$. Hence in the absolute worst case, the minimum eigenvalue of the LN network is scaled from the minimum eigenvalue of the no-LN network by a factor of $\frac{1}{h_{LN}+\max_{i}\lVert X_i\rVert}$. Under standard data normalisation practices, this is not likely to be an issue. That said, to alleviate any further concerns the reviewer may have we provide analyses of the two proposed regularisation methods to mitigate this issue. For the case of observation noise regularisation (an L2 regulariser from initialisation in weight space) of strength $\sigma_n^2$, the effective eigenvalues used by the theorem are given by $\lambda_{eff}=\lambda+\sigma_n^2$, trivially and arbitrarily lower-bounding the minimum eigenvalue in question. For the case of early stopping at time $T=n_{\text{iter}}\eta$ for some learning rate $\eta\approx0$, the effective eigenvalues are given by $\lambda_{eff}=\frac{\lambda}{1-\exp(-\lambda T)}$ (or an equivalent expression for discrete gradient descent (large $\eta$)), which lower bounds the minimum eigenvalue in question by $\frac{1}{T}$. Finally, we again direct the reviewer’s attention to the empirical results presented in the paper, which not only directly show that in practice this boundedness effect is always observed, but that the bound itself is very rarely reached due to its worst-case nature. We have thus demonstrated that the shrinking of $\lambda_{\min}$ can occur, but in practice this does not have implications for the extrapolatory properties of a layer norm network, even aside from the RKHS arguments, and can easily be mitigated (in a sense this is always the case as training is always for a finite time). We trust this expanded analysis is to the reviewer’s satisfaction.

I thank the authors for the additional explanations. However, if the minimum eigenvalue decreases inversely proportional to ||x||, won't it invalidate the whole boundedness result? Specifically, your finiteness bound is supposed to hold for arbitrary large ||x||. If after LN, the minimum eigenvalue becomes so small, I suppose that your bound will no longer be finite for arbitrary large ||x||?

We would like to thank the reviewer for reading our submission and providing feedback. We are glad to know the reviewer thinks our experiments effectively support the theoretical results. We address all questions/ concerns below.

Results expected given the homogeneous activation functions

As we show in Theorem 3.1, a network with homogenous activation functions, but without layer norm does not enjoy the stability guarantee. As such the homogeneity assumption by itself does not suffice to guarantee bounded outputs. As we then present in Theorem 3.3, it is the presence of at least one layer norm in the architecture that grants the stability property.

It is thus the presence of Layer Norm that guarantees the stability, not the homogeneity. We would like to thus politely disagree with the statement that the homogeneity makes the results expected. In fact, when it comes to our theoretical contributions, Reviewer QsAy described our results as “surprising”, Reviewer uiPW listed novelty as one of the strengths of our submission and Reviewer WTJo described our results as "interesting", while judging our assumptions to be “weak”.

The motivation for studying this problem is unclear

On any finite domain $D$, one can always bound the network outputs by $\max_{x \in D}|f(x)|$ and as such, basically any network is stable in this sense. What matters is how large is the bound $\max_{x \in D}|f(x)|$.

The result of Theorem 3.1 states that the supremum over an infinite domain is infinity, meaning that we can always find an input that causes an arbitrary large output. This means that if the inputs are restricted to some domain $D$, then as the domain size grows, the maximum output of a network with ReLU activation functions, $\max_{x \in D}|f(x)|$, also grows, by the same argument as used in the proof of Theorem 3.1. However, for a network with a LayerNorm, there is an upper bound that is unaffected by the domain size (and thus by the distance from training data).

A more important consideration here, however, is the prevalence of significantly out of distribution inputs at test-time etc. Standard normalisation practices assume the training distribution is representative of the true distribution encountered throughout the model’s use, however in practice this is well known to be untrue. This is precisely the notion of extrapolation, which is a problem of universally acknowledged importance. Furthermore, in high-dimensional problems it is well-known that almost all real inputs will be extrapolatory (for example, see “Learning in High Dimension Always Amounts to Extrapolation”, Balestriero et al. 2021), thus understanding the behaviour of models in this setting is of considerable importance. We direct the reviewer towards our experiments section for empirical demonstration of this.

Layer Norm not defined

The reviewer is absolutely right, that for completeness we should explicitly define the layer norm operation in the main body. We will add this definition in the camera ready revision.

Establishing the existence of the network in Proposition 3.4

This is a straightforward application of the well-known universal approximation theorem, which guarantees the existence of a network capable of learning any function. For a non-degenerate training set, it is clear that any function perfectly interpolating the training points will achieve the minimum MSE loss, and for a degenerate training set, i.e. one containing repeated training inputs with varying training outputs, there will be an optimal “virtual” point to be interpolated corresponding to the maximum likelihood estimator for the true function function value there under a homoscedastic Gaussian noise model. This establishes the existence of a function achieving minimum training error. By the universal approximation theorem, there exists a network capable of learning such a function, and thus achieving the minimum possible training error. It then follows that the addition of further parameters cannot lower this training loss, the rest of the proof follows.

We thank the reviewer for the clarification of their intuition on homogeneous activation functions. Without specific criticism to address beyond a disinterest in the paper, which we again note is an opinion held in isolation from the other reviewers, we instead offer the following further details:

Boundlessness of (n>0)-homogeneous networks

While it is indeed intuitive that a network equipped only with homogeneous activation functions will grow without bound, this is by no means guaranteed when disregarding training dynamics; indeed, it is trivial to construct a network by hand which achieves minimum training loss on any given finite dataset while extrapolating at a constant. The power of NTK theory, which we again stress is one of a very few respected theories explaining the behaviour of neural networks, is in allowing us to apply rigorous mathematical treatments to these networks, leading to results such as the ones presented in this paper which provide theoretical support to existing intuitions, and in many cases can provide new and useful ones.

As a final comment, focusing excessively on Theorem 3.1 somewhat misses the point of the paper, as while the theorem is indeed a novel extension of existing analyses of standard neural networks, its main purpose in this work is providing context for the later theorems and experiments.

Boundedness of 0-homogeneous activations

We agree with the reviewer that trained networks equipped with only (asymptotically) 0-homogeneous activation functions are trivially bounded for all but noiseless, exactly linear datasets. One can take this a step further and pass the network’s output through a bounded-range mapping, or even clip the network output directly. However, these measures are very rarely applied in modern deep learning (with the obvious exception of classification tasks), which overwhelmingly favours asymptotically positive 1-homogeneous activation functions.

What is, however, exceptionally pervasive in modern deep learning is the layer norm operation which features by default in every pytorch-implemented transformer, for example. It is hard to believe that any practitioner of such a network made the conscious decision to bound the extrapolation of their network; much more likely is the case that they wished to take advantage of the well-studied benefits to training stability (noting the very different notion of stability here). It should then be clear why understanding the implication of such an unassuming architecture choice is absolutely essential for informed application of such networks -- for instance, unstable extrapolation may accurately reflect the practitioner’s prior for the problem at hand? In that case, layer norm operations should be avoided, and yet without the results in this paper, there is no formal reason why this should be the case. Finally, we note that unlike the trivially bounding cases mentioned above, without the NTK-based analysis no such guarantee can be derived for the case of layer norm -- a notion we formalise in Proposition 3.4.

We trust this additional discussion will be of some aid in exposing the nuances of the paper and the significance of our theoretical results.

We would like to thank the reviewer for spending the time to read our submission. We are happy to know the reviewer finds our results strong and the experiments thorough. We address all questions/ concerns below.

Extending results to finite width NTK regime

As is standard within all NTK-based research, our main arguments make use of the central limit theorem, law of large numbers, and constancy of the empirical neural tangent kernel as a means of simplifying analysis to the point of tractability, which require the infinite-width setting for deterministic application. While some methods exist for studying the finite-width regime, these would considerably complicate an already substantial analysis, and thus a full theoretical treatment of this setting is beyond the scope of this work. We point the reviewer instead towards the empirical evidence for this behaviour in finite networks, again as is standard in all NTK-based research.

Layer Norm not defined

The reviewer is absolutely right, that for completeness we should explicitly define the layer norm operation in the main body. We will add this definition in the camera ready revision.