Year: 2020

Abstract:
Existing analyses of the neural tangent kernel (NTK) for infinite-depth networks show that the kernel typically becomes degenerate as the number of layers grows. This raises the question of how to apply such methods to practical “infinite depth” architectures such as the recently-proposed deep equilibrium (DEQ) model, which directly computes the infinite-depth limit of a weight-tied network via rootfinding. In this work, we show that because of the input injection component of these networks, DEQ models have non-degenerate NTKs even in the infinite depth limit. Furthermore, we show that these kernels themselves can be computed by an analogous root-finding problem as in traditional DEQs, and highlight methods for computing the NTK for both fully-connected and convolutional variants. We evaluate these models empirically, showing they match or improve upon the performance of existing regularized NTK methods.

text: 1 INTRODUCTION
Recent works empirically observe that as the depth of a weight-tied input-injected network increases, its output tends to converge to a fixed point. Motivated by this phenomenon, DEQ models were proposed to effectively represent an “infinite depth” network by root-finding. A natural question to ask is, what will DEQs become if their widths also go to infinity? It is well-known that at certain random initialization, neural networks of various structures converge to Gaussian processes as their widths go to infinity (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 infinite 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 (Arora et al., 2019; Jacot et al., 2018; Yang, 2019; Huang et al., 2020).
However, as the depth goes to infinity, Jacot et al. (2019) showed that the NTKs of fully-connected neural networks (FCNN) converge either to a constant (freeze), or to the Kronecker Delta (chaos). In this work, we show that with input injection, the DEQ-NTKs converge to meaningful fixed points that depend on the input in a non-trivial way, thus avoiding both freeze and chaos. Furthermore, analogous to DEQ models, we can compute these kernels by solving an analogous fixed point equation, rather than simply iteratively applying the updates associated with the traditional NTK. Moreover, such derivations carry over to other structures like convolution DEQs (CDEQ) as well. We evaluate the approach and demonstrate that it typically matches or improves upon the performance of existing regularized NTK methods.
2 BACKGROUND AND PRELIMINARIES
Bai et al. (2019) proposed the DEQ model, which is equivalent to running an infinite depth network with tied weight and input injection. These methods trace back to some of the original work in recurrent backpropagation (Almeida, 1990; Pineda, 1988), but with specific emphasis on: 1) computing the fixed point directly via root-finding 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. The DEQ algorithm finds the infinite depth fixed point using quasi-Newton root finding methods, and then backpropagates using implicit differentiation without storing the derivatives in the intermediate layers, thus achieving a constant
memory complexity. Furthermore, although a traditional DEQ model does not always guarantee to find a stable fixed point, with careful parameterization and update method, monotone operator DEQs can ensure the existence of a unique stable fixed point (Winston & Kolter, 2020).
On the side of connecting neural networks to kernel methods, Neal (1996) first 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 one can imagine these networks as having fixed parameters throughout the training process, except for the last classification layer. Following the naming convention of Arora et al. (2019), we call these networks weakly-trained, and networks where every layer is updated are called fully-trained. Weakly-trained nets induce the kernel Θ(x, y) = Eθ∼N [f(θ, x) · f(θ, y)], where x, y ∈ Rd are two samples, θ represents the parameters of the network, N is the initialization distribution (often Gaussian) over θ, and f(θ, ·) ∈ R is the output of the network. One related topic in studying the relation between Gaussian process kernel and depth is the meanfield theory. Poole et al. (2016); Schoenholz et al. (2016) showed that the correlations between all inputs on an infinitely wide weakly-trained net become either perfectly correlated (order) or decorrelated (chaos) as depth increases. This aligns with the observation in Jacot et al. (2019). 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-field 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. Such structure guarantees the converged nets depend nontrivially on the inputs, as we will see later in the experiments.
It can be unsatisfying that the previous results only involve weakly-trained nets. Interestingly, similar 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 infinite width network is the following:
Θ(x, y) = Eθ∼N
[〈 ∂f(θ, x)
∂θ , ∂f(θ, y) ∂θ
〉] . (1)
They also gave a recursive formulation for the NTK of FCNN. Arora et al. (2019); Yang (2020) later provided formulation for convolution NTK and other structures.
One may ask what happens if both the width and the depth go to infinity. It turns out that the vanilla FCNN does not have a meaningful convergence: either it gives constant kernels or Kronecker Delta kernels (Jacot et al., 2019). On the bright side, this assertion is not always the case for other network structures. For example, the NTK induced by ResNet provides a meaningful fixed point in the large depth limit (Huang et al., 2020). This may seem to give one explanation why ResNet outperforms FCNN, but unfortunately they also show that the ResNet NTK with infinite depth is no different from the ResNet NTK with just depth one. This conclusion makes the significance of infinite depth questionable.
2.1 NOTATIONS
Throughout the paper, we write θ as the parameters for some network fθ or equivalently, f(θ, ·). 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) = √ 2max(0, x) as the normalized ReLU and σ̇ its derivative (which only need to be well-defined almost everywhere). The symbol σ2a with subscript is always used to denote the variance of some distribution. We write N (µ,Σ) as the Gaussian distribution with mean µ ∈ Rd and covariance matrix Σ ∈ Rd×d. We let Sd−1 be the unit sphere embedded in Rd.
3 DEQ-NTK OF FULLY-CONNECTED NEURAL NETWORKS
In this section, we show how to derive the NTK of the fully-connected DEQ (DEQ-NTK). From now on, we simplify fully-connected DEQs as DEQs. Recall that DEQs are equivalent to infinitely deep fully-connected neural nets with input injection (FCNN-IJ), and one can either exactly solve the fixed point using root-finding (up to machine epsilon and root-finding algorithm accuracy) or approximate the DEQs by just doing finite depth forward iterations. In section 3.1, we show the NTK of the approximated DEQ using finite depth iteration, and in section 3.2 we demonstrate how to get the exact convergence point of DEQ-NTK. The details of this section can be found in appendix A.
3.1 FINITE DEPTH ITERATION OF DEQ-NTK
Let d be the input dimension, x, y ∈ Rd be a pair of inputs, Nh be the width of the hth layers where h ∈ [L+ 1]. Let N0 = d and NL+1 = 1. Define the FCNN-IJ with L hidden layers as follows:
f (h) θ (x) = √ σ2W Nh W (h)g(h−1)(x) + √ σ2U Nh U (h)x+ √ σ2b Nh b(h), h ∈ [L+ 1]
g(L)(x) = σ(f (L) θ (x))
where W (h) ∈ RNh×Nh−1 , U (h) ∈ RNh×d are the internal weights, and b(h) ∈ RNh are the bias terms. These parameters are chosen using the NTK initialization. Let us pick σW ,σU ,σb ∈ R arbitrarily in this section.
NTK initialization. We randomly initialize every entry of every W,U, b from N (0, 1). Without loss of generality (WLOG) we assume the width of the hidden layer Nh = N is the same across different layers. We remark the readers to distinguish FCNN-IJ from a recurrent neural network (RNN): our model injects the original input to each layer, whereas a RNN has a sequence of input (x1, . . . , xT ), and inject xt to the tth-layer.
Here is a crucial distinction between finite width DEQs and infinite width DEQs:
Remark 1. In the finite width regime, one typically has to assume the DEQs have tied weights, that is, W (1) = . . . = W (L+1). Otherwise it is unlikely the network will converge at all. In fact, one needs to be very careful with the parametrization of the weights to guarantee the fixed point is unique and stable. This is not the case in the infinite width regime. As we shall see soon, even with distinct weights in each layer, the convergence of DEQ-NTKs only depend on σ2W ,σ 2 U ,σ 2 b , and the nonlinearity σ. Assuming untied weights makes the analysis easier, but the same argument can be made rigorously for tied weights as well, see Yang (2019; 2020).
Our main theorem is the following:
Theorem 1. Recursively define the following quantities for h ∈ [L]:
Σ(0)(x, y) = x⊤y (2)
Λ(h)(x, y) = ( Σ(h−1)(x, x) Σ(h−1)(x, y) Σ(h−1)(y, x) Σ(h−1)(y, y) ) ∈ R2×2 (3)
Σ(h)(x, y) = σ2WE(u,v)∼N (0,Λ(h))[σ(u)σ(v)] + σ2Ux⊤y + σ2b (4)
Σ̇(h)(x, y) = σ2WE(u,v)∼N (0,Λ(h))[σ̇(u)σ̇(v)] (5)
Then the L-depth iteration to the DEQ-NTK can be expressed as:
Θ(L)(x, y) =
L+1∑
h=1
  ( Σ(h−1) (x, y) ) · L+1∏
