O N THE N EURAL T ANGENT K ERNEL OF E QUILIBRIUM M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
This work studies the neural tangent kernel (NTK) of the deep equilibrium (DEQ) model, a practical “inﬁnite-depth” architecture which directly computes the inﬁnite-depth limit of a weight-tied network via root-ﬁnding.
Even though the
NTK of a fully-connected neural network can be stochastic if its width and depth both tend to inﬁnity simultaneously, we show that contrarily a DEQ model still enjoys a deterministic NTK despite its width and depth going to inﬁnity at the same time under mild conditions.
Moreover, this deterministic NTK can be found efﬁciently via root-ﬁnding
I NTRODUCTION
Implicit models form a new class of machine learning models where instead of stacking explicit “layers”, they output z s.t g ( x, z ) = 0 , where g can be either a ﬁxed point equation (Bai et al., 2019), a differential equation (Chen et al., 2018b), or an optimization problem (Gould et al., 2019).
This work focuses on deep equilibrium models, a class of models that effectively represent a “inﬁnitedepth” weight-tied network with input injection.
Speciﬁcally, let f θ be a network parameterized by θ , let x be an input injection, DEQ ﬁnds z ∗ such that f ( z ∗ , x ) =
z ∗ , and uses z ∗ as the input for downstream tasks.
One interesting question to ask is, what will DEQs become if their widths also go to inﬁnity?
It is well-known that at certain random initialization, neural networks of various structures converge to Gaussian processes as their widths go to inﬁnity (Neal, 1996; Lee et al., 2017;
Yang, 2019; Matthews et al., 2018; Novak et al., 2018; Garriga-Alonso et al., 2018).
Recent deep learning theory advances have also shown that in the inﬁnite width limit, with proper initialization (the NTK initialization), training the network f θ with gradient descent is equivalent to solving kernel regression with respect to the neural tangent kernel (NTK)
(Arora et al., 2019; Jacot et al., 2018;
Yang, 2019; Huang et al., 2020).
These kernel regimes provide important insights to understanding how neural networks work.
However, the inﬁnite depth (denote depth as d ) regime introduces several caveats.
Since the
NTK correlates with the inﬁnite width (denote width as n ) limit, a question naturally arises as how do we let n, d → ∞ ?
Hanin & Nica (2019) proved that as long as d/n ∈ (0 , ∞ ) , the
NTK of vanilla fully-connected neural network (FCNN) becomes stochastic.
On the other hand, if we ﬁrst take the n → ∞ , then d → ∞ 1 , Jacot et al.
(2019) showed that the NTK of a
FCNN converges either to a constant ( freeze ), or to the Kronecker Delta ( chaos ).
In this work, we prove that with proper initialization, the NTK-of-DEQ enjoys a limit exchanging property
[Equation]
Θ
[Equation]
) n ( x, y ) with high probability, where Θ ( d ) n denotes the empirical NTK of a neural network with d layers and n neurons each layer.
Intuitively, we name the left hand side “DEQ-of-NTK” and the right hand side “NTK-of-DEQ”.
The NTK-ofDEQ converges to meaningful deterministic ﬁxed points that depend on the input in a non-trivial way, thus avoiding the freeze vs. chaos scenario.
Furthermore, analogous to DEQ models, we can compute these kernels by solving ﬁxed point equations, rather than iteratively applying the updates as for traditional NTK.
We evaluate our approach and demonstrate that it matches the performance of existing regularized NTK methods.
B ACKGROUND AND P RELIMINARIES
A vanilla FCNN has the form g ( t )
=
σ ( W ( t ) g ( t − 1) + b ( t ) ) for the t -th layer, and in principle t can be as large as one wants.
A weight-tied FCNN with input injection (FCNN-IJ) makes the bias term related to the original input and ties the weight in each layer by taking the form z ( t ) :
= f ( z ( t − 1) , x ) = σ ( Wz ( t − 1) + Ux + b ) .
Bai et al.
(2019) proposed the DEQ model, which can be equivalent to running an inﬁnite-depth FCNN-IJ, but updated in a more clever way.
The forward pass of DEQ is done by solving f ( z ∗ , x ) =
z ∗ .
For a stable system, this is equivalent to solving lim t →∞ f ( t ) ( z (0) , x ) .
The backward iteration is done by computing df ( z ∗ , x ) /dz ∗ directly through the implicit function theorem, thus avoiding storing the Jacobian for each layer.
This method traces back to some of the original work in recurrent backpropagation (Almeida, 1990; Pineda, 1988), but with speciﬁc emphasis on: 1) computing the ﬁxed point directly via root-ﬁnding rather than forward iteration; and 2) incorporating the elements from modern deep networks in the single “layer”, such as self-attention transformers (Bai et al., 2019), multi-scale convolutions (Bai et al., 2020), etc.
DEQ models achieve nearly state-of-the-art performances on many large-scale tasks including the
CityScape semantic segmentation and ImageNet classiﬁcation, while only requiring constant memory.
Although a general DEQ model does not always guarantee to ﬁnd a stable ﬁxed point, with careful parameterization and update method, monotone operator DEQs can ensure the existence of a unique stable ﬁxed point (Winston & Kolter, 2020).
The study of large width limits of neural networks dates back to Neal (1996), who ﬁrst discovered that a single-layered network with randomly initialized parameters becomes a Gaussian process (GP) in the large width limit.
Such connection between neural networks and GP was later extended to multiple layers (Lee et al., 2017; Matthews et al., 2018) and various other architectures (Yang, 2019; Novak et al., 2018; Garriga-Alonso et al., 2018).
The networks studied in this line of works are randomly initialized, and the GP kernels they induce are often referred to as the NNGP.
A line of closely-related yet orthogonal work to ours is the mean-ﬁeld theory of neural networks.
This line of work studies the relation between depth and large-width networks (hence a GP kernel in limit) at initialization.
Poole et al.
; Schoenholz et al.
(2016) showed that at initialization, the correlations between all inputs on an inﬁnitely wide network become either perfectly correlated ( order ) or decorrelated ( chaos ) as depth increases.
They suggested we should initialize the neural network on the “edge-of-chaos” to make sure that signals can propagate deep enough in the forward direction, and the gradient does not vanish or explode during backpropagation (Raghu et al., 2017;
Schoenholz et al., 2016).
These mean-ﬁeld behaviors were later proven for various other structures like RNNs, CNNs, and NTKs as well (Chen et al., 2018a; Xiao et al., 2018; Gilboa et al., 2019;
Hayou et al., 2019).
We emphasize that despite the similar appearance, our setting avoids the order vs. chaos scheme completely by adding input injection.
The injection guarantees the converged
NTK depends nontrivially on the inputs, as we will see later in the experiments.
While previous results hold either only at initialization or networks with only last layer trained, analogous limiting behavior was proven by Jacot et al.
(2018) to hold for fully-trained networks as well.
They showed the kernel induced by a fully-trained inﬁnite-width network is the following:
[Equation]
where N represents the Gaussian distribution.
They also gave a recursive formulation for the NTK of FCNN.
Arora et al.
; Alemohammad et al.
Yang (2020) later provided formulation for convolutional NTK, recurrent NTK, and other structures.
One may ask what happens if both the width and the depth in a fully-trained network go to inﬁnity.
This question requires careful formulations as one should consider the order of two limits, as Hanin & Nica (2019) proved that width and depth cannot simultaneously tend to inﬁnity and result in a deterministic NTK, suggesting one cannot always swap the two limits.
An interesting example is that
Huang et al.
(2020) showed that the inﬁnite depth limit of a ResNet-NTK is deterministic, but if we let the width and depth go to inﬁnity at the same rate, the ResNet behaves in a log-Gaussian fashion (Li et al., 2021).
Meanwhile, the inﬁnite depth limit of NTK does not always present favorable properties.
It turns out that the vanilla FCNN does not have a meaningful convergence: either it gives a constant kernel or the Kronecker Delta kernel (Jacot et al., 2019).
Our contributions.
We ﬁrst show that unlike the inﬁnite depth limit of NTK to FCNN, the DEQof-NTK does not converge to a degenerate kernel.
This non-trivial kernel can be computed efﬁciently using root-ﬁnding.
Moreover, the NTK-of-DEQ coincides with the DEQ-of-NTK under mild conditions.
Although the proofs here involved inﬁnite limits, we also show numerically that reasonably large networks converge to roughly the same quantities as predicted by theory, and we show the
NTK-of-DEQ matches the performances of other NTKs on real-world datasets.
2. N OTATION
We write capital letter W to represent matrices or tensors, which should be clear from the context, and use [ W ] i to represent the element of W indexed by i .
We write lower case letter x to represent vectors or scalars.
For a ∈ Z + , let [ a ] = { 1 , . . .
, a } .
Denote σ ( x ) = √ 2 max(0 , x ) as the normalized ReLU and ˙ σ its derivative (which only needs to be well-deﬁned almost everywhere).
The symbol σ 2 a with subscript is always used to denote the variance of random variable a .
We write
N ( µ, Σ ) as the Gaussian distribution with mean µ ∈ R d and covariance matrix Σ ∈ R d × d .
We let S d − 1 be the unit sphere embedded in R d .
We use n, d to denote width and depth respectively, and write G ( d ) n to stress G has depth d and width n , where G can represent either a kernel or a neural network.
We use the term empirical NTK to represent 󰀟
[Equation]
󰀠 .
We write
[Equation]
and
[Equation] [Equation]
G
[Equation]
d ) n to denote limits are taken.
All missing proofs can be found in the appendix.
NTK- OF -DEQ WITH F ULLY - CONNECTED L AYERS
In this section, we show how to derive the NTK of the fully-connected DEQ.
Let m be the input dimension, x, y ∈ S
m − 1 be a pair of inputs, n be the width of the h -th layers where h ∈ [ d ] .
Let g (0) ( x ) =
∈ R n .
Deﬁne the depth- d approximation to a DEQ as the following:
[Equation]
where h ∈ [ d ] , W ( h ) ∈ R n × n ,
U ( h ) ∈ R n × m , v ∈ R n are the internal weights and b ( h ) ∈ R n are the bias terms.
The actual DEQ effectively outputs f ( ∞ ) n
[Equation]
The forward pass is solved using root-ﬁnding or ﬁxed point iteration, and the backward gradient is calculated using implicit function theorem instead of backpropogation.
One thing to note is that usually DEQs require tied-weights: W ( h )
=
W , U ( h )
=
U .
and b ( h ) =
b for all h .
It turns out for the inﬁnite width regime, DEQ with tied weights and DEQ without tied weights will induce the same NTK.
We will discuss this point in more detail later.
Let Θ ( d ) n ( x, y ) be the empirical NTK of f ( d ) n .
In Section 3.1, we will derive for an arbitrarily ﬁxed d , the “ﬁnite depth iteration to DEQ-of-NTK” Θ ( d ) =
lim
n
[Equation]
Θ ( d ) n .
In Section 3.2, we show that Θ ( d ) converges to a deterministic DEQ-of-NTK.
Furthermore, we prove that lim d →∞ lim n →∞ Θ ( d ) n = lim n →∞ lim d →∞ Θ ( d ) n with high probability, that is, the DEQ-of-NTK equals the NTK-of-DEQ.
3. F INITE D EPTH I TERATION TO DEQ- OF -NTK
Under the expressions in the beginning of Section 3, let us pick σ W , σ U , σ b ∈ R arbitrarily in this section, and require the following NTK initialization.
NTK initialization.
We randomly initialize every entry of every W, U, b, v from N (0 , 1) .
g (2 , 2) g (2 , 1) ( x ) g (2 , 3) ( x ) g (1 , 2) ( x ) g (1 , 1) ( x ) g (1 , 3) ( x )
W (2) W (2)
W (1)
W (1) W (1) U (2)
U (2)
U (2) g (2 , 0) ( x ) g (1 , 0) ( x )
W (2)
W (1) xxx
U (1) U (1)
U (1)
The ﬁnite depth iteration to the DEQ-of-NTK can be expressed as the following:
Theorem 3.1.
Recursively deﬁne the following quantities for h ∈ [ d ] :
[Equation]
Then the d -depth iteration to the DEQ-of-NTK can be expressed as:
[Equation]
where by convention we set ˙ Σ ( d +2)
( x, y ) = 1 .
One can realize that the derivation is done as if the weights in each layers are independently drawn from the previous layers, thus violating the formulation of DEQs.
Nonetheless, it has been proven that under certain conditions, the tied-weight NN and untied-weight NN induce the same NTK, see
Remark 3.2. Remark 3.2 .
While our derivation is done on untied weights, the NTK of its weight-tying counterpart converges to the same point.
This is formally done using the Nestor program introduced in Yang (2019; 2020).
The neural architecture needs to satisfy a gradient independent assumption.
One simple check is that the output layer weights are drawn from a zero-mean Gaussian independently from any other parameters and not used anywhere in the interior of the network.
This is clearly satisﬁed in our setting.
In fact, Alemohammad et al.
(2020) has presented the recurrent NTK case with tied weights.
Using their notation, by letting g (1 , 0)
( x ) = 0 ∈ R n , x be T copies of x , and
T = d represents the depth, we exactly recover the current (ﬁnite-depth) DEQ formulation.
See
Figure 1 for a visual explanation.
Therefore, their conclusion directly applies to our setting.
We should emphasize that our work is not a trivial extension to the recurrent NTK, because we mainly study the inﬁnite-depth limit.
3. NTK- OF -DEQ EQUALS DEQ- OF -NTK
Based on Equation (8), we are now ready to show what the DEQ-of-NTK lim
d →∞ Θ ( d ) is.
Then we present the main takeaway of our paper: lim d →∞ Θ ( d ) =
lim n →∞ lim d →∞ Θ ( d ) n .
By convention, we assume the two samples x, y ∈ S d − 1 , and we require the parameters σ 2 W , σ 2 U , σ 2 b obey the following DEQ-NTK initialization:
DEQ-NTK initialization.
Let every entry of every W, U, b, v follows the NTK initialization described in Section 3.1, as well as the additional requirement σ 2 W + σ 2 U + σ 2 b = 1 .
Let the nonlinear activation function σ be the normalized ReLU: σ ( x ) = √ 2 max(0 , x ) from now on.
Using normalized ReLU along with DEQ-NTK initialization, we can derive the main convergence theorem:
Theorem 3.3.
Use same notations and settings in Theorem 3.1, the DEQ-of-NTK is
[Equation]
where Σ ∗ ( x, y ) ≜ ρ ∗ is the root of R σ ( ρ ) − ρ ,
[Equation]
and
[Equation]
Remark 3.4 .
Note our Σ ∗ ( x, y ) always depends on the inputs x and y , so the information between two inputs is always preserved, even if the depth goes to inﬁnity.
On the contrary, as pointed out by Jacot et al.
(2019), without input injection, Σ ( h ) ( x, y ) always converges to 1 as h → ∞ , even if x ∕ = y .
Theorem 3.3 provides us a way to direct calculate the DEQ-of-NTK by using root-ﬁnding algorithms.
In practice, we can solve Equation (10) by using any optimization method.
Then Σ ∗ and Θ ∗ can be computed in constant time.
Since each pair of input ( x, y ) is independent of all the other pairs, we can easily parallelize this computation process.
Our derivation can be extended to more complicated structures like DEQ with convolution layers, see appendix for more detail.
One caveat of Theorem 3.3 is the order of limits, notice that we ﬁrst take the limit of the width, then the limit of the depth.
Nonetheless, with sufﬁcient conditions, one can indeed show that the limits can be exchanged, and the NTK-of-DEQ and the DEQ-of-NTK are equivalent.
Theorem 3.5.
Let σ
[Equation] [Equation]
󰁈 be the empirical NTK with depth d and width n .
Then lim n →∞ lim d →∞ Θ ( d ) n = lim d →∞ lim n →∞ Θ ( d ) n with high probability in probability.
Proof sketch.
We ﬁrst use a well-established random matrix result to conclude that σ 2 W < 1 /guarantees us that σ ◦ 󰁳 σ 2 W /nW is a contraction with high probability.
Using this contraction property, we can then show that the empirical NTK Θ ( d ) n converges.
More importantly, it presents an “uniform convergence” property in n : a larger d does not need a larger n for the limit to converge.
This is the crucial difference between this result and the results in untied-weight network.
Intuitively, suppose contrarily our network has untied weights, to make our proof work we would need every layer’s weight becomes a contraction.
As d increases, this clearly needs larger n to use a union bound, which breaks if d → ∞ .
Finally, we prove a probabilistic version of Moore-Osgood theorem to conclude that our limit exchange result holds.
Remark 3.6 .
In Theorem 3.5, for a ﬁxed depth d , Θ ( d ) :
= lim n →∞ Θ ( d ) n converges almost surely, hence we can view Θ :=
lim d →∞ Θ ( d ) as a constant.
On the other hand, for a ﬁxed n , Θ n := lim d →∞ Θ ( d ) n exists with probability at least 1 − e − c 󰂃 2 n for some constant c , and 󰂃 ≜ 1 − 2 √ 2 σ 2 W √ 2 σ 2 W .
Formally, for any 󰂃 > 0 , we have
[Equation]
which converges in probability by deﬁnition.
Remark 3.7 .
We remark that Theorem 3.5 requires a more stringent σ 2 W than Lemma B.1.
This is indeed expected.
For the actual DEQ to converge, one usually needs I − W ≽ mI for some m > 0 .
It seems that σ 2 W ≤ 1 / 2 exactly reﬂects I − W ≽ 0 , we leave this as an interesting future work.
While Hanin & Nica (2019) also discussed about the relation between width and depth, and they concluded that the NTK may not even be deterministic if d/n ≫ 0 , our result does not contradict with theirs because their n has to depend on d , but our proof decouples the dependency using uniform convergence thanks to weight-tying.
C ASE S TUDY : L INEAR DEQ
Theorem 3.5 shows a quite surprising result that we can safely exchange the limits, which is not at all straightforward to see.
Consider the following linear DEQ case:
[Equation]
Assuming the iteration converges (this can be guaranteed with high probability picking a suitable σ W ).
Equivalently, we can also write this network as
[Equation]
Following the same derivation in Section 3, one can easily see that ˙ Σ ( h ) ( x, y )
=
σ 2 W for all h , and show that lim d →∞ lim n →∞ Θ ( d ) n ( x, y ) =
[Equation]
− σ 2 W .
However, taking the inﬁnite width limit of the network f n ( x ) , it does not obey a Gaussian nature owing to the inverse of a shifted Gaussian matrix.
It is not straightforward to see the limit exchange argument works.
In this section, we aim to solve this linear DEQ case as a sanity check.
In Section 5 we include numerical approximation that indicates the NTK-of DEQ-behaves as we expect.
Theorem 4.1.
Let f n ( x ) be deﬁned as in Equation (14) and Θ ( d ) n be the empirical NTK associated with the ﬁnite depth approximation of f n in Equation (13) .
Let σ 2 W < 1 / 4 and σ 2 W
[Equation]
σ 2 U = 1 .
We have
[Equation]
with high probability.
Proof sketch.
Let H :
= 󰀕
I − 󰁴 σ 2 W n W 󰀖 − 1 .
Such H is well-deﬁned with high probability if σ 2 W < 1 / 4 .
A straightforward derivation gives:
[Equation]
where the ﬁrst convergence happens with high probability (Arora et al., 2019), and the second convergence holds for almost every realization of a sequence of W .
This follows from the weak convergence of probability measure µ n d −→ µ a.s. and Portmanteau lemma, where µ n is the empirical distribution of the eigenvalue of the matrix
[Equation]
󰀖
T 󰀕
I − 󰁴 σ 2 W n W 󰀖 .
More precisely, µ n = 1 n 󰁓 n i =
δ λ i , δ λ
i is the delta measure at the i th eigenvalue λ i .
Next, we show that
[Equation]
From Capitaine & Donati-Martin (2016), we learn that the
Stieltjes transform g of µ is a root to the following cubic equation:
For z
[Equation]
C
[Equation] [Equation]
We then apply the inverse formula of Stieltjes transformation to derive the density
[Equation]
This now involves a one-dimensional integration, which can be computed numerically and shown to be identical to the desired quantity.
Similarly, we can compute that
[Equation]
Summing the three relevant terms and use the fact that σ 2 U + σ 2 W = 1 , we get the claimed result.
S IMULATIONS
freezing of NTK takes more depths to happen.
In this section, we perform numerical simulations on both synthetic data and real-world datasets including MNIST and CIFAR-10 to demonstrate our arguments.
In particular, we show that (a)
The
NTK-of-DEQ and DEQ-of-NTK coincides, for both linear and non-linear cases, (b) A vanilla NTK of FCNN is degenerate while the NTK-of-DEQ escapes the freeze vs. chaos scheme, (c)
The
NTK-of-DEQ delivers reasonable performances on real-world datasets as a further evidence to its nondegeneracy.
5. NTK- OF -DEQ VS DEQ- OF -NTK
Recall in Section 4, the distribution µ in Equation (16) is that of the eigenvalues of H − T H − 1 ≜ ( I − 󰁳 σ 2 W /nW )
T ( I − 󰁳 σ 2 W /nW ) as n → ∞ .
The exact limiting eigenvalue distribution µ when
[Figure]
[Figure] [Figure]
σ 2 W = 0 .
75 is shown in Figure 4a.
Keep in mind that dµ depicts the probability density of how large an eigenvalue of our random matrix can be.
For σ 2 W = 0 .
5 , 0 , 75 we include an empirical eigenvalue distribution of H − T H − 1 ∈ R n × n for n = 1000 in Figure 3.
One can see that the empirical density is sufﬁciently close to the limiting distribution for large enough n , verifying the computation in Equation (16).
We calculated the empirical trace of 1 n tr H T H where H is of size 5000 × 5000 .
This expression is the key element for Equation (15).
The simulation samples H i.i.d 10 times and the results are presented in Figure 4b.
We can see that the variance of the estimator 1 / (1 − σ 2 W ) is negligible for small σ 2 W .
Note that in the proof we require that 󰀐󰀐󰀐󰁳
[Equation]
󰀐󰀐󰀐 < 1 with high probability, which holds when σ 2 W < 1 / 4 .
However, empirically the convergence of empirical trace holds for much larger σ 2 W as well.
We also test the difference between the empirical NTK-of-DEQ Θ n and the DEQ-of-NTK Θ numerically, for both linear DEQ and nonlinear DEQ with normalized ReLU.
We initialize both networks at variable width, with σ 2 v = 2 , σ 2 W = 1 / 8 , and σ 2 U = 7 / 8 .
Θ n is calculated by taking the inner product between the exact gradients 2 of a ﬁnite-width DEQ on two inputs, and Θ is computed using the DEQ-of-NTK formula in Theorem 3.3.
A pair of input ( x, y ) is randomly sampled and ﬁxed throughout the simulation.
For each width n , 10 trials are run, and we draw the mean of log | Θ − Θ n | Θ in Figure 5.
The convergence of the relative residue indicates that the NTK-of-DEQ and the DEQof-NTK coincide as proven.
5. S IMULATIONS ON CIFAR-10 AND MNIST
Hyperparameter sensitivity.
We have three tunable parameters: σ 2 W , σ 2 U , σ 2 b .
We try three random combinations listed in Table 3.
As the results suggest, the performances of NTK-of-DEQ are insensitive to these parameters.
This observation aligns with the description in Lee et al.
[Table]
Training details and results.
For NTK-of-DEQ, following the theory, we normalize the dataset such that each data point has unit length.
The ﬁxed point Σ ∗ ( x, y ) is solved by using the modiﬁed
Powell hybrid method (Powell, 1970).
Notice these root ﬁnding problems are one-dimensional, hence can be quickly solved.
After obtaining the NTK matrix, we apply kernel regressions (without regularization unless stated otherwise).
For any label y ∈ { 1 , . . . , n } , denote its one-hot encoding by e y .
Let 1 ∈ R n be an allvector, we train on the new encoding − 0 .
1 · 1 + e y .
That is, we change the “ 1 ” to 0 .
9 , and the “ 0 ” to − 0 .
1 , as suggested by Novak et al.
The results are listed in Table 3.
These results prove that the NTK-of-DEQ is indeed non-degenerate.
On a smaller dataset with 1000 training data and 100 test data from CIFAR-10, we evaluate the performance of NTK and the ﬁnite depth iteration of NTK-of-DEQ, as depth increases.
See Figure 2.
When the depth increases, the performance of ﬁnite depth NTK gradually drops, eventually to 0.
with 0 standard deviation.
Also with larger σ 2 W , the degeneration of NTK occurs slower.
This shows that large σ 2 W preserves information from previous layers.
Figure 6 also shows that the vanilla NTK becomes independent of the input inner product x T y as the depth increases.
As proven in Jacot et al.
(2019), the NTK will always “freeze” using the sets of parameters in Figure 2.
In this scenario, the
NTK Gram matrix becomes linearly independent as the depth increases, and its kernel regression does not have a unique solution.
To circumvent this unsolvability, we add a regularization term r ∝ 󰂃 Θ ( x,x ) n , where n is the size of the training data.
C ONCLUSION
We derive NTKs for DEQ models, and show that they can be computed efﬁciently via root-ﬁnding based on a limit exchanging argument.
This argument is proven theoretically for non-linear DEQs and an extra sanity check is done on linear DEQs, exploiting random matrix theory.
Numerical simulations are performed to demonstrate that the limit exchanging phenomenon holds for both linear and non-linear NTK-of-DEQs.
Our analysis also shows that one can avoid the freeze and chaos phenomenon in inﬁnitely deep NTKs by using input injection.
Additions experiments are conducted to show that NTK-of-DEQs are non-degenerate on real-world datasets, while ﬁnite depth
NTKs gradually degenerate as their depth increases.
A F ORMAL DERIVATION OF WEIGHT - TIED NETWORK
In this section we formally derive the NTK of a DEQ (weight-tied) model, and show that they converge to the same limit as derived in Section 3.
The argument is nearly identical to that of
Alemohammad et al.
(2020), which heavily depends on the N ESTER ⊤ program (Yang, 2020).
We will ﬁrst give a brief introduction, and then adapt to our setting.
Deﬁnition A.1.
N ESTER ⊤ program is a program (as in type system) of which the variables take three-types: A -vars, G -vars, and H -vars.
Any variables are generated by one of the rules in MatMul (matrix multiplication), NonLin (nonlinearity), LinComb (linear combination), or Trsp (matrix transpose).
We also sometimes explicitly express the dimensionality of a variable in the following way:
• If x ∈ R n , and is of type G , H , we write x : G ( n ) or x : H ( n ) .
• If A ∈ R n × m , we write A : A ( n, m ) .
The program goes as following:
Input A set of G -vars and A -vars.
Body Any variable is introduced by the following rules:
• Trsp .
If A : A ( n, m ) , then A ⊤ : A ( m, n ) .
• MatMul .
If A : A ( n, m ) and x : H ( m ) , then Ax : G ( n ) .
• LinComb .
If g 1 , . .
, g k : G ( n ) and a 1 , . .
. , a k ∈ R , then
󰁓
[Equation]
1 a i g i : G ( n ) .
• NonLin .
If x 1 , . .
, x k : G ( n ) , and φ : R k
→ R is a coordinate-wise nonlinear function, then φ ( x 1 , . . .
, x k ) :
H ( n ) .
Output The program outputs a scalar of the form
[Equation]
󰀔 for h 1 . . .
h k : H ( n ) .
For example, a depth- d approximation to a DEQ model is provided in Algorithm 1.
For simplicity, we left out the scaling σ 2 W / √ n (as was done in Yang (2020)).
Algorithm
1 N ESTER ⊤ program Depth- d approximation to a DEQ model
Require:
Ux, Uy : G ( n ) , W : A ( n, n ) , b : G ( n ) , v : G ( n ) .
Polynomially-bounded coordinatewise nonlinear function φ .
for h = 1 , . . .
, d do for z ∈ { x, y } do f ( h ) ( z ) =
Wg ( h − 1) ( z ) + Uz + b : G ( n ) .
g
[Equation]
The network outputs f
( d +1) ( z ) := v ⊤ g ( d ) ( z ) n , but we don’t express this in the program.
// Backprop, for varible u , let du := √ n ∇ u f ( d +1) ( z ) .
[Equation]
⊲ We use ⊙ for Hadamard product.
[Equation]
end for end for
One can express many neural network architectures into a N ESTER ⊤ program, but not all.
The required regularity condition is the so-called BP-like :
Deﬁnition A.2 (BP-like) .
A N ESTER ⊤ program is BP-like if there exists a non-empty set of input
G ( n ) -vars v 1 , . .
, v k s.t:
If W ⊤ z is used in the program for some z : H ( n ) , and W : A ( n, m ) is an input A
-var, then z must be an odd function of v 1 , . .
. , v k .
That is, z 󰀓
[Equation]
v
[Equation]
, − v k , all other G -vars
[Equation]
all other
[Equation] [Equation]
If Wz is used in the program for some z : H ( m ) , and W : A ( n, m ) is an input A
-var, then z cannot depend on any of v 1 , . . .
v k . 3.
v 1 , . .
. , v k are sampled with zero mean and independently from all other G -vars.
Deﬁnition A.3 (Polynomially-bounded) .
We say a function f : R k
→ R is polynomially-bounded if | φ ( x ) | ≤ C
󰀂 x 󰀂 p + c for some c, C, p > 0 , for all x ∈ R k .
Note that ReLU and inner product are polynoimially-bounded.
Recall that the simple gradient independence assumption (GIA) check we give in Section 3:
Condition A.4 (Simple GIA check) .
Gradient independence assumption is a heuristic that for any matrix W , we assume W ⊤ used in backprop is independet from W used in the forward pass.
We can regard this assumption holds in the NTK computation if the following simple check holds: the output layer is sampled independently with zero mean from all other parameters and it not used anywhere else in the interior of the network, that is, if the output of the network is v ⊤ x , then v is independent of x .
Apparently our DEQ formulation satisfy the simple GIA check, notice that by formulation, the second and third condition in Deﬁnition A.2 are trivially satisﬁed.
Also since v is the last layer weight, any G -var of the form W ⊤ z only shows up in the backpropogation, and is linear (thus odd) in v as well.
Hence the ﬁrst condition is also satisiﬁed.
So any network structure that satisﬁes the simple GIA check is automatically BP-like.
Setup A.5.
For N ESTER ⊤ program, we assume that each entry in W : A ( n, m ) is sampled from
N (0 , σ 2 W /n ) , and any input G -vars x ∼ N ( µ in , Σ in ) .
We remark that this does not contradict with the parameterization that we mentioned in the main text where the entries of input A -vars W, U are standard Gaussians.
One just needs to properly scale their variables.
Theorem A.6 (BP-like N ESTER ⊤ program Master theorem) .
Fix any BP-like N ESTER ⊤ program that satisﬁes Setup A.5, and all its nonlinearities are polynomially-bounded.
If g 1 , . .
, g M are all
G -vars in the program, then for any polynomially-bounded ψ : R M → R , as n → ∞ , we have
[Equation]
where Z = { Z g 1 , . .
, Z g M } ∈ R M , µ = { µ ( g i ) } i ∈
[ M ] ∈ R M , Σ
= { Σ ( g i , g j ) } Mi,j =1 ∈ R M × M are given by µ ( g ) =
󰀻󰀿 󰀽 µ in ( g ) if g is input, 󰁓 ki =
a i µ ( g i ) if g = 󰁓 ki =1 a i g i otherwise Σ
[Equation]
if
[Equation]
g ′ are inputs
[Equation]
We are now equipped to derive the NTK of a depth- d approximation to a DEQ.
Particularly, we have
[Equation]
hence
[Equation]
n .
From this point, we need to calculate
E W 󰀗
[Equation]
󰀘 and E W 󰁫 g
[Equation]
󰁬 .
In the end, applying the Master theorem with ψ ( x, y )
= x ⊤ y on df ( h ) ⊤ df ( l ) n and
[Equation]
n shows that these empirical averages converge to the expectations.
Remark A.7 .
Notice that the Master theorem talks about G -vars, while df ( h ) and g ( h ) are H
-vars.
We can always compose ψ ′ =
ψ ◦ φ , where ψ is the inner product and φ is coordinate-wise nonlinearity (such as ReLU), and apply the Master theorem on ψ ′ , as long as it is still polynomially-bounded.
[Equation] [Equation]
󰁿
[Equation]
By the Master theorem and GIA, φ ′ ( f ( h ) ) and dg ( h ) are introduced by different A -vars ( W and
W ⊤ ), hence their coviance is 0 .
This justiﬁes the last step above.
When h, l < d , by the Master theorem we have
[Equation]
Notice that this gives a recursive expression, WLOG we assume that h < l , this induction will lead to
[Equation]
for some t > 0 .
The reason why this is zero is still due to the Master theorem, as df ( h + t ) ( x ) and df ( d ) ( y ) are G -vars involved with different A -vars W and v .
This shows that when h ∕ = l , E W
[Equation]
󰁬 = 0 .
Hence we only have to consider the case
h = l .
By the Master theorem we have
[Equation]
where
[Equation]
Notice this exactly recovers the calculation of NTK when the weights are un-tied.
The exact same argument can be applied to ∇ U f and ∇ b f .
Since such equivalence holds for all depth d , it also holds in the limit of d → ∞ .
Key takeaway The N ESTER ⊤ program allows us to calculate the NTK of a weight-tied network in exactly the same way as the weight-untied network.
B D ETAILS OF S ECTION
In this section, we give the detailed derivation of DEQ-of-NTK.
There are two terms that are different from NTK: Σ ( h ) ( x, y ) and the extra E θ 󰀗󰁇 ∂ f
( θ ,x ) ∂ U , ∂ f ( θ ,y ) ∂ U 󰁈󰀘 in the kernel.
Let us restate the depth- d approximation to DEQs here:
Let m be the input dimension, x, y ∈ R
m be a pair of inputs, n be the width of the h th hidden layers.
Deﬁne the depth- d approximation to DEQ as follows:
[Equation]
( x
[Equation]
where W ( h ) ∈ R n × n , U ( h ) ∈ R n × m , and v ∈ R n are the internal weights, and b ( h ) ∈ R n are the bias terms.
These parameters are chosen using the NTK initialization.
Let us pick σ W , σ U , σ b ∈ R arbitrarily in this section.
Theorem 3.1.
Recursively deﬁne the following quantities for h ∈ [ d ] :
[Equation]
Then the d -depth iteration to the DEQ-of-NTK can be expressed as:
[Equation]
where by convention we set ˙ Σ ( d +2)
( x, y ) = 1 .
Proof of Theorem 3.1.
First we note that
[Equation]
where the ﬁrst line is by expansion the original expression and using the fact that W, U, b are all independent.
The last line is from the strong law of large numbers.
This shows how the covariance changes as depth increases with input injection.
Recall the splitting:
[Equation]
The following equation has been proven in many places:
[Equation]
and 4 = σ 2 v E ( u,v ) ∼ N (0 , Λ ( h ) )
[ σ ( u ) σ ( v )] .
For instance, see Arora et al.
So we only need to deal with the second term E θ 󰀗󰁇 ∂ f
( θ ,x ) ∂ U , ∂ f ( θ ,y ) ∂ U 󰁈󰀘 .
Write f = f θ ( x ) and ˜ f
= f θ ( y ) , by chain rule, we have
[Equation]
where the last line uses the existing conclusion that 󰁇
[Equation]
x
[Equation]
) , this convergence almost surely holds when N → ∞ by law of large numbers.
Finally, summing
[Equation]
󰁈 over h ∈ [ d ] we conclude the assertion.
Lemma B.1.
Use the same notations and settings in Theorem 3.1.
With input data x, y ∈ S
d − 1 , parameters σ 2 W , σ 2 U , σ 2 b following the DEQ-NTK initialization, Θ ( d ) ( x, y ) in Equation (8) converges absolutely if σ 2 W < 1 .
Proof.
Since we pick x, y ∈ S d − 1 , and by DEQ-NTK initialization, we always have Σ ( h ) ( x, y ) <for x ∕ = y .
Let ρ = Σ ( h ) ( x, y ) , by Equation (5) and Equation (19), if σ 2 W < 1 , then there exists c such that ˙ Σ ( h ) ( x, y ) <
c < 1 for ﬁnite number of pairs x ∕ = y on S d − 1 , and large enough h .
This is because lim
[Equation]
Use comparison test,
[Equation]
, y ) 󰀏󰀏󰀏󰀏
[Equation]
Since c < 1 , the geometric sum converges absolutely, hence Θ ( d ) ( x, y ) converges absolutely if σ 2 W < 1 , and the limit exists.
Theorem 3.3.
Use same notations and settings in Theorem 3.1, the DEQ-of-NTK is
[Equation]
where Σ ∗ ( x, y ) ≜ ρ ∗ is the root of R σ ( ρ ) − ρ ,
[Equation]
and
[Equation]
Proof of Theorem 3.3.
Due to the fact that x ∈ S d − 1 , σ being normalized, and DEQ-NTK initialization, one can easily calculate by induction that for all h ∈
[ L ] : Σ ( h ) ( x, x ) =
σ 2 W E u
[Equation]
N (0 , 1) [ σ ( u ) 2 ]+ σ 2 V x ⊤ x + σ 2 b =
This indicates that in Equation (3), the covariance matrix has a special structure Λ ( h ) ( x, y )
=
󰀕 ρ ρ 1 󰀖 , where ρ = Σ ( h − 1) ( x, y ) depends on h, x, y .
For simplicity we omit the h, x, y in Λ ( h ) ( x, y ) .
As shown in Daniely et al.
[Equation]
[Equation]
Adding input injection and bias, we derive Equation (10) from Equation (18), and similarly, Equation (12) from Equation (19).
Notice that iterating Equations (2) to (4) to solve for Σ ( h ) ( x, y ) is equivalent to iterating ( R σ ◦ · · · ◦ R σ )( ρ ) with initial input ρ = x ⊤ y .
Take the derivative
[Equation]
For x ∕ = y we have − 1 ≤ ρ < c < 1 for some c (this is because we only have ﬁnite number of inputs x, y ) and by DEQ-NTK initialization we have σ 2 W < 1 , so the above inequality hold.
Hence
R σ ( ρ ) is a contraction on [0 , c ] , and we conclude that the ﬁxed point ρ ∗ is attractive.
By Lemma B.1, if σ 2 W < 1 , then the limit of Equation (8) exists, so we can rewrite the summation form in Equation (8) in a recursive form:
[Equation]
Directly solve the ﬁxed point iteration for the internal representation:
[Equation]
Solving for lim d →∞ Θ ( d ) ( x, y ) we get Θ ∗ ( x, y ) =
Σ ∗ ( x,y )− ˙ Σ ∗ ( x,y ) .
Finally, we process the classiﬁcation layer and get Θ
= ˙ Σ · Θ ∗ + Σ , where ˙ Σ =
σ 2 v ˙ ρ ∗ and Σ
= σ 2 v ρ ∗ .This concludes the proof
B. DEQ- OF -NTK VS . NTK- OF -DEQ
In this section we discuss Theorem 3.5 in detail.
Recall that the NTK is the kernel matrix formed by an inﬁnitely-wide network.
To be more precisely, if the network has depth d , then
[Equation]
It is straightforward to deﬁne its width- n approximation:
[Equation]
where θ ( h ) is the parameter of the h th layer with width n .
The name of lim d →∞ lim n →∞ Θ ( d ) n being the DEQ of NTK is intuitive: because we ﬁrst 3 bring width to inﬁnity, that is, the NTK is ﬁrst derived.
Then we talk about the NTK’s inﬁnite-depth limit.
This is in distinction to our desired quantity, lim n →∞ lim d →∞ Θ ( d ) n , which is the NTK of DEQ naturally.
In this section we show they are indeed equivalent under certain conditions.
First we introduce some notations.
Consider a ﬁnite depth iteration of a NTK with depth d + 1 , and for simplicity let the bias term b ( h )
=
for all h ∈ [ d + 1] .
A straightforward calculation show that
[Equation]
( θ , x ) dU (
[Equation]
=
[Equation]
⊤ where p ( h ) ( x ) =
󰀻󰁁󰀿 󰁁󰀽 1 ∈ R n , h =
d +󰁴
[Equation] [Equation]
N h diag 󰀕
[Equation]
σ
[Equation]
f
[Equation]
h )
[Equation]
x ) 󰀔󰀖󰀓
W ( h +
[Equation]
󰀔
[Equation]
p
[Equation]
h +
[Equation]
x
[Equation]
≤ d
Here diag 󰀕 ˙
[Equation]
󰀓 f
[Equation] [Equation]
󰀔󰀖 ∈ R
N h × N h .
Let N h
= n for all h , and W ( h +1) :
= v .
Notice that diag
[Equation]
and we use these terms interchangeably.
For simplicity, we omit all the x in the terms and write f ( h ) :
= f ( h ) ( x ) , etc.
Write ˙ σ ( h ) = ˙
σ 󰀓 f ( h ) ( x ) 󰀔 .
Notice that applying σ ( · ) or Hadamard product with ˙ σ ( h ) only decreases norms.
Lemma B.2 (Probablisitc Moore-Osgood for double sequence) .
Let a n,d be a random double sequence in a complete space.
Assume for any 󰂃 > 0 , δ ∈ (0 , 1) , there exists N ( δ ) >
and D ( 󰂃 ) >such that for all n > N and d > D , with probability at least 1 −
δ we have | a n,d − a n | < 󰂃
(we may refer to this property as uniform convergence with high probability).
And for any d ∈ N we have
[Equation]
n
[Equation]
= a d almost surely, then with high probability:
[Equation]
Proof.
We sometimes also write a d ( n ) to stress that we consider the sequence as a function of n .
By assumption, for any δ ∈ (0 , 1) , 󰂃 > 0 , there exists N, D such that for all n > N , d, e > D ,
| a d ( n ) − a e ( n ) | < 󰂃 with probability at least 1 − δ .
Since here N does not depend on D , let n → ∞ we get the following statement holds almost surely:
[Equation]
This shows that a d := lim n →∞ a n,d is a Cauchy sequence and have a ﬁnite limit lim d →∞ a d = L .
Now deﬁne a ( n ) :
= a n = lim d →∞ a n,d , for d > D ( 󰂃 ) :
[Equation]
By assumption, pick large enough n , we have A < 󰂃 with probability at least 1 − δ .
By the
Cauchy sequence argument above, we have C < 󰂃 with high probability.
Finally since a d ( n ) → a d pointwise for every d , we can choose n large enough such that B < 󰂃 .
This concludes our proof.
We want to remark that the above Lemma B.2 relies on a more general notion of “conditional almost sure convergence”.
In particular, we only assume that | a n,d − a n | < 󰂃 almost surely conditioned on an event with probability at least
[Equation]
E 󰀖 = 1 , where P ( E ) > 1 − δ for all large enough n.
Notice here we are not explicit about how δ evolves with n .
When we use this lemma in Theorem 3.5, we have δ =
o ( n ) which will instead gives us a convergence in probability result.
To be complete, we also provide the weaker result and its proof here.
Lemma B.3 (Another probablisitc Moore-Osgood for double sequence) .
Let a n,d be a random double sequence in a complete space.
Assume for any 󰂃 > 0 , there exists D ( 󰂃 ) > 0 such that for all d > D , with probability at least 1 − o ( n ) we have | a n,d − a n | < 󰂃 .
And for any d ∈ N we have lim n →∞ a n,d = a d almost surely, then the following convergence holds in probability:
[Equation]
Proof.
We sometimes also write a d ( n ) to stress that we consider the sequence as a function of n .
By assumption, let n → ∞ we get the following statement holds with probability 1 :
[Equation]
This shows that a d := lim n →∞ a n,d is a Cauchy sequence and have a ﬁnite limit lim d →∞ a d = L .
Now deﬁne a ( n ) :
= a n = lim d →∞ a n,d , for d > D ( 󰂃 ) :
[Equation]
By assumption, pick large enough n , we have A < 󰂃 with probability at least 1 − o ( n ) .
By the
Cauchy sequence argument above, we have C < 󰂃 with probability 1 .
Finally since a d ( n ) → a d pointwise for every d , we can choose n large enough such that B < 󰂃 with probability at least 1 − o ( n ) .
Overall this gives
[Equation]
which concludes our proof
By standard high-dimensional probability (Vershynin, 2019), the following lemma holds:
Lemma B.4.
Let A ∈ R n × m be a random matrix whose entries are sampled from i.i.d standard
Gaussian distribution, then for t ≥ 0 , with probability at least 1 − e − ct 2 for a constant c > 0 , there is:
[Equation]
We are now ready to give the formal proof.
Theorem 3.5.
Let σ 2 W ≤ 1 / 8 , Θ (
[Equation]
n
[Equation]
x, y
[Equation]
=
[Equation]
󰁈 be the empirical NTK with depth d and width n .
Then lim n →∞ lim d →∞ Θ ( d ) n = lim d →∞ lim n →∞ Θ ( d ) n with high probability in probability.
Proof of Theorem 3.5.
For any ﬁxed d , we write Θ ( d ) =
lim n →∞ Θ ( d ) n , notice this is just a ﬁnitedepth NTK (possibly with input injection).
We condition on the event that lim d Θ ( d ) n exists.
A sufﬁcient condition for this event to hold with high probability is σ 2 W < 1 / 8 .
With such σ 2 W , by
Lemma B.4, σ ◦ 󰁳 σ 2 W /nW has a Lipschitz constant less than 1 with high probability.
Recall that σ ( x ) = √ 2 max { 0 , x } is the normalized ReLU nonlinearity.
Conditioned on such event, we have
[Equation]
WLOG let g (0) =
x ∈ S
d − 1 , and 󰀂 g (0) 󰀂 ≤ 1 be our base case.
Note that U ( h ) x is ﬁxed for weighttied network, let’s denote it as C , and also overload the notation that 󰀂 C 󰀂 = C .
By induction:
[Equation]
By Lemma B.4, with probabiliy at least 1 − e − O ( t 2 ) , we have 󰀂 W 󰀂 op ≤ 2 √ n + t .
This shows that for all 󰂃 > 0 , let σ W < 1 2 √ 2+ 󰂃 , with probability at least 1 − e − O
( 󰂃 2 n ) , we have
[Equation]
Consequently:
[Equation]
which is geometric and converges absolutely as h → ∞ .
Therefore, there exists a constant Q >s.t 󰀂 g ( h ) 󰀂 < Q for all h ∈ N .
By the same spirit, using induction, we have
[Equation]
Combining the above two derivations, we have
[Equation]
Similar convergence result can be derived for dfdU as well.
Use the terminology introduced in Lemma B.2, lim d →∞ Θ ( d ) n
[Equation]
󰁓 ∞ h =1 ∂ f ( x ) ∂θ ( h ) T ∂ f ( x ′ )
∂θ ( h ) converges uniformly in n with high probability.
For a ﬁxed d , we know that lim n →∞ Θ ( d ) n = Θ ( d ) by the tensor program (Yang, 2019).
Therefore conditioned on the event that σ ◦ 󰁳 σ 2 W /nW has a Lipschitz constant less than 1 , by Lemma B.2, we can swap the limit and indeed lim
d →∞ lim n →∞ Θ ( d ) n = lim n →∞ lim d →∞ Θ ( d ) n .
This shows that the NTK-of-DEQ and the DEQ-of-NTK coincide.
One should note that it merely requires σ 2 W < 1 for the DEQ-of-NTK to converge as in Theorem 3.3, but our above proof requires σ 2 W < 1 / 8 to make sure that the NTK-of-DEQ and DEQ-of-NTK are equivalent.
Our current analysis relies heavily on a contraction argument.
However, in the actual
DEQ setting, it sufﬁce to have W being strongly monotone to guarantee convergence.
That is, one only needs the largest eigenvalue of W to be less than 1 .
This corresponds to have σ 2 W < 1 / 2 (again, this is because we use the normalized ReLU, so there is an extra factor of √ 2 ) by the semicircular law.
We leave the gap to future works.
C D ETAILS OF S ECTION
Theorem 4.1.
Let f n ( x ) be deﬁned as in Equation (14) and Θ ( d ) n be the empirical NTK associated with the ﬁnite depth approximation of f n in Equation (13) .
Let σ 2 W < 1 / 4 and σ 2 W + σ 2 U
= 1 .
We have
[Equation]
with high probability.
Proof of Theorem 4.1.
Recall that we deﬁne H :
= 󰀕
I
[Equation]
󰁴
[Equation]
n
[Equation]
󰀖
[Equation]
This inverse matrix exists with high probability if σ 2 W < 1 / 4 , due to a well-known random matrix theory result Lemma B.4.
straightforward derivation gives:
[Equation]
The ﬁrst convergence happens with high probability (Arora et al., 2019).
Note that B = E U,v [ A ] .
One needs to apply the Gaussian chaos of order 2 lemma (Boucheron et al., 2013) to show the concentration.
This was done rigorously down in Arora et al.
(2019) Claim E.2.
Their proof works for our case as well since we have 󰀂 H T H 󰀂 2 bounded independently of n and d with high probability.
The second convergence holds for almost every realization of a sequence of W .
Recall that µ n is the empirical distribution of the eigenvalue of the matrix
[Equation]
More precisely, µ n = 1 n 󰁓 ni =
δ λ i , δ λ
i is the delta measure at the i th eigenvlue value λ i .
We can rewrite
[Equation]
We will show that µ n → µ weakly a.s 4 .
Then by Portmanteau lemma, we have 󰁕
[Equation]
󰁕
[Equation]
for every bounded Lipschitz function.
Here we have f = 1 / λ deﬁned when λ has non-zero support in µ ( λ ) .
Since by Lemma B.4, our assumption σ 2 W < 1 / 8 guarantees 󰀐󰀐󰀐󰀐󰁴
[Equation]
󰀐󰀐󰀐󰀐 < 1 w.h.p, the support of µ ( λ ) is bounded away from 0 , and f is indeed Lipschitz and bounded on its domain.
Next, we show that 󰁕
[Equation]
From Capitaine & Donati-Martin (2016), we learn that the
Stieltjes transform g of µ is a root to the following cubic equation:
For z ∈ C + : g
[Equation]
z
[Equation]
Deducing the probability density from g by using the inverse formula of Stieltjes transformation, we have
[Equation]
Finally we can compute 󰁕 u l 1 λ p ( λ ) d λ .
Notice to let p ( · ) be well deﬁned, we need 9 σ 8 W b
[Equation]
−
[Equation]
σ
[Equation]
W b
[Equation]
+ 18 σ 6 W b 2 ≥ 0 , which amounts to l =
[Equation]
󰁳
[Equation]
󰀔 and u =󰀓 − σ 4 W + 20 σ 2 W + 󰁳 σ 8 W + 24 σ 6 W + 192 σ 4 W +
[Equation]
󰀔 .
This now involves a onedimensional integral, which an be solved numerically for all values of σ W , and shown be be arbitrarily close the desired quantity 1 / (1 − σ 2 W ) .
Similarly, we can compute that
[Equation]
and
[Equation]
Summing the three relevant terms and use the fact that σ 2 U + σ 2 W = 1 , we get the claimed result.
D DEQ WITH C ONVOLUTION L AYERS
In this section we show how to derive the NTKs for convolution DEQs (CDEQ).
Although in this paper only the CDEQ with vanilla convolution structure is considered, we remark that our derivation is general enough for other CDEQ structures as well, for instance, CDEQ with global pooling layer.
The details of this section can be found in the appendix.
Unlike the fully connection network with input injection, whose intermediate NTK representation is a real number.
For convolutional neural networks (CNN), the intermediate NTK representation is a four-way tensor.
In the following, we will present the notations, CNN with input injection (CNN-IJ) formulation, the CDEQ-NTK initialization, and our main theorem.
Notation.
We adopt the notations from Arora et al.
Let x, y ∈ R P × Q be a pair of inputs, let
q ∈ Z + be the ﬁlter size (WLOG assume it is odd as well).
By convention, we always pad the representation (both the input layer and hidden layer) with 0 ’s. Denote the convolution operation for i ∈
[ P ] , j ∈ [ Q ] : [ w ∗ x ]
[Equation]
󰁓
[Equation]
Denote
D
[Equation]
󰁱󰀃
[Equation]
󰀄
[Equation]
󰁲 .
Intuitively, D ij,i ′
j
′ is a q × q × q × q set of indices centered at ( ij, i ′ j ′ ) .
For any tensor T ∈
R P × Q × P × Q , let [ T ] D ij,i ′
j
′ be the natural sub-tensor and let Tr( T ) =
󰁓 i,j T ij,ij .
Formulation of CNN-IJ.
Deﬁne the CNN-IJ as follows:
• Let the input x (0) =
x ∈ R P × Q × C 0 , where C 0 is the number of input channels, and C h is the number of channels in layer h .
Assume WLOG that C h = C for all h ∈ [ d ] • For h = 1 , . . .
, d , let the inner representation
[Equation]
where W ( h ) ( α ) , ( β ) ∈ R q × q represent the convolution operator from the α th channel in layer h −to the β th channel in layer h .
Similarly, U ( h ) ( α ) , ( β ) ∈ R q × q injects the input in each convolution window.
S ∈ R P × Q is a
normalization
matrix
Let W, U, S,
σ
2 U , σ 2 W be chosen by the CDEQNTK initialization described later.
• The ﬁnal output is deﬁned to be f θ ( x ) =
[Equation]
, where W ( d +1) ( α ) ∈ R P × Q is sampled from standard Gaussian distribution.
CDEQ-NTK initialization.
Let 1 q ∈ R q × q , X ∈ R P × Q be two all-one matrices.
Let ˜ X ∈
R ( P +2) × ( Q +2) be the output
of
zero padding X .
We
index
the rows of ˜ X by { 0 , 1 , . .
, P + 1 } and columns by { 0 , 1 , . .
, Q + 1 } .
For position i ∈ [ P ] , j ∈ [ Q ] , let 󰀃 [ S ] ij 󰀄
[Equation]
=
[Equation]
q ∗ ˜ X ] ij in
Equation (23).
Let every entry of every W, U be sampled from N (0 , 1) and σ 2 W + σ 2 U = 1 .
Using the above-deﬁned notations, we now state the CDEQ-NTK.
Theorem D.1.
Let x, y ∈ R P × Q × C 0 be s.t 󰀂 x ij 󰀂 2 = 󰀂
y ij 󰀂 2 = 1 for i ∈
[ P ] , j ∈ [ Q ] .
Deﬁne the following expressions recursively (some x, y are omitted in the notations), for ( i, j, i ′ , j ′ ) ∈
[ P ] × [ Q ]
[Equation]
[ P ]
[Equation]
[ Q ]
[Equation] [Equation]
[Equation]
Deﬁne the linear operator L : R P × Q × P × Q → R P × Q × P × Q via [ L ( M )] ij,i ′ j ′
=
Tr 󰀓 [ M ] D ij,i ′ j ′ 󰀔 .
Then the CDEQ-NTK can be found solving the following linear system:
[Equation]
where K ∗ ( x, y ) =
lim d →∞ K ( L ) ( x, y ) , ˙ K ∗ ( x, y ) =
lim d →∞ ˙ K ( d ) ( x, y ) .
The limit exists if σ 2 W < 1 .
The actual NTK entry is calculated by Tr( Θ ∗ ( x, y )) .
Theorem D.1 highlights that the convergence of CDEQ-NTK depends solely on the CDEQ-NTK initialization.
The crucial factor here is the normalization tensor S , which guarantees the variance of each term is always 1 across the propogation.
This idea mimics that of the DEQ-NTK initialization.
Our theorem shows that CDEQ-NTK can also be computed by solving ﬁxed point equations.
We ﬁrst explain the choice of S in the CDEQ-NTK initialization.
In the original CNTK paper (Arora et al., 2019), the normalization is simply 1 /q 2 .
However, due to the zero-padding, 1 /q 2 does not normalize all
[Equation]
as expected: only the variances that are away from the corners are normalized to 1 , but the ones near the corner are not.
[ S ] ij is simply the number of non-zero entries in 󰁫 ˜
X
󰁬 D ij,ij .
Now we give the proof to Theorem D.1.
Proof of Theorem D.1.
Similar to the proof of Theorem 3.1, we can split the CDEQ-NTK in two terms:
[Equation]
Omit the input symbols x, y , let
[Equation]
As shown in Arora et al.
(2019), we have
[Equation]
Write f = f θ ( x ) and ˜ f
= f θ ( y ) .
Following the same step, by chain rule, we have 󰀭
[Equation]
󰁄 Rewrite the above two equations in recursive form, we can calculate the L -depth iteration of CDEQNTK by:
• For the ﬁrst layer Θ (0) ( x, y ) =
Σ (0) ( x, y ) .
• For h = 1 , . . .
, d − 1 , let 󰁫
[Equation]
• For h = d , let
[Equation]
• The ﬁnal kernel value is Tr( Θ ( d ) ( x, y )) .
Using Equation (31) and Equation (32), we can ﬁnd the following recursive relation:
[Equation]
The rest of the proof is stated in the main text.
For readers’ convenience we include them here again.
At this point, we need to show that K ∗ ( x, y ) ≜
lim d →∞ K ( d ) ( x, y ) and ˙ K ∗ ( x, y ) ≜
lim d →∞ ˙ K ( d ) ( x, y ) exist.
Let us ﬁrst agree that for all h ∈ [ d ] , ( ij, i ′ j ′ ) ∈
[ P ] ×
[ Q ] × [ P ] × [ Q ] , the diagonal entries of Λ ( h ) ij,i ′
j
′ are all ones.
Indeed, these diagonal entries are 1 ’s at h = 0 by initialization.
Note that iterating Equations (26) to (29) to solve for [ Σ ( h ) ( x, y )] ij,i ′ j ′ is equivalent to iterating f : R P × Q × P × Q → R P × Q × P × Q :
[Equation]
where
[Equation]
is applied to P ( h ) entrywise.
Due to CDEQ-NTK initialization, if P (0) ij,ij = 1 for i ∈
[ P ] , j ∈ [ Q ] , then P ( h ) ij,ij = 1 for all iterations h .
This is true by the deﬁnition of S .
Now if we can show f is a contraction, then Σ ∗ ( x, y ) ≜ lim h →∞ Σ ( h ) ( x, y ) exists, hence K ∗ and ˙
K ∗ also exist.
We should keep the readers aware that f : R P × Q × P × Q → R P × Q × P × Q , so we should be careful with the metric spaces.
We want every entry of Σ ( h ) ( x, y ) to converge, since this tensor has ﬁnitely many entries, this is equivalent to say its ℓ ∞ norm (imagine ﬂattenning this tensor into a vector) converges.
So we can equip the domain an co-domain of f with ℓ ∞ norm (though these are ﬁnite-dimensional spaces so we can really equip them with any norm, but picking ℓ ∞ norm makes the proof easy).
Now we have f =
L ◦ 1 [ S ] ij [ S ] i ′ j ′ R σ : ℓ ∞ → ℓ ∞ .
If we ﬂatten the four-way tensor P ( h ) into a vector, then L can be represented by a ( P × Q × P × Q ) ×
( P × Q × P × Q ) dimensional matrix, whose ( kl, k ′ l ′ ) -th entry in the ( ij, i ′ j ′ )
-th row is 1 if ( kl, k ′ l ′ )
∈ D ij,i ′ j ′ , and 0 otherwise.
In other words, the ℓ 1 norm of the ( ij, i ′ j ′ )
-th row represents the number of non-zero entries in D ij,i ′ j ′ , but by the CDEQ-NTK initialization, the row ℓ 1 norm divided by [ S ] ij · [ S ] i ′ j ′ is at most 1 !
Using the fact that 󰀂 L 󰀂 ℓ ∞ → ℓ ∞ is the maximum ℓ 1 norm of the row, and the fact R σ is a contraction (proven in Theorem 3.3), we conclude that f is indeed a contraction.
With the same spirit, we can also show that Equation (32) is a contraction if σ 2 W < 1 , hence Equation (30) is indeed the unique ﬁxed point.
This ﬁnishes the proof.
D. C OMPUTATION OF CDEQ-NTK
One may wish to directly compute a ﬁxed point (or more precisely, a ﬁxed tensor) of Θ ( d ) ∈
R P × Q × P × Q like Equation (10).
However, due to the linear operator L (which is just the ensemble of the trace operator in Equation (29)), the entries depend on each other.
Hence the system involves a ( P × Q × P × Q ) ×
( P × Q × P × Q ) -dimensional matrix that represents L .
Even if we exploit the fact that only entries on the same “diagonal” depend on each other,
L is at least P × Q × P × Q , which is 32 4 for CIFAR-10 data.
Moreover, this system is nonlinear.
Therefore we cannot compute the ﬁxed point Σ ∗ by root-ﬁnding efﬁciently.
Instead, we approximate it using ﬁnite depth iterations, and we observe that in experiments they typically converge to 10 − 6 accuracy in ℓ ∞ within 15 iterations. [Table]
[Table]
We test CDEQ-NTK accuracy on CIFAR-10 dataset with just 2000 training data.
The result is shown in Table
[Table]
[Table]