Just One Layer Norm Guarantees Stable Extrapolation
Authors: Juliusz Ziomek
George Whittle
Michael A Osborne
Venue: NeurIPS 2025 poster ()
Keywords: neural networks; layer normalisation; extrapolation; neural tangent kernel
Primary area: deep_learning

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.

Just One Layer Norm Guarantees Stable Extrapolation
Juliusz Ziomek*1, + George Whittle', Michael A. Osborne
Machine Learning Research Group, University of Oxford Equal Contribution Corresponding Author
{juliusz, george, mosb} @ robots.oz ox.ac........................... ac uk
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-ba 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 bounded-variane 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 produce pathologically large outputs for certain inputs. We support these theoretical findings with empirical experiments standard NNs often exhibit uncontrolled growth outside the training domain, a single LN layer effectively mitigates this instability. Finally, 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.
kernel. As a result, the output of an infinitely wide
winth networks, demonstrating that while
Introduction
Neural Networks (NN) have dominated the Machine Learning landscape for more than decade. Various deep architectures became the state-of-the-art approaches across countless applications. However, despite this empirical success, we still have a limited understanding of how these models learn and how they behave when making predictions.
What poses particular difficulty is analysing the network's prediction as inputs get further away from the training data. Clearly, a network achieving low training set eтror should output values close to the ground truth labels on the training points. Yet, it difficult to say how the network will behave outside of the training domain, as the training process does not necessarily constrains the network's s output there. As such, the extrapolatory behaviour of these networks very much undefined.
Lack of understanding of this extrapolatory behaviour can be dangerous. For example, a deep learning- based control system that applies voltages robot arm based on camera inputs can encounter pixel values with brightness exceeding those found in the training set. If this causes the network to output extreme values, this could easily lead a serious safety hazard in the system. As such, understanding how networks behave outside of the training domain and whether their predictions are stable out of distribution is critical.
One of the most complete theoretical approaches to studying the behaviour of Deep Neural Networks is the Neural Tangent Kernel (NTK) theory 19 23 Previous work in that stream 23 managed to
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

show that as the network width approaches infinity, the trained network's s output is equivalent to the posterior mean of Gaussian Process (GP) (29 with a specific kernel function that depends only on the network's architecture. While any real neural network has, of course, a finite width, in many cases, the predictions made by the NTK theory matched the empirical observations [9
13
In this work, we use NTK theory to analyse the effect on extrapolatory, or out-of-distribution (OOD), predictions of including Layer Normalisation (LN) (2 operations in network. We show that for an infinitely-wide, fully-connected standard network with positive nt-homogenoous nonlincarities, there always exists a fixed training dataset such that their predictions get pathologically large, i.e. "explode". On the other hand, we show that the NTK of infinitely-wide fully- fulle-connected architectures containing least one LN operation admits an upper bound for any imput, no matter how far from the training domain. Using this, we show that such networks trained until convergence have a bounded output even infinitely far from the training domain, for any training set. We then show how these properties are useful across different practical applications, including extrapolating residue size prediction to proteins larger than those seen during training and age prediction for unrepresented ethnicities absent from training data. We detail our contributions below:
We show with Theorem 3.] that for wnde fully-comnected positive re-homogeneous nonlinearities, but without any LN, the absolute value the expectation (with expectation taken over initialisation) can be arbitrarily high outside of the training domain.
We show with Theorem 3.3 that the inclusion of single LN anywhere in the network imposes an upper bound on the absolute value of expected output anywhere in the domain. That is, they enjoy stable extrapolation.
We verify these theoretical results number of toy problems and show that standard neural networks "explode" outside of the training domain, whereas the presence of at least one LN prevents this explosion from happening.
We show how these properties are useful when making duttt--distribution predictions, such as extrapolating residue size prediction to proteins larger than those seen during training and age prediction for unrepresented ethnicities absent from training data.
Furthermore, we extend our extrapolation stability results with Theorem C.3 to a more general class of nonlincarities and with Theorem C.I to, under certain circumstances, any architecture within the
Tensor Program framework [35 36 37
2 Preliminaries and Problem Setting
We assume we are solving supervised learning problem given a dataset Dimin = ((r,, y:)] where for all = 1,...,t, by a; Re we denote an nogddimensional input datapoint and by Vi Ih. corresponding to the target variable. We assume there are no repeated datapoints, that is, 3; 3 to train a NN So R. + R. parametrised by some 0 E Ω C Rd. We measure the success on this learning problem by the MSE loss function L(; Dunin)
j. We will denote by Xmain { and Vinin {1 Jn We use this data
1 []ff(())) y4 12 We assume the parameters of NN are initialised as described in Appendix A They then evolve via gradient descent update as Ht = f 1 observed that (even finite) NNs during training follow a kernel gradient descent with a kernel:
V L6; Durain) with learning rate 1) 0 at each step > 0. The work of [19
(xx, x':0) 0) = (v), Vofe(x²))
also known as the empirical NTK. They then show that for an MLP, as the width of the smallest layer approaches infinity, the NTK converges to a limit ((x, x²) that stays constant throughout the training process. Furthermore, the work of [23] showed that if the network is trained with learning rate small enough then in the limit co, the network converges to:
fox (x = ((x", Kurin (X Kisin
-Q foo + foo (x*),
where (xx) is the output of freshly initalised network with parameters 0 for input x.. As for the standard parameterisation gl (xx) = O, we thus also get that:
ff (x*)) O(x" Xuain ) (rraain Xerain) 2 Vrrain

---

which is precisely the posterior mean of a GP with kernel function ( æ²) conditioned on data Drain evaluated point x". While the original result of (23] assumed x < 1, this assumption can be lifted without affecting the limit, as we explain in Appendix B This result is valid, as long as the infinite width limit (a:, a²') exists. The work of [55 36 37 showed this limit exists for any combinations of basic operations (called Tensor Programs), which include Convolutions, Attention, and Layer Normalisations, to name a few. For completeness, we define the n-dimensional Layer Normalisation (LayerNorm) operation LN R" + R" below:
LN(z) - m
where μ A, E.= 2, and = and scalars and is executed for each element of z.
Ei= V
μ2² and the division and subtraction of vector
3 Theoretical Results
We now proceed to analyse how the inclusion of LNs in the network's architecture affect the NTK, and therefore the behaviour of the trained networks. We conduct our analysis under the assumption that the network's activation function (-)) satisfies the following assumptions:
Assumption 3.1. Activation functions element-wise.
R R are almost-everywhere differentiable and
Assumption 3.2. Activation functions o IR + R are positive n-homogeneous with >> 0.5, i.c. (Axx) (²)(z) AA
O.
Assumption 3.3. The minimum eigenvalue Amia of the Gram matrix of the NTK over the training dataset, Orain is strictly positive, i.c. Amia > 0.
Note that the former of the the activation function assumptions is a requirement of all activation functions, and the latter is true for the popular ReLU [12] activation function and its Leaky (25) and
Parametric [T3] variants, which are specifically positive 1-homogencous. This can easily be seen by noting that, for À > 0, ΦReLU (xx) max(0, Xr) = À max(0, x) = X (x)).(x)) The significance of n 0.5 becomes clear when considering an important consequence of Assumption [3.2] the activation
derivative d(.). which is positive n 1-homogeneous, is square integrable with respect to the standard normal measure, i.c. é(-) E LR, only if the monomial degree of the derivative close to is greater than -0.5, which is essential for NTK analysis.
While other popular activation functions, such as GELU and Swish 18 do not satisfy the latter criteria due to their behaviour for small inputs, one can see that for large inputs the behaviour of these nonlinearities approaches positive n- n-homogeneity. In Appendix C we extend our analysis to rigorously consider broader class of nonlinearities containing almost all popular activation functions 10. observing that our most important result holds under these relaxed assumptions too.
The assumption on the cigenvalues of the Gram matrix is standard in all NTK literature, and essentially ensures neither the dataset nor the kernel is degenerate.
3.1
Standard Neural Networks are unbounded
We first proceed by analysing the extrapolatory behaviour of standard networks without LN operations:
Theorem 3.1. Consider an infinitely-wide network fo(x)) with nonlinearities satisfying Assumption and Assumption 3.2] and fully-connected layers. Then, there exists a finite dataset D Korann Vinain) such that any such network trained upon D until convergence has
sup |E [f([)]]] =
where the expectation taken over initialisation.
'hen proving statements about infinite networks fa(Ø)), we refer to the limiting network. Eg. in this case the statement equivalent to suP E limm oo fo(x))]. where is the width of smallest layer.
3

---