h′=h
Σ̇(h ′) (x, y)

 , (6)
where by convention we set Σ̇L+1(x, y) = 1 for the L-depth iteration.
Proof Sketch. The first step is to show that at each layer h ∈ [L], the representation f (h)θ (x) is associated with a Gaussian process with kernel eq. (3) as N → ∞. Then use the characterization in eq. (1), calculate the NTK by:
Θ(L)(x, y) = Eθ
[〈 ∂f(θ, x)
∂θ , ∂f(θ, y) ∂θ
〉]
=Eθ
[〈 ∂f(θ, x)
∂W , ∂f(θ, y) ∂W
〉]
︸ ︷︷ ︸ 1
+Eθ
[〈 ∂f(θ, x)
∂U , ∂f(θ, y) ∂U
〉]
︸ ︷︷ ︸ 2
+Eθ
[〈 ∂f(θ, x)
∂b , ∂f(θ, y) ∂b
〉]
︸ ︷︷ ︸ 3
.
Calculating each term using the chain rule, we get eq. (6).
3.2 FIXED POINT OF DEQ-NTK
Based on eq. (6), we are now ready to answer what the fixed point of Θ(L) is. By convention, we assume the two samples x, y ∈ Sd−1, and we require the parameters σ2W ,σ2U ,σ2b obey the DEQNTK initialization.
DEQ-NTK initialization. Let every entry of every W,U, b follows the NTK initialization described in section 3.1, as well as the additional requirement σ2W + σ 2 U + σ 2 b = 1. Let the nonlinear activation function σ be the normalized ReLU: σ(x) = √ 2max(0, x).
Definition 3.1 (Normalized activation). We call an activation function σ : R → R normalized if Ex∼N (0,1)[σ(x)2] = 1.
Using normalized activations along with DEQ-NTK initialization, we can derive the main convergence theorem: Theorem 2. Use same notations and settings in theorem 1. With input data x, y ∈ Sd−1, parameters σ2W ,σ 2 U ,σ 2 b follow the DEQ-NTK initialization, the fixed point of DEQ-NTK is
Θ∗(x, y) ≜ lim L→∞
Θ(L)(x, y) = Σ∗(x, y)
1− Σ̇∗(x, y) , (7)
where Σ∗(x, y) ≜ ρ∗ is the root of:
Rσ(ρ)− ρ, where Rσ(ρ) ≜ σ2W
(√ 1− ρ2 + ( π − cos−1 ρ ) ρ
π
) + σ2Ux ⊤y + σ2b , (8)
and
Σ̇∗(x, y) ≜ lim h→∞ Σ̇(h)(x, y) = σ2W
( π − cos−1(ρ∗)
π
) . (9)
Proof. Due to the fact that x ∈ Sd−1, σ being a normalized activation, and DEQ-NTK initialization, one can easily calculate by induction that for all h ∈ [L]:
Σ(h)(x, x) = σ2W E u∼N (0,1) [σ(u)2] + σ2V x ⊤x+ σ2b = 1
This indicates that in eq. (3), the covariance matrix has a special structure Λ(h)(x, y) = ( 1 ρ ρ 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. (2016):
E (u,v)∼N (0,Λ) [σ(u)σ(v)] =
√ 1− ρ2 + ( π − cos−1(ρ) ) ρ
π (10)
E (u,v)∼N (0,Λ)
[σ̇(u)σ̇(v)] = π − cos−1(ρ)
π (11)
Adding input injection and bias, we derive eq. (8) from eq. (10), and similarly, eq. (9) from eq. (11). Notice that iterating eqs. (2) to (4) to solve for Σ(h)(x, y) is equivalent to iterating (Rσ ◦ · · ·◦Rσ)(ρ) with initial input ρ = x⊤y. Take the derivative
∣∣∣∣ dRσ(ρ)
dρ
∣∣∣∣ = ∣∣∣∣∣∣ σ2W ( 1− cos −1(ρ) π )∣∣∣∣∣∣ < 1, if σ2W < 1 and −1 ≤ ρ < 1.
For x ∕= y we have −1 ≤ ρ < c < 1 for some c (this is because we only have finite number of inputs x, y) and by DEQ-NTK initialization we have σ2W < 1, so the above inequality hold. Hence Rσ(ρ) is a contraction on [0, c], and we conclude that the fixed point ρ∗ is attractive.
By lemma 1, if σ2W < 1, then the limit of eq. (6) exists, so we can rewrite the summation form in eq. (6) in a recursive form:
Θ(0)(x, y) = Σ(0)(x, y)
Θ(L+1)(x, y) = Σ̇(L+1)(x, y) ·Θ(L)(x, y) + Σ(L+1)(x, y), and directly solve the fixed point iteration:
lim L→∞ Θ(L+1)(x, y) = lim L→∞
( Σ̇(L+1)(x, y) ·Θ(L)(x, y) + Σ(L+1)(x, y) )
=⇒ lim L→∞ Θ(L+1)(x, y) = Σ̇∗(x, y) · lim L→∞ Θ(L)(x, y) + Σ∗(x, y)
=⇒ lim L→∞ Θ(L)(x, y) = Σ̇∗(x, y) · lim L→∞ Θ(L)(x, y) + Σ∗(x, y).
(12)
Solving for limL→∞ Θ(L)(x, y) we get Θ∗(x, y) = Σ∗(x,y)
1−Σ̇∗(x,y) .
Remark 2. 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 infinity. 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.
4 DEQ WITH CONVOLUTION LAYERS
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 in experiments, we remark that our derivation is general enough for other CDEQ structure as well, for instance, CDEQ with global pooling layer. The details of this section can be found in appendix B.
Unlike the FCNN-IJ, whose intermediate NTK representation is a real number. For convolution 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. (2019). Let x, y ∈ RP×Q be a pair of inputs, let q ∈ Z+ be the filter 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 as following:
[w ∗ x]ij = q−1 2∑
a=− q−12
q−1 2∑
b=− q−12
[w]a+ q+12 ,b+ q+1 2 [x]a+i,b+j for i ∈ [P ], j ∈ [Q].
Denote
Dij,i′j′ = {( i+ a, j + b, i′ + a′, j′ + b′ ) ∈ [P ]× [Q]× [P ]× [Q] :
− (q − 1)/2 ≤ a, b, a′, b′ ≤ (q − 1)/2 } .
Intuitively, Dij,i′j′ is a q × q × q × q set of indices centered at (ij, i′j′). For any tensor T ∈ RP×Q×P×Q, let [T ]Dij,i′j′ be the natural sub-tensor and let Tr(T ) = ∑ i,j Tij,ij .
Formulation of CNN-IJ. Define the CNN-IJ as follows:
• Let the input x(0) = x ∈ RP×Q×C0 , where C0 is the number of input channels, and Ch is the number of channels in layer h. Assume WLOG that Ch = C for all h ∈ [L]
• For h = 1, . . . , L, let the inner representation
x̃ (h) (β) =
Ch−1∑
α=1
√ σ2W Ch W (h) (α),(β) ∗ x (h−1) (α) + C0∑
α=1
√ σ2U Ch U (h) (α),(β) ∗ x (0) (α) (13)
[ x (h) (β) ]
ij =
1
[S]ij
[ σ ( x̃ (h) (β) )]
ij
, for i ∈ [P ], j ∈ [Q] (14)
where W (h)(α),(β) ∈ R q×q represent the convolution operator from the αth channel in layer h − 1 to the βth channel in layer h. Similarly, U (h)(α),(β) ∈ R q×q injects the input in each convolution window. S ∈ RP×Q is a normalization matrix. Let W,U, S,σ2U ,σ2W be chosen by the CDEQ-NTK initialization described later. Notice here we assume WLOG that the number of channels in the hidden layers is the same.
• The final output is defined to be fθ(x) = ∑CL
α=1
〈 W
(L+1) (α) , x (L) (α) 〉 , where W (L+1)(α) ∈
RP×Q is sampled from standard Gaussian distribution.
CDEQ-NTK initialization. Let 1q ∈ Rq×q, X ∈ RP×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 )2 = [1q ∗ X̃]ij in eq. (14). Let every entry of every W,U be sampled from N (0, 1) and σ2W + σ2U = 1. Using the above defined notations, we now state the CDEQ-NTK. Theorem 3. Let x, y ∈ RP×Q×C0 be s.t ‖xij‖2 = ‖yij‖2 = 1 for i ∈ [P ], j ∈ [Q]. Define the following expressions recursively (some x, y are omitted in the notations), for (i, j, i′, j′) ∈ [P ]× [Q]× [P ]× [Q], h ∈ [L]
K (0) ij,i′j′ (x, y) =
  ∑
α∈[C0]
x(α) ⊗ y(α)


ij,i′j′
(15)
[ Σ(0) (x, y) ]
ij,i′j′ =
1
[S]ij [S]i′j′
C0∑
α=1
Tr ([ K
(0) (α) (x, y)
]
Dij,i′j′
) (16)
Λ (h) ij,i′j′ (x, y) =


[ Σ(h−1)(x, x) ]
ij,ij
[ Σ(h−1) (x, y) ]
ij,i′j′[ Σ(h−1) (y, x) ]
i′j′,ij
[ Σ(h−1) (y, y) ]
i′j′,i′j′

 ∈ R2×2 (17)
[ K(h)(x, y) ]
ij,i′j′ = σ2W [S]ij · [S]i′j′ E (u,v)∼N (0,Λ(h)
ij,i′j′ )
[σ(u)σ(v)] + σ2U
[S]ij · [S]i′j′ [K(0)]ij,i′j′ (18)
[ K̇(h)(x, y) ]
ij,i′j′ = σ2W [S]ij · [S]i′j′ E (u,v)∼N (0,Λ(h)
ij,i′j′ )
[σ̇(u)σ̇(v)] (19)
[ Σ(h)(x, y) ]
ij,i′j′ = Tr
([ K(h)(x, y) ]
Dij,i′j′
) (20)
Define the linear operator L : RP×Q×P×Q → RP×Q×P×Q via [L(M)]ij,i′j′ = Tr ( [M ]Dij,i′j′ ) .
Then the CDEQ-NTK can be found solving the following linear system: Θ∗(x, y) = K̇∗(x, y)⊙ L ( Θ∗(x, y) ) +K∗(x, y), (21)
where K∗(x, y) = limL→∞ K(L)(x, y), K̇∗(x, y) = limL→∞ K̇(L)(x, y). The limit exists if σ2W < 1.
5 EXPERIMENTS
In this section, we evaluate the performance of DEQ-NTK and CDEQ-NTK on both MNIST and CIFAR-10 datasets. We also compare the performance of finite depth NTK and finite depth iteration of DEQ-NTK.
Implementation. For DEQ-NTK, aligned with the theory, we normalize the dataset such that each data point has unit length. The fixed point Σ∗(x, y) is solved by using the modified Powell hybrid method (Powell, 1970). Notice these root finding problems are one-dimensional, hence can be quickily solved. For CDEQ-NTK, the input data x has dimension P × Q × C0, and we normalize x s.t ‖xij‖2 = 1 for any i ∈ [P ], j ∈ [Q]. We set q = 3 and stride 1. The fixed point Σ∗(x, y) ∈ RP×Q×P×Q is approximated by running 20 iterations of eq. (17), eq. (18), and eq. (20). The actual CDEQ-NTK Θ(x, y) is then calculated by solving the sparse linear system eq. (21).
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 ey . Let 1 ∈ Rn be an all-1 vector, we train on the new encoding −0.1 · 1 + ey . That is, we change the “1” to 0.9, and the “0” to −0.1, as suggested by Novak et al. (2018).
Result. On MNIST data, we test the performance of DEQ-NTK with σ2W = 0.25,σ2U = 0.25,σ2b = 0.5 and achieve 98.6% test accuracy. The results are listed in table 1.
On CIFAR-10, we trained DEQ-NTK with three different sets of random initializations. These initializations are not fine-tuned, yet we can still see they are comparable, or even superior, to the finite-depth NTK with carefully chosen regularization. For CDEQ-NTK, we compute the kernel matrix on 2000 training data and tested on 10000 samples. See the results in table 2.