Proof sketch. We show that such a network has an NTK which, for a large input to even one of its urguments, approximately positive n nomoogggesssss, where L is the number of layers in the network, and thus grows without bound when nontrivial. Furthermore, we show that this behaviour strictly positive, and thus indeed nontrivial, for any pair of inputs with positive cosine similarity. Using the Representer theorem to express the prediction of the trained network as weighted sum of NTK evaluations of the input with all training inputs, and by constructing the training dataset with nontrivial targets such that all pairs of training inputs have positive cosine similarity, we show that there exists a direction within the training set along which the magnitude the trained network prediction must be nontrivial and therefore grow without bound with input magnitude.
In the language of kernel methods, this is equivalent to showing that the reproducing kernel Hilbert space (RKHS) of the NTK, the span of which forms the set of all functions learnable by the network, contains only functions which grow without bound. For the training dataset described above, the function learned by the network does not enjoy perfect cancellation of these growths across training inputs. thus itself grows without bound. This result generalises existing theoretical results [34] (also built on NTK theory), and presents a clear problem: previously unseen inputs can generate pathologically large outputs, which could easily lead to a serious safety hazard in critical systems.
3.2 Inclusion of even a single Layer Norm bounds network predictions
By analysing the changes induced in the NTK, and therefore the predictions of the trained network, by the inclusion of an LN operation, we proceed to derive our main result:
Theorem 3.2. The NTK of a network fe( with nonlinearities satisfying Assumption 3.1 and Assumption 3.2 fully-connected layers, and at least one Layer Norm anywhere in the network, enjoys the property that:
VaER (x, C
for some constant C > 0.
Proof sketch. We prove that inclusion of Layer Norm at layer hEN of the network causes the variance NTK, i.e. O(x, x), to take the form of a ratio of two (strictly positive) functions of input norm, each leading order monomial of order ||> independent of the direction of thus for large input magnitudes approaches a finite limit, also independent input direction. Furthermore, we show that by the assumptions on the nonlinearities, both functions remain bounded for finite input magnitude. Thus, the variance NTK itself is bounded.
Equipped with this result, now proceed to study the extrapolatory behaviours of infinitely-widd networks with LN.. We have previously proven that for standard -nfie networks, one can always find a dataset causing arbitrarly large outputs. We now prove the opposite for network with Lh-that is, for any dataset and any test point the output of networks with LN remains bounded with the bound being finite for any finite dataset.
Theorem 3.3. Given an infinitelywwide network fo(Øz) satisfying Assumption 3.1 and Assumption 3.2 fully-comnected layers, and at least one layer-norm anywhere in the network, trained until convergence on any training data Dirain = (Krrains Viroin), we have that for any x E "
C Caaxyevva lyl |Drrinl
[B]f( S ((Drroin)
-V
Amii
where the expectation is taken over initialisation and Amii is the smallest eigenvalue of O and C the kernel-dependent constant from Theorem C.2 or Theorem 3.2
Proof sketch. Note that the GP posterior mean at a single test point x is a product of two vectors the absolute value of dot product of two
(zx, ( 1 (Xrnin Xumin) and y. By Cauchy- Clearly lyl S / /Draii max ly/ and vectors cannot be greater that the product of their norms. due to bounded variance property of the kernel we have (0(x, Xuain) ) I | Prain C, as explained in Proposition E. The proof is finished by observing that multiplication with '(Xusin Xusin) cannot increase the norm by more than
A which by Assumption 3.3 is finite.

---

Again in the language of kernel methods, the RKHS of the NTK. and therefore the set of functions learnable by such a network, contains only bounded functions. As the training targets are finite, by the self-regularising (with respect to the RKHS norm) property of kernel regressions, the trained network learns finite combination of these functions, and therefore must itself bounded. Note the bound is derived for the worst-case dataset and such might not be tight in the average case. However, we would like to emphasise again that it is remarkable such a bound exists at all, in contrast to Theorem 3.1 This is significant result and the first of its kind; by including even a single Layer Norm, one guarantees that the trained network's ' predictions remain bounded, providing a certificate safety for critical tasks. We add here that, while the above only considers the case of fully connected networks with nonlinearities satisfying Assumption 3.2 in Appendix C we extend our analysis to a broader class of architectures and nonlinearities, providing this same guarantee for almost all popular activation functions and arbitrary architectures the Tensor Program framework starting with a linear layer followed by Layer Norm operation.
3.3 Is considering learning dynamics necessary?
Within the analysis above, we analysed the behaviours of randomly initialised neural networks trained with gradient descent until convergence via the NTK theory. Given the promising properties of LN networks, one might reasonably ask if possible to arrive at such a bound without considering learning dynamies. However, now show that in this relaxed setting, one may choose network parameters optimal with respect to the training loss such that even with a Layer Norm, no such guarantee is possible:
Proposition 3.4. There exists trained network fx()) including Layer Norm layers such that
fo(x))
sup æER*O ,Ef
00,
where δ = arg mingeo (0, D Dmain) is the set of minimisers of the training loss.
Proof sketch. There exists a network fo( with fully-connected layers, operating in an over- parametrised regime such that additional linear layer parameters may not further decrease the training loss over some finite training set Dimin and sup=eR"" 00E0- [=)( [=] = oc, where Q is the set of minimisers of the training loss (for example, single linear layer network with non-trivial noiseless linear D.ain). One can construct network fz(x:) from fo(x)) containing standard Layer Norm operation, such that Va E Ra0, 30 E Ω s.l. fg(x)) fo(x) v0 Ω. As sa( already achieves the minimum training loss, it follows that the set of i exactly recovering fa (x) is subset of Ω v0 E Ω2*, any properties of trained (a( are also present in this subset of trained fz(x)).
Thus, S Reñ
fg())) co.
The above proves that, without explicitly considering training dynamics and initialisation, one cannot arrive such guarantees. This lends credence to the popular theory that the training process itself is of particular importance to network generalisation [5] 26 4 32), and demonstrates the power of NTK
theory in understanding the behaviour of Neural Networks.
4 Experiments
We now proceed to verify our theoretical results with empirical experiments, well as show how the insights we derived can be utilised in practical problem settings. For details about compute and the exact hyperparameter settings, Appendix We implemented all networks using PyTorch (27). We open- open-source our codebase
5

---

sine
quadratic
linear
20
15
0
10
-10-
-20
-10
10
20
-20
-10
0
10
20
-20
-10
10
20
X
Standard NN
LN after every
LN after 1st
LN after 2nd
Figure I: Predictions made by networks with various architectures when trained synthetic datasets. Red dot show train set datapoints. The solid lines indicate average values over 5 seeds and shaded areas are 95% confidence intervals of the mean estimator.
Standard NN
LN after every
600
15
15
0.45
0.40
5
X
300
5
55
0.30
200
-15
-15
0.25
100
0.20
25
-15
-5
5
15
-55
-15
5
15
X
X'
LN after 1st
LN after 2nd
0.45
0.50
15
0.40
15
0.45
0.35
0.40
5
5
0.30
5
0.30
0.25
-15
115
0.25
0.20
0.20
0.15
-15
-25
-15
15
X'
X'
Figure Heatmaps showing the values of empirical NTK values O(x, x²) plotted on domain æ, x E [-25, 25] with brighter colours indicating higher values. Note that the scales are different for each heatplot, with the values range for the NTK Standard NN being orders of magnitude higher than others. The displayed values are averages over 5 seeds.

---

4.1 Verifying Theoretical Results on Toy Problems
We first verify our theory on number of toy problems. We specifically choose three one-dimensional datasets, where the relationship between the target and input variables are sinusoid, a linear, and quadratic function respectively For each dataset, we fit neural networks with different architectures. We first adopt a baseline standard MLP with ReLU nonlinearities, consisting of two layers, each of size 128. We then compare three different variants including Layer Normalisation: Preceding all, the first, and the last hidden layers. We train each architecture to convergence on each dataset and plot the results in Figure | To account for randomness in initialisation, we repeat this experiment over five seeds and plot the mean predictions together with 95% confidence interval of the mean estimator.
The results clearly show that regardless of dataset or initialisation, the network without an LN extrapolates linearly and thus "explodes", verifying Theorem 3.T while the presence of even single LN operation anywhere within the network prevents that behaviour from occurring, verifying Theorem 3.3 In fact, what we typically see is that standard NN simply continues the local trend that appears near the boundaries of the training distribution, whereas the outputs of networks with at least one LN quickly saturate. Interestingly, varying the position of our number of LNs in the network has an insignificant effect on prediction. We show additional experiments in Figures and oj Appendix. These results show that the same effect is observed when changing the activation function to GELU or SiLU, and that BatchNorm does not prevent the explosion as LayerNorm does, which is expected under our theory as BatchNorm does not alter the NTK. Additionally, in Figure 7 we show that LayerNorm also prevents explosion in transformer architecture, thus verifying experimentally Theorem C.3 k
Our theory indicates the properties of the induced NTK should drastically change after the inclusion of even single LN operation throughout the network. To verify this claim, we show the empirical NTK at initialisation, that is (x) fo (æ'), as a heatplot. We note that as the network width approaches infinity, the empirical NTK should approach the limiting NTK, which we studied above. The displayed values are averages across all seeds.
We show the heatplot in Figure 2
By inspecting the plots in Figure 2 one immediately sees that there is completely different structure to the empirical NTK between Standard NN and the architectures with LN, where the latter exhibit a "chessboard" -like pattern, clearly closely related the cosine similarity of the inputs, which when low leads to minimal covariance. At the same time, this kernel value quickly saturates and reaches an upper bound that does not seem vary with the magnitude of inputs, again as expected. On the other hand, the pattern displayed by the empirical NTK of standard network is vastly different, with values steadily increasing with input's magnitude without saturation. This behaviour seems present even when the inputs have opposite signs, as indicated by the lighter corners of (25, 225) and 25, 25). As such, we again see form of explosion prevented by the presence of LN.
4.2 Extrapolating to bigger proteins
As the first real-world problem we study the UCI Physicochemical Properties of Protein Tertiary Structure dataset (28), which consists of 45,730 examples with nine physicochemical features extracted from amino acid sequences. The prediction target is the root mean square deviation (RMSD) of atomic distances in protein tertiary structures, reflecting the degree of structural similarity. As training set, we use randomly selected 90% of all proteins with surface area less than 20 thousand square angstroms. We then construct two validation sets, one in-domain with the remaining 10% of proteins with smaller surface area, and one out-of-domain with all proteins whose surface area is larger than 20 thousand square angstroms. We fit the models on the training set and evaluate them on both validation sets. We utilise the same network architectures as in the toy problem, that is, two hidden layers of size 128, with the following variants: no LN, LN preceding every hidden layer and LN preceding 1st and 2nd hidden layers only. Additionally, we compare with the well-know XGBoost [6] method. We show the numerical results in Table [
We observe that all methods perform similarly when evaluated in-distribution, but this is no longer (OOD) evaluation. Both the standard neural network and XGBoost
the case under out-of-distribunt perform substantially worse OOD and are outperformed by neural networks that include at least one layer normalisation. When inspecting the OOD coefficient of determination (R²2) across networks with different IN placements, we find that the position of the normalization layer does not appear to affect performance significantly--n all LN-equipped networks yield mean OOD R² values within each
7

---

other's confidence intervals. To complement the R² analysis, Figure presents the distribution of model predictions on OOD data for each architecture, as well as the relationship between predicted values and the total surface area of the proteins. 'he standard neural network without LN tends to predict unrealistically large residue sizes for large proteins, while XGBoost effectively extrapolates to constant value. Biophysically, larger proteins often exhibit higher RMSD values because longer chains and more complex folds are harder to model accurately. However, this relationship is not strictly linear, it depends how the protein folds and the degree of structural compactness. Consequently, while one may expect residue size to increase with protein size, a simple linear extrapolation without bounds is unlikely be accurate. Conversely, the constant extrapolation with XGBoost likely underestimates residue sizes, which should continue to grow beyond those observed in the training set. Neural networks that include at least one LN layer appear to strike more appropriate balance-contimeee continuing the upward trend without producing overly extreme predictions. Thus, incorporating LN improves the network's 's extrapolatory behaviour, leading to more stable and realistic OOD predictions for this problem.
Table I: In-distribution (ID) and out-of iit-o--diittibutionn (OOD) coefficients of determinstation R2 (higher the better) on the UCI Protein dataset. Displayed value average over 10 seeds with 95% confidence intervals (rounded to 2 decimal places). The best values in each row and overlapping methods are highlighted in bold.
Method
LN every
LN first
LN second
Standard NN
XGBoost
OOD R
0.50 0.02
0.50 + 0.02
0.50 = 0.01
0.31 0.02
0.27 = 0.04
ID R²
0.60 0.01
0.59 + 0.01
0.60 + 0.01
0.58 0.01
0.61 + 0.01
Average Model Prediction (LOWESS)
315
3.13
Density
3.10
3.00
3.22
3.00
Predicted residue size
Total protein surface
Standard NN
LN after every
LN after 1st
LN after 2nd
XGBoost
Figure 3: Histogram of predictions made by each model on the -iittr-disttitttionn data for protein experiment (left) and LOWESS 7 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%-cconfdence intervals.
4.3 Predicting Age from Facial Images of Underrepresented Ethnicities
As second roal-world experiment, target computer vision application, where the target is to predict a person's age from their facial image. We utilise the UTKFaaee datase(") a large-scale facial image dataset annotated with age, gender, and ethnicity labels. The dataset contains over 20,000 facial images spanning a wide range of ages from 0 to 116 years, captured in unconstrained conditions. Each image is cropped and aligned 200x200 pixels, and age labels are provided discrete integers. The ethnicity label has five unique values: White, Black, Asian, Indian and Others. To evaluate extrapolation capabilities,
construct the training set using only samples with ethnicity labels White, Black, Asian, or Indian. During evaluation, the an-distribtion set consists of validation images from these same four ethnicities, while the out-oo-distributiin set includes remaining validation set images.
htttps://suseee github. 10 /O/UUKFace/
8

---

We train a number of different networks, each one starting with a frozen ResNet-18 [17] as a feature extractor. After that, each architecture includes two hidden, fully-connected layers of size 128 with ReL.U non-linearities. Same as in the previous two experiments, we experiment with presence and position of L.N.. We train one network without LN, one network with LN after every fully-connected layer and with only one LN after stt and 2nd fully-connected layers respectively. We train each network until convergence on training set and then evaluate on in and out-of-distribution evaluation sets. We show the results in Table
We observe that while all architectures perform equally well on in in-distribution data, their perfor- mances vary significantly out- -ittiitributinn. The architecture performing best OOD contains an LN preceding only the first hidden layer, but its confidence intervals overlap with that of architecture with an LN preceding every hidden layer. Having an LN preceding only the second hidden layer performs worse, but still clearly outperforms having no LN in the network. In Figure f we show average model error for given OOD test input as a function of average cosine similarity of ResNet features of that input to training samples. We see that as similiarity goes to 0.0, meaning the testing input becomes dissimilar to the training data, the average error of all models grows, but this grow is fastest for the model with no LN, indicating that models with LN make more accurate extrapolation, implying the stable extrapolation property can also be useful in high-dimensional feature spaces.
Table 2:
In-distribution (ID) and the better) on the UTKFaae age prediction problem. Displayed value is average over 250 seeds with 95% confidence intervals (rounded to 2 decimal places). . The best values in each row and overlapping methods are highlighted in bold.
(OOD) coefficients of determinstation R (higher
Method
LN every
LN first
LN second
Standard NN
OOD R
0.42 + 0.05
0.47 + 0.03
0.36 + 0.04
0.27 + 0.03
ID R²
0.66 + 0.01
0.68 + 0.01
0.66 0.01
0.66 + 0.01
5 Related Work
Theoretical research in this area is limited, as is typically the case for neural network theory.
Xu et al. [34]]]]]]]] also utilise NTK theory to analyse extrapolatory behaviour for range of infinitely- wide architectures, but for our problem setting of fully connected networks they show only that networks with ReLUU activations extrapolate lin- early; a special case of Corollary E. Jused in the proof of Theorem 3.1 Hay et al. [14] provide extrapolation bounds when learning functions on manifolds, but this requires a specific train- ing method and prior knowledge of the function, rather than general property of networks as we
Randars Mk
Average Model (LOWESS)
every
19 aitor Est.
21
20
prove. Courtois et al. 9 discuss the general implications of the existence of NTK theory on network extrapolation, but contribute no bounds or quantitative results.
Average Cosine Similarity with Training Samples
Figure 4: LOWESS (7) 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.
Empirical research is more abundant; Kang et al. (20) empirically observe that in very high- dimensional input spaces trained networks tend to extrapolate at the optimal constant solution, that is, the constant solution which minimises error on the training dataset, including with Layer Norm. This relates to our Proposition E.2 and the proof of Corollary E.2 which suggest that for network with ReLU activation function and a very high number of dimensions, the NTK is approximately constant for relatively low input magnitudes, regions of which these dimensions may still have the majority of their volume described as out- of-distribution, and therefore extrapolatory. Additionally, in very high dimensions the criterion
9

---

the dataset to guarantee prediction instability occurs with very low probability, with most training input pairs having cosine similarity close to O. Vita et al. [3]] show that extrapolatory behaviour affected by the choice of optimiser for training. Fahlberg et al. [[]] investigate network extrapolation in protein engineering. finding inductive biases of networks drive extrapolatory behaviour, which is complimentary to our results.
Regarding work understanding Layer Norm in general, Xu et al. [33] found the effect of the Layer Norm operation on gradients during backpropagation to be more important than its effect on the forward pass, which matches our observations as these gradients are the mechanism by which Theorem 3.2 is possible. Zhu et al. [38] empirically found replacing Layer Norm operations with a scaled tanh function to be effective in large transformers; we speculate that this success comes about from enforeing the sigmoidal behaviour of the infinite width Layer Norm (obvious from our Equation 10 even at lower network widths.
Conclusions
Within this work, we have shown how the inclusion of even single LN operation throughout the network fundamentally changes its extrapolatory behaviour. We explained this phenomenon from the NTK point of view and verified our claims with a number of empirical studies, showing how this property may be useful in real-world scenarios. We believe the conclusions drawn from our work can be useful to practitioners. While LNs are already very popular, there are still many architectures that do not include it. However, if one wishes for the output to be bounded even outside of the training domain, one should consider utilising (at least one) LN somewhere within the network's architecture.
We note that extrapolation is, in general, an incredibly difficult problem. As we go outside of data distribution, what exactly consistutes 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. For example, there might be some datasets when we want to continue linear trend even as we go far away from training data. This pertains to the broader problem of selecting a kernel based on limited observations, such that the function of interest lies in the RKHS, which has been extensively studied in GP literature [24] 39 40]. Within this work, we do not claim that LN universally "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.
References
[1] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in neural information processing systems, 32, 2019.
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1rXiv:1607.06450, 2016.
[3] Felix Berkenkamp, Angela P Schoellig, and Andreas Krause. No-regret bayesian optimization with unknown hyperparameters. Journal of Machine Learning Research, 20(50): 2019.
[4] Satrajit Chatterjee. Coherent gradients: An approach to understanding generalization in gradient descent-based optimization. arXiv preprint arXiv:2002.10657, 2020.
[5] Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Chris- tian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019 8, 2019.
[6] Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd aem sigkdd international conference knowledge discovery and data mining, pages 785-794, 2016.
[7] William S Cleveland. Lowess:
A program for smoothing scatterplots by robust locally weighted regression. The American Statistician, 35(1):54, 1981.
[8] Vlad-Raul Constantinescu and lonel Popeseu. Approximation and interpolation of deep neural networks, 2024.
10