heading: MNIST Method Model size Acc.
text: We should emphasize that the calculation of NTK requires a huge amount of computing resource, even for shallow networks. On the other hand, our method provides an efficient way to compute a special NTK with infinite depth. Typically, training a DEQ-NTK on all CIFAR-10 data takes around 400 CPU hour, and the training hour of CDEQ-NTK halves that of its finite-depth CNTK counterpart, as we only need to calculate Σ(h), whereas the actual CNTK needs to calculate both Σ(h) and Θ(h).

heading: CIFAR-10
text: With a smaller dataset with 1000 training data and 100 test data from CIFAR-10, we evaluate the performance of NTK and the finite depth iteration of DEQ-NTK, as depth increases. See fig. 1 As proven in Jacot et al. (2019), the NTK will always “freeze” in our setting. Therefore the NTK starts to become linearly independent as the depth increases, and its kernel regression does not have a unique solution. To circumvent this issue, we add a regularization term r ∝ %Θ(x,x)n , where n is the size of the training data. Such regularization is known to guarantee uniform stability (Bousquet & Elisseeff, 2002), and it still interpolates data in the classification sense (training accuracy is 100%).
6 CONCLUSION
We derive NTKs for both fully-connected DEQs and convolution DEQs, and show that they can be computed more efficiently than finite depth NTK and CNTK, especially when the depth is deep. Moreover, the performance of DEQ-NTK and CDEQ-NTK is comparable to their finite depth NTK counterparts. Our analysis shows that one can avoid the freeze and chaos phenomenon in infinitely deep NTKs by using input injection. One interesting question remained open is to further understand the role of σ2W ,σ 2 U ,σ 2 b in the fixed point computation, and how they affect generalizations of DEQNTKs.
REFERENCES Luis B Almeida. A learning rule for asynchronous perceptrons with feedback in a combinatorial
environment. In Artificial neural networks: concept learning, pp. 102–111. 1990.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems, pp. 8141–8150, 2019.
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. In Advances in Neural Information Processing Systems, pp. 690–701, 2019.
Shaojie Bai, Vladlen Koltun, and J Zico Kolter. Multiscale deep equilibrium models. arXiv preprint arXiv:2006.08656, 2020.
Olivier Bousquet and André Elisseeff. Stability and generalization. Journal of machine learning research, 2(Mar):499–526, 2002.
Minmin Chen, Jeffrey Pennington, and Samuel S Schoenholz. Dynamical isometry and a mean field theory of rnns: Gating enables signal propagation in recurrent neural networks. arXiv preprint arXiv:1806.05394, 2018a.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in neural information processing systems, pp. 6571–6583, 2018b.
Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity. In Advances In Neural Information Processing Systems, pp. 2253–2261, 2016.
Adrià Garriga-Alonso, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional networks as shallow gaussian processes. arXiv preprint arXiv:1808.05587, 2018.
Dar Gilboa, Bo Chang, Minmin Chen, Greg Yang, Samuel S Schoenholz, Ed H Chi, and Jeffrey Pennington. Dynamical isometry and a mean field theory of lstms and grus. arXiv preprint arXiv:1901.08987, 2019.
Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. Mean-field behaviour of neural tangent kernel for deep neural networks. arXiv preprint arXiv:1905.13654, 2019.
Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks generalize better than deep feedforward networks?–a neural tangent kernel perspective. arXiv preprint arXiv:2002.06262, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, pp. 8571– 8580, 2018.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Freeze and chaos for dnns: an ntk view of batch normalization, checkerboard and boundary effects. arXiv preprint arXiv:1907.05715, 2019.
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.00165, 2017.
Jaehoon Lee, Samuel S Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. arXiv preprint arXiv:2007.15801, 2020.
Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. arXiv preprint arXiv:1804.11271, 2018.
Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pp. 29–53. Springer, 1996.
Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are gaussian processes. arXiv preprint arXiv:1810.05148, 2018.
Fernando J Pineda. Generalization of back propagation to recurrent and higher order neural networks. In Neural information processing systems, pp. 602–611, 1988.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In Advances in neural information processing systems, pp. 3360–3368, 2016.
Michael JD Powell. A hybrid method for nonlinear equations. Numerical methods for nonlinear algebraic equations, 1970.
Maithra Raghu, Ben Poole, Jon Kleinberg, Surya Ganguli, and Jascha Sohl-Dickstein. On the expressive power of deep neural networks. In international conference on machine learning, pp. 2847–2854. PMLR, 2017.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv preprint arXiv:1611.01232, 2016.
Ezra Winston and J Zico Kolter. Monotone operator equilibrium networks. arXiv preprint arXiv:2006.08591, 2020.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel S Schoenholz, and Jeffrey Pennington. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks. arXiv preprint arXiv:1806.05393, 2018.
Greg Yang. Tensor programs i: Wide feedforward or recurrent neural networks of any architecture are gaussian processes. arXiv preprint arXiv:1910.12478, 2019.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006.14548, 2020.
A DETAILS OF SECTION 3
In this section, we give the detailed derivation of DEQ-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 FCNN-IJ here:
Let d be the input dimension, x, y ∈ Rd be a pair of inputs, Nh be the width of the hth hidden layers. Let N0 = d and NL+1 = 1. Define the FCNN-IJ with L layers as follows:
f (h) θ (x) = √ σ2W Nh W (h)g(h−1)(x) + √ σ2U Nh U (h)x+ √ σ2b Nh b(h), h ∈ [L]
g(L)(x) = σ(f (L) θ (x))
where W (h) ∈ RNh×Nh−1 , U (h) ∈ RNh×d are the internal weights, and b(h) ∈ RNh are the bias terms. These parameters are chosen using the NTK initialization. Let us pick σW ,σU ,σb ∈ R arbitrarily in this section.
Proof of theorem 1. First we note that
E [[ f (h+1)(x) ] i · [ f (h+1) (y) ] i | f (h) ]
= σ2W N
N∑
j=1
σ ([ f (h)(x) ]
j
) σ ([ f (h) (y) ]
j
) +
σ2U N
N∑
j=1
x⊤y + σ2b
→Σ(h+1)(x, y) a.s
where the first 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:
Θ(L)(x, y) = Eθ
[〈 ∂f(θ, x)
∂θ , ∂f(θ, y) ∂θ
〉]
=Eθ
[〈 ∂f(θ, x)
∂W , ∂f(θ, y) ∂W
〉]
︸ ︷︷ ︸ 1
+Eθ
[〈 ∂f(θ, x)
∂U , ∂f(θ, y) ∂U
〉]
︸ ︷︷ ︸ 2
+Eθ
[〈 ∂f(θ, x)
∂b , ∂f(θ, y) ∂b
〉]
︸ ︷︷ ︸ 3
.
The following equation has been proven in many places:
1 = L+1∑
h=1