---

[9] Adrien Courtois, Jean-Michel Morel, and Pablo Arias. Can neural networks extrapolate? discussion of theorem by Pedro Domingos. Revista de Real Academia de Ciencias Exactas, Físicas y Naturales. Serie A. Matemáticus, 117(2):79, 2023.
[10] Shiv Ram Dubey, Satish Kumar Singh, and Bidyut Baran Chaudhuri. Activation functions in deep learning: A comprehensive survey and benchmark. Neurocomputing. 503:92-108, 2022.
[11] Chase R Freschlin, Sarah A Fahlberg, Pete Heinzelman, and Philip A Romero. Neural network extrapolation to distant regions of the protein fitness landscape. Nature Communccations. 15(1):6405, 2024.
[12] Kunihiko Fukushima. Cognitron: A self-organizing multilayered neural network. Biological cybernetics, 20(3):121-136, 1975.
[13] Eugene Golikov, Eduard Pokonechnyy, and Vladimir Korviakov. Neural tangent kernel: A survey. arXiv preprint arXiv:2208. /3614, 2022.
[14] Guy Hay and Nir Sharon. Function extrapolation with neural networks and its application for manifolds. arXiv preprint arXiv:2405. 10563, 2024.
[15] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, In Proceedings of the IEEE international conference on computer vision, pages 1026-1034, 2015.
[16] Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-1034, 2015.
[17] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[18] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415. 2016.
[19] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[20] Katie Kang. Amrith Setlur, Claire Tomlin, and Sergey Levine. Deep neural networks tend to extrapolate predictably. arXiv preprint arXiv:2310.00877 2023.
[21] Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint
14Xiv:1412.69880 2014.
[22] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:1711.0016 2017.
[23] Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
[24] James Lloyd, David Duvenaud, Roger Grosse, Joshua Tenenbaum, and Zoubin Ghahramani. Automatic construction and natural-language description of nonparametric regression models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 28, 2014.
[25] Andrew L. Maas. Rectifier nonlinearities improve neural network acoustic models. null, 2013.
[26] Vaishnavh Nagarajan and Zico Kolter. Generalization in deep networks: The role of distance from initialization. arXiv preprint arXiv: 2019.
[27] A Paszke. Pytorch: An imperative style, high- perforarform deep learning library. arXiv preprint arXiv:1912.017000 2019.
[28] Prashant Rana. Physicochemical Properties of Protein Tertiary Structure. UCI Machine Learning Repository, 2013. DOI:
[29] Carl Edward Rasmussen. Gaussian processes in machine learning. In Summer school on machine learning, pages 63-71. Springer, 2003.
[30] David E Rumelhart, Geoffrey E Hinton, and Ronald Williams. Learning representations by back-propagating errors. nature, (2366689) 1986.

---