σ2W E (u,v)∼N (0,Λ(h))
[σ(u)σ(v)] · L+1∏
h′=h
Σ̇(h ′) (x, y)
  , 3 = L+1∑
h=1
 σ2b · L+1∏
h′=h
Σ̇(h ′) (x, y)


For instance, see Arora et al. (2019). 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
〈 ∂f
∂U (h) ,
∂f̃
∂U (h)
〉
=
〈 ∂f
∂f (h) ∂f (h)) ∂U (h) , ∂f̃ ∂f̃ (h) ∂f̃ (h)) ∂U (h)
〉
=
〈 ∂f (h)
∂U (h) , ∂f̃ (h) ∂U (h)
〉 · 〈 ∂f
∂f (h) ,
∂f̃
∂f̃ (h)
〉
→σ2Ux⊤y · L+1∏
h′=h
Σ̇(h ′)(x, y)
where the last line uses the existing conclusion that 〈
∂f ∂f(h) , ∂f̃ ∂f̃(h)
〉 → ∏L+1 h′=h Σ̇
(h′)(x, y), this convergence almost surely holds when N → ∞ by law of large numbers. Finally, summing 〈
∂f ∂U(h) , ∂f̃ ∂U(h)
〉 over h ∈ [L] we conclude the assertion.
We now proceed to explain more about the fixed point convergence in theorem 1. Let us first show the limit converges.
Lemma 1. Use the same notations and settings in theorem 1 and theorem 2. Θ(L)(x, y) in eq. (6) converges absolutely if σ2W < 1.
Proof. Since we pick x, y ∈ Sd−1, and by DEQ-NTK initialization, we always have Σ(h)(x, y) < 1 for x ∕= y. Let ρ = Σ(h)(x, y), by eq. (5) and eq. (11), if σ2W < 1, then there exists c such that Σ̇(h)(x, y) < c < 1 for all finite number of x ∕= y on Sd−1, and large enough h. This is because limh→∞ Σ̇ (h)(x, y) = Σ̇∗(x, y) < Σ̇∗(x, x) < 1.
Use comparison test,
lim L→∞
L+1∑
h=1
∣∣∣∣∣∣ ( Σ(h−1) (x, y) ) · L+1∏
h′=h
Σ̇(h ′) (x, y) ∣∣∣∣∣∣ < 1 + lim L→∞ L+1∑
h=1
cL+1−h.
Since c < 1, the geometric sum converges absolutely, hence Θ∗(x, y) converges absolutely if σ2W < 1, and the limit exists.
B DETAILS OF SECTION 4
We first explain the choice of S in the CDEQ-NTK initialization. In the original CNTK paper (Arora et al., 2019), the normalization is simply 1/q2. However, due to the zero-padding, 1/q2 does not normalize all [ Σ(h)(x, x) ]
ij,i′j′ 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̃ ]
Dij,ij .
Now we give the proof to our main theorem.
Proof of theorem 3. Similar to the proof of theorem 1, we can split the CDEQ-NTK in two terms:
Θ(L)(x, y) = Eθ
[〈 ∂f(θ, x)
∂θ , ∂f(θ, y) ∂θ
〉]
=Eθ
[〈 ∂f(θ, x)
∂W , ∂f(θ, y) ∂W
〉]
︸ ︷︷ ︸ 1
+Eθ
[〈 ∂f(θ, x)
∂U , ∂f(θ, y) ∂U
〉]
︸ ︷︷ ︸ 2
.
Omit the input symbols x, y, let [ K̂(h) ]
ij,i′j′ = σ2W [S]ij · [S]i′j′ E (u,v)∼N (0,Λ(h)
ij,i′j′ )
[σ(u)σ(v)].
As shown in Arora et al. (2019), we have
〈 ∂fθ(x)
∂W (h) , ∂fθ (, y) ∂W (h)
〉 → Tr
 K̇(L) ⊙ L ( K̇(L−1) ⊙ L ( · · · K̇(h) ⊙ L ( K̂h−1 ) · · · )) 
Write f = fθ(x) and f̃ = fθ(y). Following the same step, by chain rule, we have 〈
∂f
∂U (h) ,
∂f̃
∂U (h)
〉 → Tr
 K̇(L) ⊙ L ( K̇(L−1) ⊙ L ( · · · K̇(h) ⊙ L ( K(0) ) · · · )) 
Rewrite the above two equations in recursive form, we can calculate the L-depth iteration of CDEQNTK by:
• For the first layer Θ(0)(x, y) = Σ(0)(x, y).
• For h = 1, . . . , L− 1, let [ Θ(h) (x, y) ]
ij,i′j′ = Tr
([ K̇(h) (x, y)⊙Θ(h−1) (x, y) +K(h) (x, y) ]
Dij,i′j′
) (22)
• For h = L, let
Θ(L) (x, y) = K̇(L) (x, y)⊙Θ(L−1) (x, y) +K(h) (x, y) (23)
• The final kernel value is Tr(Θ(L)(x, y)).
Using eq. (22) and eq. (23), we can find the following recursive relation:
Θ(L+1)(x, y) = K̇(L+1)(x, y)⊙ L ( Θ(L)(x, y) ) +K(h+1)(x, y) (24)
At this point, we need to show that K∗(x, y) ≜ limL→∞ K(L)(x, y) and K̇∗(x, y) ≜ limL→∞ K̇
(L)(x, y) exist. Let us first agree that for all h ∈ [L], (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 eqs. (17) to (20) to solve for [Σ(h)(x, y)]ij,i′j′ is equivalent to iterating f : RP×Q×P×Q → RP×Q×P×Q:
P (h+1) = f(P (h)) ≜ L (
1
[S]ij [S]i′j′ Rσ(P
(h)) ) , P (0) = K(0) (25)
where
Rσ(P (h) ij,i′j′) ≜ σ2W


√ 1− ( P
(h) ij,i′j′
)2 + ( π − cos−1 ( P
(h) ij,i′j′
)) P
(h) ij,i′j′
π

 + σ2UK (0) ij,i′j′ (26)
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 definition of S.
Now if we can show f is a contraction, then Σ∗(x, y) ≜ limh→∞ Σ(h)(x, y) exists, hence K∗ and K̇∗ also exist. We should keep the readers aware that f : RP×Q×P×Q → RP×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 finitely many entries, this is equivalent to say its ℓ∞ norm (imagine flattenning this tensor into a vector) converges. So we can equip the domain an co-domain of f with ℓ∞ norm (though these are finite-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 flatten 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′) ∈ Dij,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 Dij,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 2), we conclude that f is indeed a contraction.
With the same spirit, we can also show that eq. (23) is a contraction if σ2W < 1, hence eq. (21) is indeed the unique fixed point. This finishes the proof.
B.1 COMPUTATION OF CDEQ-NTK
One may wish to directly compute a fixed point (or more precisely, a fixed tensor) of Θ(L) ∈ RP×Q×P×Q like eq. (8). However, due to the linear operator L (which is just the ensemble of the trace operator in eq. (20)), 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 324 for CIFAR-10 data.
Moreover, this system is nonlinear. Therefore we cannot compute the fixed point Σ∗ by root-finding efficiently. Instead, we approximate it using finite depth iterations, and we observe that in experiments they typically converge to 10−6 accuracy in ℓ∞ within 15 iterations.

Review summaries:
This paper studies the double infinite-width + infinite-depth limit of fully connected and convolutional neural nets from an NTK angle, when input injections enter the picture. The techniques mix NTK techniques with Deep Equilibrium (DEQ) model techniques to directly compute the infinite-depth limit of the infinite-width limit of such neural nets. They show that there is no freeze/chaos transition for such networks (unlike the case without input injections). The writing is reasonably clear, although the size of the formulas is not very pleasant. 

This paper studies the neural tangent kernel (NTK) of fully-connected neural networks with input injection (defined in the first set of display in Section 3.1), and the infinite depth limit of the NTK. The calculations are further carried out for the convolution neural networks with input injection (defined at the beginning of page 6). Those kernels are empirically evaluated on MNIST and CIFAR-10 datasets, and are compared with the usual NTKs without input injection.

The paper shows the deep equilibrium model has non-degenerate neural tangent kernel in the infinite depth setting. The neural tangent kernel can be computed by a similar root-finding problem as that in the deep equilibrium problem itself. Some experiments have been performed to compare the performance of deep equilibrium neural tangent kernel with that of finite depth neural tangent kernel.