[31] Joshua A Vita and Daniel Schwalbe-Kode Data efficiency and extrapolation trends in neural network interatomic potentials. Machine Learning: Science and Technology, 4(3):035031, 2023.
[32] Puyu Wang, Yunwen Lei, Di Wang, Yiming Ying, and Ding-Xuan Zhou. Generalization guarantees of gradient descent for multi-layer neural networks. arXiv preprint arXiv:2305. 2023.
[33] Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. Advances in neural information processing systems, 32, 2019.
[34] Keyulu Xu, Mozhi Zhang, Jingling Li, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv preprint MXiv:2009.118444 2020.
[35] Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes, Advances in Neural Information Processing Systems, 32, 2019.
[36] Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint ftfi:2000.145445 2020.
[37] Greg Yang. Tensor programs ii:: Neural matrix laws. arXiv preprint arXiv:2009.10663 2020.
[38] Jiachen Zhu, Xinlei Chen, Kaiming He, Yann LeCun, and Zhuang Liu. Transformers without normalization. arXiv preprint 25Xiv:2503.10622, 2025.
[39] Juliusz Ziomek. Masaki Adachi, and Michael A Osborne. Bayesian optimisation with un- known hyperparameters: regret bounds logarithmically closer to optimal. Advances in Neural Information Processing Systems, 37:86346-86377 2024.
[40] Juliusz Ziomek, Masaki Adachi, and Michael A Osborne. Time-varying gaussian process bandits with unknown prior. In The 28th International Conference Artificial Intelligence and Statistics, 2025.
A Parametrisation and Initialisation
Throughout this paper, the following parametrisations and, where applicable, initialisations for network components:
A.1 Linear Layers
Denoting a general layer input as (li-1) E 1R and output as 200) E 12,,, we parametrise linear layer as
2 1 g(i) + ozb())
(1)
Vns
with parameters
( E Ru ((0) R"
and hyperparameter
Th 0.
Note that d 3 is necessary ensure that the kernel is non-degenerate, which is requirement of Assumption 3.3 Furthermore, we initialise these parameters
12

---

1 E R. N (0,I) b (O,I),
Vj 1, . 14
(2)
(3)
where the row vector w is the jth row of W such that
w = (( w
(4)
A.2 Layer Norm
We consider Layer Norm operate without additional scaling parameters. Thus, denoting general layer input 2(6) IR and output g(0 R
(1) 3
p
Vj = 1.....7..
(5)
11
(6)
( 2
74
A.3 Activation Functions
consider activation functions, or nonlincarities, ó R - I, to act elementwise. Thus, denoting general layer input as z(i) E IR and output as g()) R",
5
Vea (2) 22j
Vj 1,..,,,,,,
(8)
where co is the activation function-dependane constant defined in Equation 19
B Limit for ||= 1| >
While Theorem 2.1 of 223 assumes ||æ||| < 1, we note this is only neccessary for establishing the inequality labelled as S85, which states that -1/2 ||a. || < Ki, where of is the post-activation of Ith layer with x0 being network input, n the size of smallest layer and K; is some constant. For any finite n the set of inputs for which n 1/2 ||>²) || < K1 some subset of the entire input space R". However, for any finite input all post-activation are also finite and thus taking >
max Ki
sufficies. As such, for an arbitrary input æ E _0 we can always find sufficiently
large network and even for points x l| the limiting result holds.
C Extended Results
C.1 The soft-cosine similarity
Definition C.1. The soft-cosine similarity between two vectors x; x'' E R parametrised by o² 2 is given by
x x + 0²
S (w. (e:,a';03)
(e x + ) xo. xo + 02))
13

---

Remark C.1. For o² = O, the soft-cosine similarity between two vectors 2, x' E R", S (x. reduces to the standard cosine similarity between a and x', cos (7 (xz x: x²)).
Proof. S (x, x0;
cos (< (x, x²)) by definition. 2 |
Remark C.2. The soft-cosine similarity between two vectors æt, x' E R", S (æx, x'; :²²). is bounded to the set 1-, and Vo² > 0, achieves the maximum value of 1 only for a = x'.
Proof. S o';² can be rewritten as the cosine similarity between the two augmented vectors ã = ( and n' ( .b) E R so is trivially bounded to the set [-1, 1]. The cosine similarity achieves the maximum value of 1 only for á = Aá' for some À > 0. When o² > 0. the matching last clement o² constrains À = 1, and thus æ = x'. Note that for the same reason, when 2² > 0 the minimum value of -1 is unattainable.
C.2 Asymptotically Positive n-homogeneous Functions
Definition C.2. A function p(-) R IR is asymptotically positive n-homogeneous
(Axx) ((x) Ve E R.
lim X
O,
An
where 8(-) F R is positive n-hhomogeneous, that is, Ō(Az) x²0(x) Vo R,
Furthermore, we assume throughout that ó(.) is nontrivial.
Corollary C.1. All positive n-homogeneous functions o R2 - R are also asymptotically positive n--homogeneous.
Proof.
@(Ax))
(x((z)
lim A
lim
A
X0000
A
((x).
By assumption of the theorem, ((x) is positive n--oomggenooss, thus 6(x) ) is asymptotically positive n-homogennous by definition.
Note that many popular activation functions 10 are asymptotically positive ne-homogenooss, which we explicitly show for a small selection:
Remark C.3. The ReLU activation function 112 (R) = max(0, æ). and its Leaky [25]
ΦLReLU (z) max(0.01æ, .)) and parametric (15) (x) max(aæ, 2), a [0, 1) variants
are asymptotically positive I-homogeneous (and indeed positive 1 l-homogencous).
Proof. To begin, note that both ReLU and Leaky ReLU are special cases of Parametric ReLU, with (x) ( 0) and (1))LU()) ÓPReLU 0.01), thus it suffices to show that these properties hold for Parametric ReLU. We proceed by showing that Parametric ReLU is positive n-homogencous:
ÓPRAL (Xx; doReLuuAA
= A max(ar, ..)
= A (eReLu
Va E
IR, A 0,
which shows positive 1-homogeneity by definition. Thus due to Corollary and Parametric ReLU are all asymptotically positive 1 1-homogeneous.
C.1 ReLUU, Leaky ReLU,
L
14

---

Remark C.4. The GELU and Swish (or SiLU) activation functions [18] given by OGELL (2) = (x),,,
where ((x) = Jxx V2= exp
(-z) 2 dt is the standard normal cumulative distribution function and
@swish (z; 5) = zo z (Bx), 3 2 O, where o(x) [- is the sigmoid function respectively, are
asymptotically positive 1 l-homogeneous.
Proof. To begin, note that both GELU and Swish/SiL.U take the form <(x) = aw((x) with
1 0
lim w(xxr=
and
2
0
I < 0
lim xw(x)
---00
where (1x) (xx) for GELU and w(x) o(Bx) I+exp(-Da) for Swish/SiLUU. It then follows that
((Ax)
lim (Ax(Ax)
lim
A
X
A
lim aw (Axr)
lim w(hx)
A
x 20
0 < max(0,
= Φx))
Vx E IR,
where the second transition follows from the dominance of (xx over in the negative limit and the boundedness of w()).
see that (x) is positive I-homogeneous. Hence by definition, both GELU and Swish/SiLU are asymptotically positive 1-homogencous.
From Remark C.3
Remark C.5. The tanh duant (Ir) tanh(x) and sigmoid dssgmood (2) = (Ix) 1-exp(-x) activation
functions 30 are asymptotically positive 0-homogencous.
Proof. Noting that (sigmoia (dran (2) ($) 1), it suffices show this for chunh (x) alone:
Prant (Ax)
lim A
lim tanh(Az) A
10
x
0
T 0
-1
<
Orant (x).
1
Az >
Granh (Az))
0
Ax
Az
1. o
x > 0
0 2
= 0 1 <1
X Sman (x)
Vr E T 0,
15

---

thus Granth (-) is positive O-homogeneous and therefore both Oranh (-) and @sigmoid are asymptotically positive 0-homogeneous by definition.
Note that beyond what explicitly show here, many other popular activation functions are also asymptotically positive n-homogeneouss.
Assumption C.I. Activation functions Φ IR + are asymptotically positive n-homogeneous with n 0.5, such that by Lemma E.1 the derivatives 6--) are asymptotically positive n 1-homogeneous with, and both o()), (()) E (², 7), i.e. they are square-integran with respect to the standard normal measure 7
Note that Assumption C.1 directly implies Assumption 3.2 but the same cannot be said about the inverse of this.
C.3 Extended worst-case bound on Layer Norm Networks
Theorem C.1. The NYK of an infinitelywiide network fo(x) starting with a linear layer follows by Layer Norm operation, depends only the soft-cosine similarity between the two inputs and xe' parametrised by of i.e.
Ve. xB O (x, (x, = (S (x(xx'':08))
Corollary C.2. The NTK of an infinitely-wide network starting with linear layer followed by layer- norm enjoys the property that:
VZERTO (ax,
for some constant C > O.
Proof. Due to Theorem C.I ( = θ (S (;);))))) However, due to Remarks C.1 and C.2
S(x, æ;; of))) Va "h o² 0. Hence, @(r, 6(1) C Vo E RR where C is positive constant depending only on the architecture, Th, and the initialisation of the rest of the network.
Theorem C.2. The NTK of network fo( with nonlinearities satisfying Assumption 3.1 and Assumption[C.] fully-comnected layers, and at least one layer-norm anywhere in the network, enjoys the property that:
YeERTo (x, C
for some constant > O.
Theorem C.3. Given an infinitely-wide network fo(æ)) starting with a linear layer followed by Layer Norm operation, with asymptotically positive n-homogeneous, nonlinearities satisfying Assumption Narr the condition of positive n--homogeneity, fully-connected layers, and at least one layer-norm anywhere in the network, trained until convergence on any training data D prain we have that for any x E "
C' Cmaayyeoo [
E [lff(()]] S B(Dirain)
Drrainls
Amin
where the expectation taken over initialisation and Amin the smallest eigenvalue of O and C is the kernel-dependent constant from Theorem C.2
D Proof of Theorem
Theorem C.I. The NTK of an infinitel-widde network fe(x) starting with a linear layer follows by Layer Norm operation, depends only the soft-cosine similarity between the two inputs and x! parametrised by oZ, i.e.
0x, = θ (2, 2 x':0)))
16

---

Proof. Let us denote the neural network as fo (01, 02...... ................................ ,L) to be collection of operations 01. 02 ........0 or parametrised by some 8 E 5 { 0r) such that:
fo(æ) (. (x)).
Let then, o₁ be a linear layer and 02 be the Layer Norm operator, as described by Parametrisation | respectively. As such 81 = (W E R Xng, bb())) E R and 82 0. Notice
and Parametrisation 5 that we can always represent the output of this neural network as fā()), where f = (og.-..,0)) OL.) and
(2 (a=) Let denote by (xx, æ)) the NTK of neural network fa evaluated at points and
x'. Using the decomposition in Equation 10 of [36,, we get that:
+ (x, x²) Σ (e.f(x)) Volfe(a()) Vo,fe())
(Vo. So, + Volr). Vo
[Vo, Vo + Vo V V (a (2'))
(a Jo (z))) (
where the treson line transition it true because the layer norm operation has no parameters and second-to-third is true because by construction fo(x) = fa (x). Let us denote x(1) 01 (x) = w + ob Due to Initialisation 2 and Initialisation 3 each 2) (for j l..... n2)) an
iid Gaussian variable with moments:
(1)
W (1)
x + os.b91)
B
EE
TTo
E (1)
x + E 6,80)
1/o
=0
Vj = 1, 11
(9)
Var (z(²) = Var
1
w x + )rb
Vno
Var (w/!i) a + Var (b())
70g
x x + of
Vj = 1, n1.
(10)
We can similarly denote (1) 01 (z'), which equivalently has elements iid with moments given as above. Each pair 2; (1) and 2 (for j 1,..., n1) are Gaussian distributed and are linearly dependent through the random variables (11) and b(1), so are joint-Gaussian distributed with covariance given by:
Cov (2 (1) (2 E
(1) (1)
E
1
w
(1)
x + oob (1)))
(W 2 tondg")
x TE
W(1) wr
x + oF p,
no
1 1i0
x "xx'' 8ij + o5o
x'
+ o3 8zj.
n2p
17

---

where the first line follows from the fact that both 2.4) and 2/ are zero-mean scalars, the third line follows from the independence of ((1) and W and the fact that W x x w and dig denotes the Kronecker delta.
We now consider the Layer Norm operation as described in Parametrisation providing
き - )=)(111))
(11)
where μ and as defined in Equation and Equation respectively are the maximum likelihood estimators for E
{3] V VVar (z)")) respectively.
(1) and
Due to asymptotic consistency of MLE, we have that with probability 1 1 -> [=]00] and a
Var x
n + co. As such, in the limit ni + oo, we have that for both and æ' with
V probability
E
~N(OII),
(12)
Var
(2)
E 3
(1)"
(O.
(13)
(2))
Var
The joint distribution over à and '' is then fully specified by the covariance between ã; and d/y
Cov
2,110 ;2] 1²2)²²
(1))
Cov (zi, xJ)
(1) Var 2x Var (2)
1 Tox x' + oz dis
x x + o?) ( Lo xr' o3
no
(a,, 2e; 8ij
where S (ax (e; is the soft-cosine similarity from Definition C.I this is now also exactly the correlation between á and .'
Note that due to normalisation,
Due to result of (36, we have that ( ø') converges to limit that is a function of inputs ă, ø'. However, note that o and ä' are now only related to æ and x' jointly through S (ax, xt;; af)) As such, Ô (#. á²' must converge to a limit dependent on æ and x' only through S (x, xx: ozf). To finish the proof, it is thus sufficient to show the same happens for the term (o(w), Vo, (x))). We have that
(LxT a' + +o?
(V, Vo, (x))) (V. (oz)) Sz (x'))
n1
aã az
0i' Oz'
(L +
v a
(.)
V (i))
V
( 11xTx x + oî) is due to the multipliers of the weight and bias parameters in the
where the term layer.
18

---

For the Layer Norm operator, we have:
dā
I
11
110 (z()) (2
111
0z(1)
n0
0'no
1
(1) (11)
11
Rg
1 0
-110))111)
03
Due to law of large numbers have lim +00 ng
1 (1) (1) ([x21)=(1)) (ar(2(2))
IVar (1)) Applying asymptotic consistency MIF again, we get:
Var (1)
aæ
(I 1 11 I
- 111T)
(2(1)
n2
Var 3/2
/Va (2)))
120
Var
Plugging this expression for both and xr' in to the full gradient formula:
Lx T Tax + oB
Ve.fo (x))) = nn
((=-11) 11
Vafo (i),
- (-11 V fo (i')
no
Var (2ftt)
Var (2(1))
To + ob
((=111)
1 11
Vala (x),
1 11 Far
no
n0
v
L No
o x + +o?)
(to of w + oz
((
((111T) Vafa (i),
11 no
1 11 no
V Vo fo (6')
s (1x,
We that the expression depends only on the gradient Vafa (ă) of the network fa (ø) w.r.t. input T and the soft-cosine similarity S (x,,²; but as we discussed before, ã and क' are independent of the inputs and x' except for jointly through S (ar, ar; ag). Thus, we conclude that the NTK for any network of this form can only depend on the soft-cosine similarity between the two inputs æ and x'. (x,x'; x'; oz).
E Auxillary Results
Proposition E.1. For a hounded-varianee kernel (i.e. VeE x) < for some C > 0), we have that the corresponding kernel feature map v6(:) (i.e. k(x,x')= = v(x) ((((()))) has the property of:
VEERTO [[((x)] S vc
Proof. Clearly for all E R
((x,x)= = = (x) [x(x)=)<>-))))))))))())))) (1)()) s.
Taking square root of both side of last equation completes the proof.
Lemma F.1. Consider an asymptotically positive n--oooggneoos function (with > 0.5) function 6: R satisfyyng Assumption 3.1 Then, the derivative of 6. ò(-) will also be asymptotically
positive (n - 1)- homooggeeos Moreover, if ((-) is positive n--omogeneous, then Φ(-) will be positive (n homoogggeos
19

---

Proof. We can decompose ó as o(.) é(-) + o(-), where δ(z) lim 000 6lz) = 0, ie o(.) has
((Az) is positive n
homogeneous, both o(-) and d(.) satisfy Assumption 3.1, and lim An
bidirectional asymptotic growth bounded by polynomial of order ñ < Moreover, ((x) ||=||" ((sig(())) where f(-) (-1,0,1) - is bounded.
rewrite
By linearity of differentiation, <(-) = x()) + ((-). sign(-) is constant almost-everywhere, and is multiplied by 0 where it is not, thus does not contribute to the derivative of (. Then, Ø(x) (n
1) ||=||"" sign(z) (z) ||~||~ fi (sign(x)), which is clearly positive (n - 1) h-mmmogennouus. As ob-) is almost-every differentiable by Assumption 3.1 <(-) exists and is defined on R E where F. is set of measure
Moreover, as (() has bidirectional asymptotic growth bounded by polynomial of order n, ó has bidirectional asymptotic growth bounded by polynomial of order ñ < n 1. Thus,
lim 6(Ax) lim lim e(Ax) X--1 x000 x=-1
)(va))
x0000 A
Ö(x)
where the second line follows from the fact that d(x) is positive (n - 1) J-homogeenoos and ñ. Φ has bidirectional asymptotic growth bounded by polynomial of order < n 1. Hence, o( ((Az) is asymptotically positive (n 1) homoogenoous..
Clearly, if ((-) is positive n n-homogencous, then o(-)= 2(-). Thus, é(-) = 2(-), ), so will be positive (n homoogggess.
Lemma E.2. Consider an asymptoticallly positive n-homogeneous (with n 0), nonlinear func- tion <(-) R + R E (R, where is the standard normal measure. Then, for zero- mean, joint-distributed scalar random variables with correlation ie (u, v) ~ (0, where
A a o.0up
J.0uP
0
(e [o(u)(()) = (a.0.)" x(p) + ((au,ov.P)
Bun:N(,o"))(aaa)) g
where co
-lima-oo
<((p) [-1,1] + [-1,1] satisfies x(1) 1, and
limo o'
Rla.'' .p) 0 Vo, > 0. 1-111.
(oo'))
Proof. First, let us change variables from (u, v) to (x, y). where (u, v) = (aww, owy), and thus (x, v) joint normally distributed with unit variance and covariance p. That is, (x,y) N'0, Aay).
1 p Thus.
where Ary
p 1
((E.)) [o(u)((u)] = E(aw) (o(a.=))(ovv))
As (6-) is asymptotically positive n-hommgeneouss above decompose & as
<(-) = 3(-) + 3(-).
(14)
where )(x)) limx ((Ar) A nn-homogeneous, and lim Loc
((A)) 0, ((.) has bidirectional
x
asymptotic growth bounded by polynomial of order n. Then,
20

---

((.,))((o.))
(ō(ant) + ö(one) ((lon))+(o0.))))
((ow))((o.)) + + townddo + ((aw)))(o.))
+((..))((o.)) (00 +o0 =
(a.0)) (0=))))(())(()) + R(z,y,o..
where
R(z, 0. 0b p) o')(()((()))) + o( + =((
(15)
clearly has asymptotic growth in daTo bounded by polynomial of order ñ < n, such that
((0,y;0,0') = Vx, E 20,eE-- (ool)"
lim
Thus,
(aa)) [((0))))((.1)] E( ((z))(() + R(a,
(ono,) (B(,s) (( Beews [(u,
= (ono,)"(((p) + R(o,,, o..
(16)
where
k(p) = E(2,,) (0(0)6(1)]
(17)
trivially depends only on p and
R( ov = E(
[(xx
(18)
We have that,
R(o,o'sp)
((
[
lim 400
lim
(0ol)"
0,0 +00 E(ay)
(oo')"
R(x, ((,y)
lim a. +000
(ao')"
where the second line follows from the dominated convergence theorem. We then have that
E, ~ N (0,
[6(u)
2 (())( + o(
lim 0-000
lim 0-330
2n
22n
0²n8(1)
1(0,0,1) 22n
lim
lim
I2n
=A(1).
21

---

Thus,
Ey ~ N (0,0²) [
1
CΦ =
lim 0-100
(19)
22n
R(1)
Finally, defining
x(p) R(P)
(20)
R(I)'
ii(ous o p) R(1)
Ov.P)
(21)
where k (-1,
(n,) [d(u)((v)] (ouoo)) x(p) + R(ou.ov.P)
EunN(o,r")) (o(w))) 0
-1 k(p) |-1,1] [-1,1] satisfies x(1) and
where co.
[imo
limoro +000 (oo')"
1( p) 0 E required, completing the proof.
Proposition E.2.
The NTK of neural network with L fulll--comectedd layers and nonlinearities satisfying Assumption 3.1 and Assumption C.I takes the following form:
n² R (x, + R (2x
(-11 ||ح||| ||22')
θ (x,
110
where R (x,xx) x²) [-C, R (xx = C Va, x' E Rn0 for some constant C' > 0, and R(x0,2)) Va, x' E R ie R(, x') has bidirectional asymptotic growth bounded by
1i a polynomial of order n. <
(XX')"
Proof. Recall from [22 that in the indinite-width limit, the pre-activations f(h (xx) every hidden
layer [4] have all their coordinates tending to i.i.d. centered Gaussian processes with covariance
( Rd x RR +
defined recursively as follows:
1
( (x, x²)
T x + o
x.
10
(b 1) x))
1) (a, x²)
A (h) (x, 2x)
E 2x2x2
(22)
11- (x' xx)
E(h 1) (x' xx)
(h) (x, x²) = (u.
[o(u)((v)] + oo.
where ca = (Eu>N(0.1) [o(u)))) Note that while [ use a-)) to denote nonlinearities, we use ((.)) to prevent ambiguity with standard deviations, which we denote with variants of a. Note also that co is arbitrary normalisation constant used to control the magnitude of the variance throughout the network, but is ill-defined in this form for non-exactly positive homogeneous activation functions (ie, all but powers of ReLU, Leaky ReLU, and Parametric ReLU), hence freely redefine this
Eus 2) [o(u)²]
to match the definition provided in Equation 19 namely co = limx-400
for
Xm
22

---

asymptotically positive n--homggeneous o(.). Note that this definition is equivalent to the previous one for positive n-homogeneous activation functions such as ReLU.
To give the formula of the NTK, we also define a derivative covariance:
( (x, x') CuBB(..)--W1O.MMM...... ((())(())
(x,
(23)
The final NTK expression for the fully- fully-connected neural network is I 36
L+1
+ x')
=
Σ
(or,
x²)
II M=h
SS (ar, x!)
h=1
(24)
(1.A))
where ()(++)(( (x,, (x,xx) x') for convenience and
L+1
((L.)) ()(N)) (x, x')
II
((k)
(x,x')
(25)
h'=h+1
the contribution of the /th layer to the NTK. Applying Lemma E.2 to the recursion for ( (x, x'),
Recursion 22 we get that
1( (x, (x, (a', aa)) 3 x (pit) (xxx²)) x'))
(26)
+ R ( (x, v EE(h) (x' x')), (() (0, + ob.
2
where
(h) (x, 1,
(h) (x, x²)
(27)
1 (x, x (h)) (x' x')
We therefore also get that
2((+1)(2 (x),)) " R (V>W),, v (20((x)))))))) of.
(28)
which follows from the fact that , (1)) (x, x) Vx R by definition and x(1) 1. Applying this recursion to ()) (x, x) no x of we inductively get that
n
(x))
T x
(1(N)( x))
1i0
(29)
|| 2nth
+
R( (xc,
no
where (b(b) (x, x²) is defined recursively by
23

---

kn
R(b+1) (x) =
Σ
n2
||=|||
R(h) (xx x') (m-
k
m2
k=
+
1
||
2nth
(2, 2r), / no
|1001 20hh + R(h) (x)), p (xx, at)
no
+ of
R(0) =
(30)
and satisfies lim 20 (Am,)'a)) O. To show this, we take an inductive approach. First, (Ax7)((13))
note that for the case O it trivially holds R(O) (ar, x²) of Vx, xe' RR"o, constant. Then assuming this holds for R(() (xx 2'). we have that
1 ||x0|| ( And (A (
R (Aze, N'xe) (Asgaattt)
lim -0000000
lim
no
+00
(AX))aitt))
R
V
1
||x
|m
+ RR( (Mxe Xx), V mg
|| 2nt (2ar)a
+ R(h) (Nz', Xxx'), (Acx, Xz'))
no
+
of
+
(Adya(att))
|20|| ||=e|] knd 22 N Nxx) ((-
(Ax
1
lim
(ad
-000
km
R
Anh
|
2m
1i0 (Mam Paah
(27m*
1 ||00||2" + Rin) (x',x'), p(h) (Xxx, Xa')
no
+
(A
And (||=||] (|''
An ROA) (XEE, (n-K)
lim
na
A.X
+00
k=0
R
nn
||ح|||
Umo |
.p(k) (Axx X Noo))
Vho
+
(AAr) malhta)
Rewh (Aaz, Xot')
(n-
11
||
[|=0||
lim A,00
k
no
Ana
k=
n
R
|
1
.limxx (N) (Axr, Xax'
vmg
vag
lim
(xx)
-0
where the transition from the 5th to the 7th line follows from the closely related property and the equivalent for X. the third-to-last line follows from direct applica- tion of the algebraic limit theorem and the fact that - k >> O, and the second-to-las line follows from
lim ( (A)
24

---

an aforementioned equivalent property of R(æ, a', ,)) the fact that p(Aæ, X²') E [-1, X, A > 0, and the fact that for An limx-+00 lim n² and the equivalent for X. Applying the above to Equation 26 we then get that
na=1 1 (200 + R(h+t)) (,)
((++1) (x, = n00
||a||| ||an1l]
Due to Lemma R.1
(6-) is asymptotically positive (n -homoggeneous. Thus, applying Lemma
E.2 to Equation 23
we get that
)1b+1) (2, = (E(h) (2, xx)x))(h))) (,', " i (ptN) xx,
(31)
+ R (VEW 2)), v E(h) (o' x') p(h) (z. xe)z
where k(-) and R(.,'' p) are defined as the analogues to ne(-) and R(,, p) for 2(.);
8(p) A(P)
(32)
R(1)'
R (.
R(
(33)
R(1)
where h(p) and (.)'.)) are defined by application of Equation 17 and Equation respectively to d(-). Note the different scaling (for example, A(1)
A(1)
f in general where R(-) is the analogue 2(1)
to i(-)). More specifically, we then have that i(-) E for ()++1 (x, x'), we get that
R(1)
A(1)
Inserting the above expression
x(1) x(1]
n ||=||||=²]
(1h) (x, x²)
(1_ 73g ||=||| ||~'|]
2 (p (cx, + R(u+1) (x,
(34)
which follows from the fact that d(-) is asymptotically positive (n h)-homogeneous and 1(++1) (x.x") x²) satisfies equivalent properties to ()++1) (x, x'), namely limx h
8th 11
R0++9()mAm)) 0 Va, x' p which can be shown by similar
O Va, x' go and x (Ax)="11- derivation.
Inserting Equation 29 and Equation 34 into Equation 25 the contribution of the /th layer of the NTK.
25

---

L+1
(LL,h) (e,,2') =x(hk) (ar, x²)
II
INND (3))
W=++1
n²
||>|||
||
(pln-1) (ce, a')) + RR(h) (xx, 22)
no
L. II #==++1
no -nhi-s
1 |,0|| |
i (pli'-1) (xo, x²)) + Reth") (ce, at))
- =0+1 nn -n
1 ||a|||
||a2"|
K (pll-1) (20, x2))
II
(pli'-1) (co, (o)
730
h-=h+1
+ (((h))
(ax, (ax,
n R((h) (re, x²) + R(()) xx)
1 |||||| ||
(35)
730
implicitly accounting for the exceptions for E((++) (xr, x') = 1 and K <((-1) (x. x')) S (x,,r'; oz) by abuse of notation, and where R((h) (ar, æ!) accounts for all remaining (lower-order) terms, thus satisfying limA,A'-200 (Am, Var xf IR h 0,..., L with equivalent derivation to the above, and
(Ax7)*
n (pth) (xxx²)) x')) -h & (p (2, x')
O.... (36)
(()) (x, (²,x')= =
R(LL) (x, = k (p(l) (xr, x'))
h=L
noting that ((a) (z, x') E
|-
(())) ((())
((1))
I.-h] A(1)
wh o..... L. Finally, applying this to
Equation by summing over all layers we get
(vx x²)
Σ (I.hi)
(0,,z')
h=0
L
n² R(th) (xx x') + R(h) (x, x')
||
||221]]
120
h=0
1
||ح|||
||x'll]
R (x, + R xx))
(37)
where
(xxx²)
Σ
R(() (x, x'),
(38)
h=0
((1)) noting that x (x, E [-C, C] where C= Em==
and
R(I))
2(1) R(1)
L.
R (ax, x²) =
Σ Rhh) (x,x')
(39)
h=0
R(Ax.\=") (
trivially satisfying X individually.
0 as this is satisfied by all components of the finite sum
26

---

Definition E.1. A function of two variables Rd x Rd - R is doubly asymptotically positive n-homogeneous if
+ (Axr Xa') (ax,)')) â (xr,a²) Var ae E Rd (AV)"
lim 00000000
where (x,x") is positive n-homogeneous in both arguments, that is, é (Xxx, Xv (xxx')) x²) Va, x' E Rd >
Corollary E.I. The NTK of a neural network with fully-connected layers and nonlinearities satisfying Assumption 3.1 and Assumption C.I. itself doubly asymptotically positive n n4-homogeneous. That
is, Vxx x' E 1
+ (Ax, (xx,) Ô (xr,z") Va, x' Ro
lim X0
(AV)
where ② (ax x²) is positive n n homogeneuus with respect to both arguments, that is,
3 (Az, x²) = - Ô (x, )x')xx))=))) = xn à (x,xx) Va, x' E R A o, and
Ô (ax, Xxx) (22)** (xx Va, x' R" X, 0.
Proof. From Proposition E.2 we have that
(x'z1)) n²
θ (xx, (xx)'x))
I
R (Axe, Nax:o))
R (Axx Nz'; (AXr)at
lim
lim --00
no
A,X
(Aagne (Axrgn
(1) ||=||| |1211]
lim x
(Mar, Na':
n2p
1 || ||a'||
( lim (Nee, Na';
Ve,, a' E Ro
no
$00
We thus define
(x,x') lim O (x,
(40)
(f
Now, note from Equation and 36 Equatior38 that R (æ, x²) is composed of summation of products of compositions x(.), and k(.). which both only act on pre-activation correlations, which in turn are given in Equation 27 by (h) (xx, x') ((h) (x)) (æ,x')
(3 In the limiting case,
27

---

(1) (Ace, Xax)
lim
lim 2.
200
v V2 (Ax, Xx) )2))) (Nw', Nxx'
N. ||az|l| (1a'1)) "*
(Maa,, + R( X'x'
lim
A,0
| 2n
+ (R(h) (Ax, Xae) (( || + +R(h) (Xox Xxx)
No
na
(A2r)*
(1|2211 || (1)" (pln-1)) Na'1) + (Xar) Xxx
= lim A,
no
na (Ax)) (
(1eel |an + ngnn (an
|1a0| ng Ria (2 Naw (aa')) (Arranh
no
22mh
(1|zo|) ||00"1))"
(0th-2)) (Ace, Xxx'))
ReW (1ce, Noo
lim A,00000
V ([|20|| || 2)
(2a)))
([|a)) || 1)
( lim (,,0-0000
(he,
([|a=|| ||=0 ())
n lim plt (Xx, Xax))
0,000000
which follows similar reasoning that used in the proof of Proposition 1.2 with the final transition assuming x(() is continuous on [-1, 1] and the limit for p(h-1) exists. Now, the initial limit correlation is given by
lim /00) (Ax, X'x')= -100
lim S (Axx, Nz' (0?)
-
XX
1 nl
x:
x.' + oz
lim
A.X
1 || Aax |
of
1 |a''' 1|2 + o3
AXxx x
oo?
lim
X.X XX
|2: |<0||
XX
x
a a' Va, E "o
where è and ø' are unit vectors the directions of and x' respectively. Note that when xx this expression is instead given by
of ( (o 11.
(0) (0.0)
In either case, the limit exists and depends only on 2 and á' (which if are equal to 0 indicate the adoption of the alternative case). Hence by induction, + (()) (æ,, x²) exists vh O,..., L. and depends only on â and C'. To solidify notation, we define
B00) (#, (e, lim (()) (XxxXXa')
(41)
2.0-000
3(0+1) (f. x6) (7500 (a. at))
(42)
(00) (0,00) [ ④ a'
a, +
2 =
28

---

where è and a' are unit vectors in the direction of x and æ' respectively. By identical logic, the same can be said about limx.X' & (p(h-1) X²x)). As both & and & have range S [-1, 1], this limiting behaviour carries over to
(át, #0)
lim T (Aar, Xxx)). +000
(43)
Thus,
Ô (Axx \'2' = (-- | [|x'a0)) IXalll n²
(Xic, Nax)
= =AXyx* (1) ||||| (a'll) n (((, k')
no
(Ax))* â (x,x') Va, R,',,''>>
where the λx and Xxx' are abuses of notation referencing unit vectors in the direction of xæ and X'x' respectively, and the penultimate transition follows from noting that for nonzero a: and X, Xa and X/xe do not depend on À X for A, X > O, and in the case where either are zero the change in behaviour is hidden by multiplication by zero. Hence, é (z. x') is positive n n'-mooogensous with respect its first argument by definition, and also with respect the second argument by symmetry. Therefore O (ar æ') is doubly asymptotically positive n homogeeneuuus by definition.
Corollary E.2. The NTK of neural network with fully-lly-ccoeced layers and positive ne-homogeneous nonlinearities satisfying Assumption 3.1 is üsself asymptotically positive n- a--mooogeeoous with respect both arguments. That is, Vx, Ro
(Aae x²)
+ (x', Aar) An
lim
lim
= (x, Va, x' E Ro
An
where è (xx x²) is positive n homogeneous with respect its first argument, that is,
(Ax, x²) x Va, x' Ro
Proof. Consider Lemma E.2 for the case of positive n-homogeneous Φ(.). Clearly, with reference to
Equation 14 the non-homogeneous part (xx)) = 0 Va: R, and as such (bx) o(x) Vo e IR. It then follows from application of these conditions to Equation 15 that
R(z, yi ou, ob, p) = Yx,y E >.,
(44)
and thus from Equation Equation T6 and Equation 21
(o p..p) =
Voa. o. > 0,p [-1,1]
(45)
B(xy)) (6(..))((..)) (o.
Vo,, Tv. 0.p p [-1,1]
(46)
R (au,
Voa, T. 0,p [-1,1]
(47)
We now consider Proposition B.2 for the case of positive n--homggeneous d(.). Recall the recursion
for (h) (x, æ'), Recursion 30
29

---

n-1
Ann
R
(1++1) (x,x')
Σ
1
||0||| |ax'1)
R(h) (a, x') (n-
no
h=0
+ R
+
(X xx),
V
1
| 2nh R(h)) (x', xl).pph) (z, (2, + oz
720
no
11
71
x
|2'1)
Ann"
Re(h) (a), .o' ( + ob.
k=0
R
=
Under these conditions,
R(b+l) (Axx, x²)
k
(
||
(2'))
R(h) (har, 22)(1-b)) + o
lim x
lim X
Anhel
Anhin
n
kn
RC(h) (Nar, (nn)(1
lim
10
||
||
k
720
A (n-A)
kn
((h) (Az,, (n-
1
|||||
||=²|
lim
Ant
120
Ann
R(()) (Jar, 20) (n-1)
1
||a|||
||>21|
lim A-+00
Azh
no
where the final transition follows from direct application of the algebraic limit theorem. Considering the case of n = O;
R(0) (Ax,
of
lim X
lim X
And
À
= Va, x' E R"
hence by induction;
R(h) (Ax, A
lim
=0
Va, x' E X = O....L L.
By identical derivation,
Re) (Aæx, x²)
lim A
Ya, E R = O....L.
Anh p
Again considering the contribution of the /thh layer of the NTK and collecting lower order terms into R(h) (xr, æ²), it is clear that
Rh (Xxx x') =0 Va, 20 R"0 h O,...,L. L. And
lim
Note that this a stronger statement than the equivalent in the proof of Proposition E.2 and owes to the fact that under these stronger assumptions R(h) vanishes under scaling one argument only, rather than requiring both as was previously the case. It follows that
30

---

R (2ar, x²:
Erzo Reth) (har,
lim 2
lim
Art
X
Ant
R((h) (Aar, xt)
lim A-200
Art
hm0
= 0 Var, x' E R
Following now the proof of Corollary E.1
no || ||a2)) nt (Aar, of)
θ (Axx x²)
R(h) (Axx, x')
lim X0000
lim
Anh
Anc
Ane
Arh (1
lim
||a0|| [Iae"))
(Nar, aa; ab)
X Art no
= -(1- ||a||| |arll) x²
1
lim in (Aar, (
(Jime
For convenience, we define the single-argument limit NTK
+ (Axz, x²)
- (xx
lim 20000
(48)
And
As before with reference to Equation 27 considering the single- angluent limiting case for the layer-h correlation,
(b)) (Aze, x²)
lim (xxe, x²)
lim
X0000
x000
(h) (ax, )xe) (1))) (xt, x')
1. ||zi|| [|22=1) nit
x (p(a- (Axx, x))) + R(()) (Xx,
lim
no
X=000
(
2ndh
1 ||
+ (Xce) (xee, )x0 200) (x', x²)
no
(+)) x² (||=|| ||2"11)"** (p(n-1) Dar, x²') + R (Axr x²)
Ah
lim X
Ant (1to na V (11001200 + nggh ne0 dre)
I(M) (x', x²)
Aan
([|2|l ||=0") ||=0"1))" r (pll-12) (120,)))20,00 xd))
Re0 (x))
lim
X
1|20 EXN) (xo' xx)
Ama
(1 [[=e"1)"""
lim v (p (xze, x²)) 20000
||=0
EE() (ac" x0))
lim ( 20))
VEUN) (xx' x²))
which follows similar reasoning to that used in the proof of Proposition E.2 with the final transition
assuming x(-) is continuous on [-1, 1] and the limit for plh-1) exists as before. Now, the initial limit correlation is given by
31

---

lim (1x x²) = lim S (Ax, o':
A0000
A-100
x
x + o
lim
A-000
((2o 1 || 1|2 + oo?)
1 |0|||
+ of
no
xo T
noog
lim
+
X
||a||| v || + +noog
A ||:0| / ||00 12 + no
T x
Vo E R
0, 20,z' R
| + nooy
where @ is unit vector in the direction of æ. Note that when æ 0, this expression is instead given by
noo
(0,
V (noof) (1x01) + mooz)
σb
Va' E R",
| 2
+ of
mg
In either case, the limit exists and depends only & and w' (the former of which if equal to 0 indicate the adoption of the alternative case). Hence by induction, lim p(h) (Aæ, æ²) exists
vh = O.....L and depends only on ø and x'. By identical logic, the same can be said about 1im is (p(n-1) (Axx 2').. To solidify notation, we then define
(16) (t, x') = lim p!h) (Axxx'),
(49)
2
(1|=011) (1)" n"
1lb+1) (át x²)
(p(h) (é, x²))
(1)) (xt', xx)
Teo'
1 E Rno 0
(50)
(00) (#, x²)
||a 1 ++ooo
|| 1
ND
As before, since both and A have range C [-1, 1J. this limiting behaviour carries over to
& x') lim (Axe, x!).
(51)
Thus,
n² (1 > ||= |1221||
è (Axr,
& (Xa)
no
n
xn (1
||ج||| [|
& 20;
no
xn θ (ax xx) Vac, x
Ro x 0,
where Xx is an abuse of notation representing unit vector in the direction of λæ, and the penultimate
transition follows from noting that λæ does not depend on X for all 18 X > 0, and in the
32

---

case where either are zero the change in behaviour is hidden by multiplication by zero. Hence, Ô (x' is positive n p'-moomgensous with respect to its first argument by definition (and also with respect to the second argument by symmetry). Therefore, θ (æ+ x²; is asymptotically positive hnmooongess in its first argument (and indeed the second) by definition.
Lemma E.3. The probabilists' Hermite polynomials satisfy
(x) [Hn( (v)] = 5mmm:
where
(x, (α, 1) ~ N (0, [
f))
Proof. We begin by noting that the differential joint density function of the given a and y is
x. + y² 2pxy
dxx,((,,))
exp
2x 2T 12
2(1 p2)
l n!
Σ
(
where y is the standard normal measure and the second line is direct statement of Mehler's formula. Hence,
/
B(a,))
f Hn(z)
L
R2
11 JR2
[ (() /
!! J
H.
π. (nlōmt) (m!5mm)
p" m6
where the second transition follows from the dominated convergence theorem, the penultimate transition follows from the orthogonality of the probabilist's 's Hermite polynomials with respect to the standard normal measure 7,
Lemma E.4. For all nontinearities Φ R. - R satisfying Assumption 3.1 and Assumption C.1 the associated k(p) and k(p)) defined as in Lemmal E.2 exist and satisfy
<(p) > 0 Vp > 0, and
x(p) > 0 Vp
Proof. First, recall the definition of A(p):
33

---

R(p) Blass) ( where
(0, N (o.
1 1 P 9 1
where ((-) is the positive nz-homogeneous part of (0-), i.e. Φ(x) = limx d(A Ve E I.
Observe that R(-) E (R, 7), where is the standard normal measure, which follows trivially from the fact that o(-) has polynomial bidirectional asymptotic growth (specifically monomial growth of order ri). Indeed, this is requirement for the mere existence of x(())
The probabilists' Hermite polynomials Hn(-) are well-known to form an orthogonal basis in T,2 (TR, 7), with
J Hm))))m()))d((( = n16mm
where 6mm is the Kronecker delta, nonzero only for Thus, write
x(-) an Hn()), where n=0
On E E [0()))n((0]
Hence,
k(p) = (ux) [
am Hm (w)) ]
= (a.v
Σ m==
On Hn(())
Σ M==
agam ((a
[H (m Hml
n.m=0
ag mmpmm 8pom
n.m=0
Σ nlp"
where
n=0
0.
1 P
1
9).
(I,w) ~ N
where the second transition follows from the dominated convergence theorem and the third transition follows from direct application of Lemma
definition of asymptotic positive re-homogeneity, &(-) is nontrivial: hence there must exist at least one a, f 0, so the corresponding term n!!p is strictly positive for all > 0. Thus noting that 1 > 0,
By assumption
k(p) > 0 Vp > 0.
Noting that under Assumption C.I d(-) L (R, 7). so the above also holds for A(p).
Finally, noting that as 1 > 0, k(1) > 0, and recalling k(p) x(U
k(p) and k(p) k(p)
K(T)
34

---

<(p) > 0 vp > 0, and
k(p) 0 Vp
Corollary E.3. For all neural network fullycfullyceeeeeeteed layers and nonlinearities satisfying Assumption and Assumption 3.2
(t,
lim R. (Ax, x²) 0 Va, x E
R
22 x >> 0,
where & (x, x') is defined as in Equation of Proposition E.2
Proof. We begin by showing that the limit correlation (át, x²) limx (0) (Aær, x²) Yx, E R x Tx' > 0, = O...., L. First, recall the recursion for the single-argument limit correlations derived in Corollary E.1 Recursion 50
(1<=11)"" /(i) (x',
0 (b++1)
(0,x')
K (pth) (tc (i, x')
# TTx
(00) (Ø,
|| 1
On
x 0
Due to Lemma E.4 and the strict positivity of Σ(M) (x', as it a valid kernel and x' > by assumption of the theorem, if ((h) (#, x') > then it follows that p(++1) (#, > as the nonlinearity 66-) satisfies the conditions for this Lemma by assumption of the theorem. Note also that by assumption of the theorem ((h) (ft, x²) > 0, hence by induction
pth) (#, x')
Va, x IR" TT 0, 0, O,...,,. L.
(52)
Again due to Lemma E.4
(pe0) (#, x0)
>
Va, E Ro TT > 0, h O...., L,
(53)
(pth) (xc, xe))
>
Va, a E no x. Ta' > 0, h 0...., L.
(54)
To conclude the proof, recall from Equation 38 (it, x²) lim n Aar, x' is composed of summation of products of R(h) (œ, x²) and ( (æ, x') for various h .,... L. all of which are existing limits and are strictly positive. Thus concluding the proof,
& (0. x²) lim R (Aæ, x') > 0
Va, x' E n x T
xT
(55)
X-200
Proof of Theorem C.2 and Theorem 3.2
Theorem C.2. The NTK of network fo( with nonlinearities satisfying Assumption 3.1 and Assumption C.T] fully-connected layers, and at least one layer-norm anywhere in the network, enjoys the property that:
VeER 10 (z, x) < C
for some constant C 0.
35

---

Proof. We first consider the effect of placing a first layer-norm operation between the linear com- ponent and nonlinearity of arbitrary layer hn,, with 0 I ,n I L of such network. The first observation is that this destroys the variance for both inputs and x' subsequent to this layer, while maintaining correlation. Thus, the activations and therefore )(h x") and 5)() (ac x') can have no dependence on ||a=|| and |a='] for all hLN I h I L.. Not accounting for the effect of the layer-norm on the backwards pass, following the reasoning of the proof of Corollary E.I, the contribution of layer h, where S AEN to the NTK would be asymptotically positive n nLw -omogeneous with respect to both æ and æ'. That is, at leading order they are polynomials of order 1 in both æ and a'. Similarly, the contributions of layers h will depend only on the correlation of the activations layer hLN.
We must now consider the effect of the layer-norm operation on the backwards pass. From the proof of Theorem C.2 we see that the layer-norm operation essentially divides the backwards pass by the square root of the variance directly preceding the operation, which in this case is given by the square roots of g(hLx) (æz æ)) and (hLN)) (x', x²) for inputs and x' respectively, both of which are asymptotically positive n ALN -homogeneous, i.e. at leading order a polynomial of order AL..
Accounting for this effect on backpropagation, to the NTK is now given by the ratio of two asymptotically positive n -homogeneus terms in both æ and x'. Thus, in the limit ||=]|, ll, ||=''| - so, the contribution to the NTK must approach finite limit dependent only on the layer correlations between the two inputs.
that for all layers & < hLN, the contribution
For the case x = x', all layer correlations are 1, thus in the limit ||>|| + co, O(x, x) approaches positive constant independent
the direction of æ. By the assumptions of the Theorem, specifically that the nonlinearities are locally-bounedd, almost-everywherr differentiable, and of locally-bounded derivative (which apply to all commonly activation functions), (ar, æ:) also bounded by a finite positive constant for all ||> || not inducing the homogeneous regime. Thus, the NTK of such network satisfies VZER ( (x, x)) || s for some finite constant C' 0. Hence, the proof is concluded.
Theorem 3.2. The NTK of a network Sa(æ) with nonlineurities satisfying Assumption 3.1 and Assumption 3.2 fully-connected layers, and least one Layer Norm anywhere in the network, enjoys the property that:
V 10 (x, C
for some constant C > 0.
Proof. By Corollary C.1 positive n-homogeneous nonlinearities satisfying Assumption 3.1 also satisfy the conditions for Theorem C.2 hence its result also applies here.
G Proof of Theorem 3.3
Theorem 3.3. Given infinitelywwide network fa satisfying Assumption 3.1 and Assumption 3.2 fully- fully-comnected layers, and at least one layer-norm anywhere in the network, trained until convergence any training data Dirain (Niraiay Virain,, have that for any E "o:
C Cmmayyeaaaa Ivl
[[]o((=)]] S B(Dirain)
--
Amin
where the expectation is taken over initialisation and Amin is the smallest eigenvalue of O and C is the kernel-dependent constant from Theorem C.2 or Theorem 3.2
Proof. For an inifinitely wide NN after convergence we have:
fa ((x, Xisin 1 No(x=) + No(x),
where (x, Xuau)) )); = (ax
az x.)) and (Quan)ss = ( 2,) and (y); = y for i = 1...., | and No( is the output of freshly initialised network with property E B[No(=)) = 0. Let denote by (b()) the feature map of network's 's NTK, such that ( x') (xx)) (x)) and let Amin be the
36

---

smallest eigenvalue of the Quain matrix. In expectation we thus have:
[ = (xx, Aurin ) mminn
O (x, Virain ) ll |1vll
(x, X ll2
Diriil max lyv Vi
Amin
Σ
((ar))((ar)) Duriin
max [y]
Ammin
V
y
ar E'E
Dimin C
Denin
max | ye
Amin
C maaxyeemm |
Derain
Amin
where the penultimate transition is due Proposition E.1
H Proof of Theorem 3.1
Theorem 3.1. Consider an infinitely-wide network fe with nonlinearities satisfying Assumption
3.1 and Assumption 3.2 and fully-connected layers. Then, there exists finite dataset D (oran Vinin) such that any such network trained upon until convergence has
sup
IE [fox))]] =
BERMO
where the expectation taken over initialisation.
Proof. For an infinitely-wied NN after convergence we have:
Sa (x) = (xx, Xrain O (y - No(#)) + No(ar),
where (O(æ, Mrain)); O (r, .x.) and (Qwain)s, = O ar,) and (y)); yt for = 1,..., |Puaial and No(x) is the output freshly initialised network with property [[No(=)] = As such, we have:
[x)) (ax, Xiran ) V
which is just the mean of a noiseless Gaussian Process conditioned on Dinin with kernel ( x'). Due the Representer Theorem, we can always rewrite this mean as:
[x)f =
ar (xx,
for Owe' R. Due the assumption of Theorem 3.1 must have least Virain such that ÿ + 0, and hence Rlllll o O, where o; Ox. for I,..., P
Let us study the predictions for the set Ar where (Aruiin); E Ro (Musin)) for some À > given by
E (N(AA = [N(Aar))
Σ ar O (Axx,
Ex
We rewrite this E [N e (Aæ,, 2.)) for i,j 1,..., | ||| As Xmai is non-degenerate and ((..')) ') is valid ker- nel, θ (Arriin Xuain) is full-rank VX > O, and thus has trivial nullspace. Therefore, as a of
θ (AArrinn Xuain) a. where θ (Arurains Nuain)
'When proving statements about infinite networks fo( we refer to the limiting network. Eg. in this case the statement equivalent to sup E limm fo(x))], where is the width of smallest layer.
37

---

0, E (A((A)))))) = θ (Aruain Nuai)) a # O. Hence, YA > 0, 3 E Xuain such that E (N) * 0.
We can arbitrarily multiply this expression by xn An 1, giving
R [N((A) = Ar
Or
And
E'Eii
Consider the where Niain is constructed such that x Tax Vax, x E Arain- Then, due to E.2 and Corollary
Corollary
(Aãe x²)
lim A-000
② (i:,x')
An
| || R x²)
Va'
Krain
where è is unit vector in the direction of Ø and the final transition holds as the condition æ x' Yx, x' E Xuain guarantees || > 0 Va E R
Invoking Corollary E.1 in the limit we then have
Ane
Aae x') An
lim E [N(A))]] = lim
x
ae Exain
lim
nn
lim A
Σ
Ocee
2 (Aăt x²)
Anz
e'Eiini
lim And
Σ Ow (ãc, x')
x=occ
ar E
where the last transition follows from the assumption that n > 0 and the first transition is valid we (0,
e ((Aā,o")) Aub
can choose a E _o such that limx-+00 Ex'e ax
Σ n' ex.... ax 200 x²) * O.
Concluding the proof, follows that
sup =E3%
JE [f( 00.
I Proof of Proposition 3.4
Proposition 3.4. There exists trained network (x) including Layer Norm layers such that
To(x)|= 00,
sup
mER*0 DE
where ir- arg minggo (0, Dimii) is the set of minimisers of the training loss.
Consider some network Sa (x: with fully-connected layers, operating in over-parametrised regime such that additional linear layer parameters may not further decrease the training loss over some finite training set D D Suppose further that
38

---

sup wER"0,BE0"
[fo(z) = 00,
where
Ω arg min L (0, Drain) 0 9 L (0, Duain)
min L (0, Dirain) BE
Bef2
Such network trivially exists; for any finite training dataset, any network interpolating the training targets (or, in the case of a degenerate training set with duplicated inputs, interpolating the minimal predictions with respect to the loss) will achieve the minimum possible loss of all networks. Moreover, there always exists a finite fully connected network achieving this g
Consider now the modified network fa(x)) formed by the insertion of a linear layer followed by an LN operation and ReLU nonlinearity anywhere in the network immediately followed by another linear layer. Specifically, we choose this additional layer to have 2n + 2 output nodes, where n number of nodes in the preceding layer.
Suppose we are free to choose the parameters of this linear layer. block diagonal form, we choose the following:
W
00,
In
In
0g]
where O, E R" is vector of zeros, In is the n X n identity matrix, and e > 0 is some positive constant.
Denoting the input to the linear layer as # E RR, the first and second moments of the output of the linear layer given by
1
Vn
Vn
+
2-
=
2n
+2
0² 1
n
+
2
2
+
11
=
n
E2 + =
2n
+2
E2
E2
t=1
i=1
where 2ims / Σ 2? is the root mean square value z. As all æ E Xiraio are finite, so are all
Eyms. hence we can upper bound Ems by the largest value of any element observed during training, mmax for all forward passes during training. Thus, setting = < Sps
7
1
for all forward passes during training. Consider the outputs from the Layer Norm during training, given by
Wz + b 1
E (W* b)
Vn
EZ -E2 -vn]"
Consider now the next linear layer, with weight matrix W E (2n+2)) Suppose in the absence of our additions (and noting the differing dimensions), this layer originally had optimal weight W e gmxn and bias b' m We can freely decompose W' as
39

---

w'
W On
In - 0g
such that the output of this linear layer is
z' W ORRLU (2) b*
W [On In In Og] ØRRLU (lvr Ex -E2 -vij ) + 6
W (ØROLU (sz) ØRALU (-sz)) + b²
W (max EZ) - max (0n, -Ex]] + b*
W (az) +
W' + b",
i.e. the unmodified output the network, where the second transition follows from the element-wise distribution of ReLU. By assumption of the Theorem, this minimises the empirical risk over the training set and the additional parameters and layer-norm do further decrease the value of this minimum. Hence, denoting a complete set of unaugmented network parameters drawn from Ω set minimisers L (æ, Drain), augmented with our additional and modified parameters, as 0 (W, b, W") it follows that
I. (0 (W, b, W') Drmin) min (0, Dirai)) 8E8
min L (0(z), Drain) deô
and so
0 (W. b, W EΩ" Vs < 2max:
where the first transition follows from the assumption of the theorem that the network is sufficiently overparametrised such that these augmentations cannot further reduce the empirical risk over the training set. Thus, when considering supremums over f2* we can consider the limiting case - Considering the dual-limiting case where - 0 faster than ||>||| - co for some E Ro it is clear that the assumption on holds for such æ, thus the Layer Norm effectively bypassed and we approach network exactly equivalent to the unmodified network. More precisely,
Va E To 30 Ô* such that fo(x)) = fo(
Noting that by assumption that
[fo( 0o, =ER"
sup
it is clear that
[5(00) 00.
sup *ER"0 BEΩ-
Experimental Details
To run all experiments we used NVIDIA GeForce RTX 3090 with 24GB of memory. For all experiments we used the same architecture consisting fully-connected layers with ReLU non- lincarities. All hidden layers have size of 128. To be consistent with the theory, initialised
40

---

weights of fully-connected layers with Kaiming initialisation 16 and biases were sampled from NO of) with of 0.01. For the UTK experiments in Section 4. frozen ResNet- 18 [77 used before the fully-connected layers. We utilised Adam optimiser 21 and MSE Loss for all experiments. See Table 3 below for exact hyperparameters settings. For XGBoost we use the XGBoost Python package with hyperparameter values set to default values and number of boosting rounds set to 100.
fl
Training took less than 10 seconds per seed for experiments in Section 4.1 and 4.2 and around minutes per seed for experiments in Section 4.3
Table 3: Hyperparameter values used throughout the experiments.
Experiment
Method
Batch size
Epochs
Learning rate
All
100 (entire dataset)
3000
0.001
4.2
All
40132 (entire dataset)
2500
0,001
4.3
Standard NN
128
10
0,001
LN after 1st
128
10
0.001
LN after 2nd
128
10
0.003
LN after every
128
10
0.003
K
Additional Experiments: Different Activations
sine
quadratic
linear
20
-5
20
10
10
20
10
10
2b
20
i0
1b
20
X
Standard NN (Relu)
Standard NN (Gelu)
Standard NN (Swish)
sine
quadratic
linear
4-
0-
-2
0-
2
-10
10
-10
-10
LN (Relu)
LN (Gelu)
LN (Swish)
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 the mean estimator.
41

---

L
Additional Experiments: Batch Norm vs Layer Norm
sine
quadratic
linear
10
-50
-100
-50
50
100
-20
-10
0
10
20
-20
-10
10
20
Standard
after every
after every
Figure 6: Predictions made by networks with various activations when trained on synthetic datasets. Red dots show the train set datapoints. The solid lines indicate average values over 5 seeds and shaded areas 95% confidence intervals of the mean estimator.
M Additional Experiments: Transformer
1750
Transformer Variant
Without LayerNorm
With LayerNorm
1500
1250
Prediction
1000
750
500
250
20
40
50
70
100
Index
Figure 7: Results transformer toy problem. We train 2-layer decoder-only transformer with 4 heads and embedding size of 64 with two versions: with and without layer normalisation after input embedding layer. The model trained to on sequences of form f()) a b i, where a. ~ (000 1) and E 21,2,... ... 30) to predict the value with index + given value with indices from to i. On the plot above we show the average prediction made by each model, when tested for index i E (20, 21,... 100} and as such beyond its training domain. In red we show model without LayerNorm, whereas blue show the model with LayerNorm. Solid lines are means over seeds and shaded areas 95% confidence intervals.
42

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA]. .
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] "provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). ). general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We made claims about deriving novel theoretical results, which we do in section Theoretical Results and about verying them empirically, which do in section Experiments.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
43

---

Justification:
We discuss the limitations in Conclusions section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approxinations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions are stated at the beginning of Thcoretical Results section all proofs are in the Appendix.
Guidelines:
The answer NA means that the paper does not include Theoretical Results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: All details described in the Experiments section.
Guidelines:
44

---

The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide full code by anonymised link in the Experiments section.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
wwt.:/oubbbiiiooooooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https c/iips.cc/public//uu
CodeeSumiiii wettsbbmmssiiioi ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
45

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Full details provided in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All tables and plots have confidence intervals.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We list compute resources in Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
46

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: The paper is concerned with foundational research and not tied to any particular application that can cause ethical concerns.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The research presented in the work is foundational and not tied to any particular application.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly. harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learn from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: The paper is not accompanied by a release of any new data sets or pre-trained models.
47

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We clearly credited the sources for datasets used in Sections 4.2 and 4.3
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperssithcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We released our code and provided README.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
48

---

Answer: [NA]
Justification:
Papers contains no experiments including crowdsourcing or human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing. editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [No]
Justification:
The paper does not describe any form of LLM usage.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non standadd components.
Please refer to LLM policy (nttps://neur)) hwtp://eeecoeeeeeee/////0////// for what should or should not be described.
49


Reviews:
[{"review_id": "SaTjTjzavR", "reviewer": "Reviewer_QsAy", "summary": "The paper studies the boundedness of neural networks with layer norm in the infinite width NTK regime. The first result is that there is a “bad” dataset such that training the network on it will result in an unbounded function. Next, it is shown that on any training set, if the network contains layer norm, then it converges to a bounded function. Finally, it is shown that if we consider only ERM classifiers (rather than considering the training dynamics), then layer norm is not enough to bound the converged network. Several experiments are given that verify the results on toy and realistic datasets.", "questions": "Would it be tractable to make the analysis in the finite width NTK? Will similar results still hold, assuming a large but finite width?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The results in this paper are strong and somewhat surprising. First, they imply that one of the major advantages of using LN, which is very popular in modern architectures, specifically transformers, is that it makes the resulting function bounded. Second, that the boundedness indeed depends on the learning procedure, and performing ERM (not through gradient methods) may result in a non-bounded function.\n- The experiments are quite thorough. They also show that not only in the NTK regime, the boundedness property holds. This is important since, in practice, networks are not trained in the kernel regime.\n\nWeaknesses:\n- I think the main weakness in the paper is that the authors analyze LN in the infinite width NTK regime. Now, I understand that using the NTK regime makes the analysis tractable, but why is the infinite width necessary? Looking at the proofs, I think they can be generalized to the finite width NTK regime, with slightly more work.\n- This is a paper about layer normalization, where layer normalization is not defined. I highly suggest devoting some space in the preliminaries to properly define LN, since even if it is “common knowledge” for the community, in theoretical papers, it would make the results easier to understand.\n\nMinor comments:\n- Typo - Line 66 should be 1 until n\n- Proposition 3.4 seems to be true for any dataset. It would be better to phrase it this way; it also makes the proposition stronger.\n\nTo conclude, I think this is a strong paper with interesting results and a thorough experimental study.", "was_revised": true, "final_justification": "I think this is a nice theory paper. It provides some insights on the boundness of layernorm, although I agree with the other reviewers, it doesn't provide new techniques, and it is limited to the infinite width NTK regime. I retain my score, but slightly lowered my confidence in light of this.", "rebuttal": "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.\n\n**Extending results to finite width NTK regime**\n\nAs 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.\n\n**Layer Norm not defined**\n\nThe 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."}, {"review_id": "p69OQaYhqy", "reviewer": "Reviewer_L8DJ", "summary": "This paper studies the neural tangent kernel (NTK) with layer normalization, showing that the corresponding kernel predictor remains bounded for any input. In contrast, without layer normalization, neural network outputs can become pathologically large. The theoretical results are validated through experiments on synthetic and UCI datasets.", "questions": "In the proof of Proposition 3.4, how do you establish the existence of such a network? The statement \"There exists a network with fully-connected layers, operating in an overparameterized regime such that...\" requires clarification.", "limitations": "The theoretical contribution is incremental.", "rating": 2, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 1, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\nThe paper is well-written with clear logic, and experiments effectively support the theoretical results.\n\nWeaknesses:\n\n1. The contribution is incremental—the result is somewhat expected given the assumption of homogeneous activation functions.\n2. The motivation for studying this problem is unclear. In practical settings, input data are typically not large due to standard normalization practices. The term \"stable\" may be inappropriate here, as large inputs can be easily controlled and are not a primary source of instability. \nStudying stability within bounded domains would be more meaningful.\n3. The definition of layer normalization should be introduced in the preliminaries.", "was_revised": true, "final_justification": "This paper is positioned as a theory paper, but I did not find any new insights or techniques in the analysis. Therefore, I tend to keep my score unchanged.", "rebuttal": "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.\n\n**Results expected given the homogeneous activation functions**\n\nAs 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.\n\nIt 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”.\n\n**The motivation for studying this problem is unclear**\n\nOn 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)|$. \n\nThe 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).\n\nA 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.\n \n**Layer Norm not defined**\n\nThe 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.\n\n**Establishing the existence of the network in Proposition 3.4**\n\nThis 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."}, {"review_id": "38l2PuQgOB", "reviewer": "Reviewer_oWLy", "summary": "This paper studies the impact of layer normalization on the output of neural network on out-of-distribution (OOD) data. The analysis is carried out using NTK (neural tangent kernel) of the neural network with layer\nnormalization. The main result is that, as long as there is one layer that performs layer normalization, the output of the trained NTK model will be bounded for out-of-distribution data, regardless of how large the\ninput is. The intuition is that, with such layer normalization, the normalized input will be bounded in expectation. Hence, the output of the network cannot be too large.", "questions": "1. The reviewer is a bit skeptical why keeping the output bounded is necessarily a desirable goal for out-of-distribution data. For example, consider Figure 1, the rightmost sub-figure. The desirable regressor appears to be a line, not a bounded function. If the line is indeed the ground truth, then having a bounded output just means that the difference from the ground truth will be unbounded! Then, why is the boundedness desirable? \n\n2. The output bound in Theorem 3.3 depends on $\\lambda_\\min$. However, the value of $\\lambda_\\min$ itself may depend on the training data and network architecture. This dependency may make the advantage of the layer normalization less clear. For example, is it possible that, with layer normalization, the $\\lambda_\\min$ value become significantly smaller than that without layer normalization?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength: \n\n1. The main result that one layer of normalization will keep the output bounded on out-of-distribution data, while not surprising, is interesting. \n\nWeakness:\n\n1. It is unclear why keeping the output bounded is the desirable goal for dealing with out-of-distribution data.", "was_revised": true, "final_justification": "The response partially addresses my earlier doubts. Thus, I am willing to increase my score to be slightly positive.", "rebuttal": "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:\n\n**Why keeping the output bounded is the desirable goal for dealing with out-of-distribution data**\n\nAs we wrote in the Conclusion section: \n“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.”. \nWithin 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: \n“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”.\n\n**Is it possible layer norm will shrink the eigenvalue**\n\nIn 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.\n\nWe 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."}, {"review_id": "xtFc2ikkos", "reviewer": "Reviewer_uiPW", "summary": "The paper proves—using Neural Tangent Kernel (NTK) theory—that inserting just one Layer Normalization (LN) anywhere in an (infinitely‑wide) fully‑connected network fundamentally changes the kernel: it becomes a bounded‑variance kernel. Consequently, the posterior mean of the corresponding Gaussian Process remains bounded for all inputs, no matter how far they lie from the training domain, while standard architectures without LN can diverge arbitrarily.\nKey theoretical results are\n1) Theorem 3.1: networks with positive n-homogeneous activations but no LN can yield unbounded outputs on some dataset/input pair.\n2) Theorem 3.2 & B.2: adding a single LN forces the variance term Θ(x,x) of the NTK to be uniformly bounded.\n3) Theorem 3.3: this bounded kernel implies a data‑dependent uniform bound on |E f(x)| for the trained network.\nThe authors confirm the theory on (i) 1‑D toy regressions, (ii) the UCI‑Protein RMSD task—testing proteins larger than seen in training, and (iii) UTKFace age‑prediction for ethnicities absent from training. In every case, architectures with ≥1 LN avoid “exploding” predictions and achieve markedly higher OOD R² than baselines.", "questions": "1) Finite‑width behaviour: Have you explored how large a width is needed for the theoretical bound to manifest? A sensitivity plot (width vs. max|f(x)| on extrapolation grid) would strengthen the practical message.\n\n2) Beyond fully‑connected layers: Can Theorem B.2’s extension via Tensor Programs be instantiated for a modern ResNet/Transformer? A small experiment or discussion would clarify applicability.\n\n3) Bound tightness: The worst‑case bound in Theorem 3.3 scales with √|D|/λ_min. In practice this may be very loose. Could the authors report empirical maxima and compare to the theoretical bound?\n\n4) Activation assumptions: Assumption 3.2 excludes GELU/Swish (except via asymptotic argument). Are the bounds still valid for popular Transformer activations numerically?", "limitations": "yes", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n1) Strong Theoretical Guarantee: Provides a general and rigorous NTK-based proof that a single LayerNorm anywhere in an (infinitely-wide) fully connected network suffices to guarantee bounded network outputs for all inputs, regardless of distance from the training set.\n\n2) Novelty: This is the first work to formally prove that LayerNorm fundamentally changes the extrapolatory behavior of deep networks, rather than merely showing this effect empirically.\n\n3) Empirical Validation: The paper supports its theory with systematic experiments on toy regression tasks and two real-world benchmarks (protein RMSD and UTKFace age prediction), using multiple seeds and reporting confidence intervals.\n\n4) Practical Safety Implications: Highlights the practical significance of output boundedness for safety-critical ML deployments, such as robot control.\n\n5) Clarity of Empirical Results: Visualizations (e.g., Figures 1–3) and tables provide clear evidence of the “exploding” vs. “stable” behaviors, and OOD gains are easy to interpret.\n\nWeaknesses\n1) Limitation to Infinite-Width Theory: All theoretical results hold strictly in the infinite-width, full-batch regime; no non-asymptotic or finite-width generalization bounds are provided.\n\n2) Limited Architectural Breadth: Experiments are restricted to shallow MLPs; the only use of a convolutional backbone (ResNet-16) is as a frozen feature extractor in the UTKFace task, so end-to-end results for CNNs or Transformers are not established.\n\n3) Boundedness ≠ Correctness: The guarantee is about boundedness of output magnitude, not correctness, calibration, or improved prediction accuracy. The practical relevance for tasks not threatened by output explosions may thus be limited.\n\n4) Dense Theoretical Presentation: Some notation and proof sketches (especially kernel recursions) are dense, and several technical steps are deferred to appendices, making parts of the analysis less accessible to non-experts.\n\n5) LN Sufficient but Not Proven Necessary: The necessity of LayerNorm is empirically indicated (i.e., no alternative mechanism is shown to work), but not formally established; other mechanisms could potentially provide similar guarantees.", "was_revised": true, "final_justification": "I appreciate the authors' efforts to address my concerns. I raised my score to 5.", "rebuttal": "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.\n\n**Limitation to Infinite-Width Theory**\n\nWhile 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.\n\n**Beyond fully‑connected layers: Can Theorem B.2’s extension via Tensor Programs be instantiated for a modern ResNet/Transformer?**\n\nYes, 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).\n\nWe 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:\n\n| Index ($i$) | Transformer w/ LN (95% CI)    | Transformer w/o LN (95% CI)     |\n|----------------|--------------------------------------|---------------------------------|\n| 50              | 55.4613 ± 11.591                 | 62.0827 ± 3.068                 |\n| 70              | 57.7766 ± 13.492                 | 146.1958 ± 59.857               |\n| 100            | 57.7941 ± 13.508                 | 412.1160 ± 276.456              |\n\nAs 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.\n \n**Activation assumptions:**\n\nWe 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.\n\nTo 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):\n\n| X        | ReLU (No Norm)        | Swish (No Norm)      | GeLU (No Norm)       | ReLU (LN)            | Swish (LN)           | GeLU (LN)            |\n|----------|------------------------|-----------------------|-----------------------|-----------------------|-----------------------|-----------------------|\n| -1000.00 | 570.08 ± 79.62         | 276.34 ± 44.08        | 271.30 ± 32.08        | 2.63 ± 0.14           | 2.65 ± 0.47           | 2.49 ± 0.51           |\n| -595.96  | 338.44 ± 46.98         | 164.36 ± 26.30        | 161.29 ± 19.21        | 2.62 ± 0.14           | 2.64 ± 0.46           | 2.48 ± 0.51           |\n| -191.92  | 106.79 ± 14.34         | 52.44 ± 8.61          | 51.28 ± 6.36          | 2.56 ± 0.14           | 2.59 ± 0.46           | 2.43 ± 0.50           |\n| 10.10    | -2.19 ± 0.20           | -2.18 ± 0.25          | -1.74 ± 0.13          | -0.78 ± 0.26          | -0.83 ± 0.16          | -1.05 ± 0.23          |\n| 212.12   | -125.70 ± 4.36         | -57.29 ± 6.52         | -53.20 ± 7.23         | -2.12 ± 0.61          | -2.36 ± 0.39          | -2.67 ± 0.30          |\n| 616.16   | -372.86 ± 12.51        | -167.79 ± 18.82       | -156.46 ± 21.32       | -2.17 ± 0.62          | -2.40 ± 0.40          | -2.71 ± 0.30          |\n| 1000.00  | -607.64 ± 20.24        | -272.77 ± 30.50       | -254.56 ± 34.72       | -2.18 ± 0.63          | -2.41 ± 0.40          | -2.72 ± 0.30          |\n\n\nWhere 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.\n\n**Boundedness ≠ Correctness:**\n\nWe 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.\n\n**Dense Theoretical Presentation**\n\nWe will use the additional page in camera-ready version to improve presentation and provide a more intuitive explanation of the derived results and proofs.\n\n**LN Sufficient but Not Proven Necessary**\n\nThe 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.\n\n**Finite‑width behaviour**\n\nAs 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:\n| **Model Size** | $\\max \\|f(x)\\|$ | **CI**       | **% Error to GP** |\n|----------------|------------------|--------------|-------------------|\n| 16             | 1.076239         | 0.049857     | 21.93%            |\n| 32             | 1.068731         | 0.019022     | 22.44%            |\n| 64             | 1.094828         | 0.025697     | 20.54%            |\n| 128            | 1.152339         | 0.022382     | 16.38%            |\n| 256            | 1.203813         | 0.036533     | 12.63%            |\n| 512            | 1.234375         | 0.019526     | 10.43%            |\n| 1024           | 1.282959         | 0.016616     | 6.88%             |\n| 2048           | 1.317124         | 0.007575     | 4.42%             |\n| NTK-GP      | 1.378103         | 0.000000     | 0.00%             |\n\n\nNTK-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.\n\n**Bound tightness:**\n\nAs 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]. \n\n| **Dataset Size** | **Theoretical Bound** | **LN** $\\max \\|f(x)\\|$ | **Standard** $\\max \\|f(x)\\|$ |\n|------------------|------------------------|--------------------------|-------------------------------|\n| 2                | 2.35                   | 0.73                     | 3304                          |\n| 3                | 5.11                   | 1.01                     | 3390                          |\n| 10               | 104.24                 | 0.61                     | 4031                          |\n\nThe 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."}, {"review_id": "vQ2T6KpIzF", "reviewer": "Reviewer_WTJo", "summary": "This paper shows that a single LayerNorm layer induces boundedness of the learned function in the infinite-width limit under neural tangent parameterization and MSE loss. It is argued that this boundedness causes extrapolatory stability outside the training domain in practice.", "questions": "- **Layernorm position.** It is interesting that one normalization layer — irrespective of its position — suffices. Can you provide an informed intuition why the position does not seem to matter and under which circumstances? Concretely analysing the resulting NTKs could enable to quantify this statement. It could also enable to quantify whether the differences in Table 2 between layernorm positions stem from the data or from the convolutional architecture.\n- **Generalizability beyond the kernel regime.** The potential statement ‘Normalizing the last-layer activations induces norm-bounded outputs, as long as the output layer has bounded operator norm throughout training.’ almost sounds tautological. Do you think it would be feasible to show similar boundedness results in much more generality beyond the kernel regime?\n- **Other stabilizing components.** Your results hold for LayerNorm. Would other stabilizing components such as BatchNorm layers or the cross-entropy loss (as would be relevant for practice) also ìnduce boundedness OOD?\n- **RKHS statement.** In lines 127 to 129 you write an equivalence statement that the RKHS only contains functions which grow without bound, but is it not enough, that there exist functions in the RKHS that grow without bound and that the coefficients learned by the training procedure on these unbounded functions are non-negligible?", "limitations": "The limitation that NTP is used for the theoretical results and that the consequence for practical settings remains unclear is not sufficiently clarified or empirically evaluated.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n\nThe paper is well written and the proof sketches are informative. It is interesting that one Layernorm suffices for boundedness of the NTK, irrespective of its position. The boundedness result holds under weak distributional and architectural assumptions. Albeit a simple and unsurprising intervention, the improved OOD generalization is an impactful consequence, even though all popular architectures already employ normalization layers so that the contribution is limited to a principled understanding of the benefits, which is still interesting. The kernel theory results look correct.\n\n\n**Weaknesses:**\n\nThe following main weaknesses should be addressed:\n\n- **Limited correspondence to practical neural networks due to NTP.** The paper uses neural tangent parameterization, as explicitly specified only in Appendix A. This results in a kernel regime in the infinite-width limit, which yields its theoretical tractability. This could have been clarified. Even Jacot et al. (2018) noted that NTP is not standard parameterization (SP) (as is ambiguously suggested in l.80), hence the implications to finite networks in SP remains unclear, and has been shown to not hold at realistic widths in various papers. Hence a closer empirical evaluation in practical settings would be crucial to quantify how well boundedness results hold in finite networks at the end of training and which impact non-vanishing feature learning has on the main conclusions in this paper. Is the soft cosine similarity still useful for describing Layernorm networks at the end of training? Do predictions of networks without Layernorm indeed diverge on realistic test and OOD inputs, or is this not a practical issue? The experimental results on real data are only tangentially related to the theory, and it remains unclear whether they are cause by the fact that the main theoretical conclusions still hold in practical finite networks after training on vision, protein or language data. Using frozen ResNet features is arguably an artificial setting. Full training would be more interesting. It would be interesting to see whether the main claims also hold in feature learning networks, such as parameterized in muP.\n- **Missing acknowledgement of related work.** After a 5-minute literature search, it looks like normalization for improved OOD performance has been previously considered, see Regmi et al. (2024) and references therein such as LogitNorm (Wei et al., 2024). It seems that this line of work should at least be mentioned in the Related Work section.\n\n**Minor point:**\n\n- MSE loss is nice for theoretical tractability, but CE loss is important for capturing practical calibration properties and could result in very different conclusions.\n\n**Typos:** l. 78: then, 95: the the, 146: NTK variance?, 179: …it is possible…, 207: respectively.\n\n**References:**\n\nRegmi, Sudarshan et al. “T2FNorm: Train-time Feature Normalization for OOD Detection\nin Image Classification.” *CVPR workshop paper, 2024.*\n\nWei, Hongxin, et al. \"Mitigating neural network overconfidence with logit normalization.\" *ICML*, 2022.", "was_revised": true, "final_justification": "The paper is technically solid, but the results are not very surprising: Layernorm induces boundedness which often improves OOD performance. The theory is limited to the kernel regime, which limits practicality. The promised changes will improve the paper but do not affect my overall slightly positive evaluation.", "rebuttal": "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:\n\n**Layernorm position**\n\nIntuitively, 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.\n\n**Missing related work**\n\nWe 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.\n\n**MSE vs CE Loss**\n\nOur 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.\n\n**Generalizability beyond the kernel regime**\n\nWhile 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.\n\n**Other stabilizing components**\n\nFor 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.\n\nTo 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):\n| X         | Standard NN        | LayerNorm              | BatchNorm              |\n|-----------|----------------------------|-----------------------------|-----------------------------|\n| -979.80   | 558.49 ± 77.99             | 2.63 ± 0.14                 | 134.86 ± 103.28             |\n| -595.96   | 338.44 ± 46.98             | 2.62 ± 0.14                 | 81.68 ± 62.56               |\n| -191.92   | 106.79 ± 14.34             | 2.56 ± 0.14                 | 25.70 ± 19.69               |\n| 10.10     | -2.19 ± 0.20               | -0.78 ± 0.26                | -0.29 ± 0.20                |\n| 212.12    | -125.70 ± 4.36             | -2.12 ± 0.61                | -15.03 ± 7.94               |\n| 616.16    | -372.86 ± 12.51            | -2.17 ± 0.62                | -44.40 ± 23.51              |\n| 1000.00   | -607.64 ± 20.24            | -2.18 ± 0.63                | -72.31 ± 38.30              |\n\nWhere 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. \n\n**RKHS statement**\n\nThe 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.\n\n**The limitation that NTP is used for the theoretical results**\n\nThis 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."}]

Author rebuttal:

Meta-review: The paper shows that in an infinite-width NTK regime 
one layer norm guarantees that the outputs of a network
trained until convergence are bounded (and not stable as 
suggested by the title), while in the absence of a 
layer norm this is not the case for certain, malign data sets.

During the discussion, two reviewers raised their score 
so that now 4 out of 5 are voting for acceptance. The fifth 
reviewer remained sceptical staying at his rating 2. 
Here are his major concerns:

- infinite-width NTK regime and parametrization
- results are not surprising in view of the assumed homogeneity of the activation function
- the paper is not about `extrapolation stability' as it only considers extrapolation boundedness, which could also
  be ensured by other means
  
All of these concerns have certainly some validity and are shared to some extend 
by the other reviewers (though one reviewer finds the results surprising). 
On the other hand, the paper also has its merits:

- rigorous analysis and novelty of results
- good supporting experiments

In summary, making a decision for this paper is really about weighting 
the pros and cons. Personally, I think, the two reviewers rating it 
by 5 may be a bit too positive, while the negative reviewer is certainly
way too negative. 

I tend to vote to accept the paper, although some comparison to 
other papers is probably a good idea.
Decision: Accept (poster)