On the Neural Tangent Kernel of Equilibrium Models
Authors: Zhili Feng
J Zico Kolter
Venue:  (2021)
Keywords: deel learning
equilibrium model
neural tangent kernel

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 root-finding.  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.

Under review as a conference paper at ICLR 2021
ON THE NEURAL TANGENT KERNEL OF EQUILIBRIUM MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
Existing analyses of the neural tangent kernel (NTK) for infinite-deple networks show that the kernel typically becomes degenerate 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-dppth limit of weight-tied network via root- finding. 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- not-finding problem as in traditional DEQs, and highlight meth- ods 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.
INTRODUCTION
Recent works empirically observe that as the depth of weight- -ied its output tends to converge to 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 their widths go to infinity (Neal, 1996; Lee et al., 2017; Yang, 2019; Matthews et al., 2018; Novak et al., 2018; Aarriga-Alonss 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 fe with gradient descent equivalent solving kernel regression with respect the neural tangent kernel (Arora et al., 2019; Jacot et al., 2018; Yang, 2019; Huang et al., 2020).
network increases,
However, as the depth goes to infinity, Jacot 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 non-trivial way, thus avoiding both freeze and chaos. Furthermore, anal- ogous 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) com- puting 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-scal 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

---

Under review as a conference paper at ICLR 2021
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 unique stable fixed point (Winston & Kolter, 2020).
On the side of connecting neural networks to kernel methods, Neal (1996) first discovered that single-layered network with randomly initialized parameters becomes 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-trainge where z, y E R2 are two samples, represents the parameters of the network, /V the initialization distribution (often Gaussian) over 0, and f(0,-) E I. is the output of the network.
nets induce the kernel (z., v) B (0, x:) f(0 y)),
One related topic in studying the relation between Gaussian process kernel and depth the mean- field theory. Poole et al. (2016); Schoenholz et al. (2016) showed that the correlations between all inputs on an infinitely wide weakly-trainne net become either perfectly correlated (onder) or decorrelated (chaos) depth increases. This aligns with the observation in Jacot et al. (2019). They suggested we should initialize the neural network on the "edge-of-chaon" 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.
to make sure that
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 fully-trained infinite width network is the following:
( (01(0.0) 05(0.)) v) ]
(zx, v) = -Eows
00 , 00
They also gave 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 0 as the parameters for some network fo or equivalently, -0. We write capital letter W to represent matrices or tensors, which should be clear from the context, and use [W]: to represent the element of W indexed by i. We write lower case letter x to represent vec- tors scalars. For Z4, let [a] {l,.... a}. Denote (z) v2 max(0, x) as the normalized ReLU and ē its derivative (which only need to be well-defined almost everywhere). The symbol o
with subseript is always used to denote the variance of some distribution. We write (p,, E) as the
Gaussian distribution with mean μ E 2d and covariance matrix Σ E dddd sphere embedded in IRd.
We let gd- be the unit
2

---

Under review as a conference paper at ICLR 2021
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 colly-connected neural nets with imput 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, 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 be found in appendix A.
3.1 FINITE DEPTH ITERATION OF DEQ-NTK
Let d be the input dimension, 2, y E IR be pair of inputs, NA be the width of the hth layers where [L + 11. Let No d and Nr.-1 1. Define the FCNN-IJ with L hidden layers as follows:
o No
of o x V NA
of b(h) he h [1+1]
100h (x) = V
w(h) ghh (x) + g
Nn
(x) (x) o(5(
where W()) E NNN, 1, (h) E RN are the internal weights, and b(h) E R. are the bias terms. These parameters are chosen using the NTK initialization. Let us pick ow, ou. ob R arbitrarily in this section.
NTK initialization. We randomly initialize every entry of every W, U,U,b from N00, 1).
Without loss of generality (WLOG) we assume the width of the hidden layer N₁. = is the same across different layers. We remark the readers to distinguish FCNN-IJ from recurrent neural network (RNN): our model injects the original input to each layer, whereas a RNN has a sequence of input (æ1,..... JT and inject o, the 4 layer.
Here 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, ((++1) Otherwise it is unlikely the network will converge at all. fact, one needs 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 ow of of and the nonlinearity a. Assuming untied weights makes the analysis easier, but the same argument can be made rigorously for tied weights well, see Yang (2019; 2020).
that is, (()) =
Our main theorem is the following:
Theorem 1. Recursively define the following quantities for [L]:
x (x. y) xT
(21-1)222)) 1(- (x,
(1)) (x. (11-1)(gz)
R22x
(3)
(h- (y.
(0 (z. y) o ow.B(a.e) (o())(o))) + o y t
(4)
5100 (z. y) our (..))))))))))) (w)
(5)
Then the I-depth iteration to the DEQ-NTK can be expressed as:
L+1 Σ
e (z)
(I(L-1) (I(A-1((.))) (I.0))
II E(()) (2.10)
(6)
h=1
R'=h
where by convention set 5 (2x for the T--epth iteration.
3

---

Under review as a conference paper at ICLR 2021
Proof Sketch. The first step is show that each layer h [L], the representation 1/ (x)()) is associated with a Gaussian process with kernel eq. (3) as N - xc. Then use the characterization in eq. (I), calculate the NTK by:
(a:
9 (x,,) Eg ( (0110,0)
0f(0,y)
00
00
Eg K @
(0
@f(0,a))
a5(0.v) au
)
@8((0m)) Öb
05(0,y)
+ Eg
OW
oW
aU
Calculating each term using the chain rule, 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 ( is. By convention, we assume the two samples I, y gd-1, and we require the parameters ow of of obey the DEQ- NTK initialization.
DEQ-NTK initialization. Let every entry of every W, U. follows the NTK initialization de- scribed in section 3.1, as well as the additional requirement ow + ot + to =.
Let the nonlinear activation function o be the normalized ReLU: (xr) = v2 max(0,
Definition 3.1 (Normalized activation). We call an activation function a
normalized
(x=)) =
Using normalized activations along with DEQ-NTK initialization, we can derive the main conver- gence theorem:
Theorem 2. Use same notations and settings in theorem With input data x. y E Sdd-1 parameters oi. o.. of follow the DEQ-NTK initialization, the fixed point of DEQ-NTK is
Σn (x, v)
+* (²,
4 lim ( (1D)
L
1 - 5- y)
where 2 (x, y) A p" the root of:
[V- p² + (t cos p)p) + + y o
Ro where Ro (p) 2 ow
(8)
and
( T (T-cos cos (p*) (p*)
(x, A lim (x,)))))) oiv h-000
π
Proof. Due the fact that x E $d-1, a being a normalized activation, and DEQ-NTK initialization, one can easily calculate by induction that for all h E [L]:
(x) x) = oi E (0.1)
[u()] + oF T x + - =1
This indicates that eq. (3), the covariance matrix has a special structure A (x, = (3
P 1).
where (x, y) depends on h, N. y. For simplicity we omit the h, x, in A (1) (x) 8). As shown in Daniely et (2016):
[o(u)o()) VI- p² + (N cos (p))
(10)
(v.
E [o(u)ō())] = π cos (p)
(11)
(u.0)~>((.A))
4

---

Under review as a conference paper at ICLR 2021
Adding input injection and bias, derive eq. (8) from eq. (10), and similarly, eq. (9) from eq. (11). Notice that iterating eqs, (2) to (4) to solve for (1h) (z., y) is equivalent to iterating (R, o o R (p)
with initial input p y. Take the derivative
dR, (p) dp
Wv 1 COS (())
)
)
< 1, of < 1 and -1 -15
For x A y we have -1 I < C < for some c (this is because we only have finite number of inputs 2,, y) and by DEQ-NTK initialization we have oy 1, so the above inequality hold. Hence Ror (p) is a contraction on [0, c). and we conclude that the fixed point p* is attractive.
By lemma 1, if ow 1, then the limit of eq. (6) exists, so can rewrite the summation form in eq. (6) in a recursive form:
(x, - (x,
(x, (1(T)) (a, v) (14)(z))) + E(Lt)) v),
and directly solve the fixed point iteration:
lim O(L+1) (x.y) v) lim
(x. (S(+))
( (z. v) + (xLt))(z.))
L-00
L-000
lim ((L+1) (x. v) = -E lim (z)) v) + (x,
(12)
L-000
1-00
= lim (x. = (x, lim @ (r,y) (z,)), + E' (z,
1000
100
Solving for lim ((T)) (z. v) we get 0 (z, (x.9)
(x,y))
Remark 2. Note our Ex 3) always depends the inputs and so the information between two imputs is always preserved, even if the depth goes to infinity. On the contrary as pointed out by Jacot et al. (2019), without input injection, (1 (x, y) always converges + co, even + y.
4 DEQ WITH CONVOLUTION LAYERS
In this section we show how to derive the NTKs for convolution DEQs (CDEQ). Although 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 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. E PxQ be a pair of inputs, let q Z4 be the filter size (WLOG assume it is odd well). By convention, always pad the representation (both the input layer and hidden layer) with 0's. Denote the convolution operation as following:
[w zlij
Σ
[ wha+ stt .u 47 (c]o+ioo for i [P], [Q].
J
Denote
Dig.e {{ (i+a. (., +o' ( +'' [P] [Q] x [P] [Q]
(4-1)/2 S a, b, a' s (q - 1)/2).
Intuitively, Dij.ig' is a q x x g x g set of indices centered (ij,''j'). For any tensor T RPxxxxxQ let [T]p be the natural sub-tensor and let Tr(T) Σig Tizas-
5

---

Under review as a conference paper at ICLR 2021
Formulation of CNN-IJ. Define the CNN-I follows:
Let the input 20 x E PxQxCo, where Co is the number of input channels, and C is the number of channels in layer h. Assume WLOG that Cj. = C' for all h E [4]
For h I...., L. let the inner representation
Ch-1
th Ch
Co
(h) (3)
Σ
V
w(h)
(a).(0)
* (1-1)
(a)
+ Σ
Ch
U
(a).(0)
(h)
((a
(0)
(13)
(00
1
(100) (3)
for e (P].5 E [Q]
(14)
(8)
1j
(2) E 12 represent the convolution operator from the th channel in layer h 1 to the 3th channel in layer h. Similarly, convolution window. S E IR PxQ is a normalization matrix. Let W; U, of: ow be chosen by the CDEQ-NTK initialization described later. Notice here we assume WI.OG that the number of channels in the hidden layers is the same.
where (h(h)
(7(a) (a)((3) E Roxq injects the input in each
The final output is defined be fo(Ø) E.e ( (wiLt)) ( (aa)/
where w (a)
PxQ is sampled from standard Gaussian distribution.
CDEQ-NTK initialization. Let lg E Roxq, E RRQ be two all-one matrices. Let X e ((P+2)x((+2) be the output of zero-padding X. We index the rows of X by {0, (O......,PP P + 1] and columns by f0, 1,.., Q + 1}. For position i [P]. [Q], let (I$),)) "/5))===>>> = [1g Xli in eq. (14). Let every entry of every W, U be sampled from (N 1) and of of = 1.
Using the above defined notations, we now state the CDEQ-NTK.
Theorem 3. Let 2, E Pxxxxx be s.t || = l1asell = 1 for i E [P].j E 1Q). Define the following expressions recursively (some y, are omitted in the notations), for (i,j,", (P) x (Q) x (P) x (Q), h [T.]
k-10)
1889 (I, =
Σ
T((a) & M(a))
(15)
LcE(Col
"j.''"'
1
Tr
(0)
(x,, v)
(16)
13.8'"
[S] 51
(a)
Dey
2) (os, x)
()- (x, v)
Aggy (x,2)
(K)
63.03
E
(17)
11h-1) (y,
[h-- 1) (v:1)
(x, v)
[S]is oi [Sley
E (u.0))
[o(u)o())] +
[S] ot
K
(18)
15.855"
[5
[S]k o2p [S]es NT..A
ix(00 (z0, v)
E
[o(u)((b)
(19)
63.055
-T((XMM (K v) Dises)
(1) 13.83
= -Tr (
(20)
Define the linear operator L: PxQxxXx - PPXXXXXQ via (L(M)) = -T((MMagr))
Then the CDEQ-NTK can be found solving the following linear system: ('(x, y) = (x y) @ (( (zx, v)) + K K*(xx,)) y), where K* (x. y) = lim K (2. y), (x, lim K(0 y). The limit exists if ow
(21)
<
1.

---

Under review as a conference paper at ICLR 2021
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" (ar, 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 has dimension P x Q x Co, and we normalize for any i E [P].j [Q]. We set = and stride The fixed point D (a, y) E PxQxxxx is approximated by running 20 iterations of eq. (17), eq. (18), and eq. (20). The actual CDEQ-NTK ((xr,) y) is then calculated by solving the sparse linear system eq. (21).
a s.t || ll2
After obtaining the NTK matrix, we apply kernel regressions (without regularization unless stated otherwise). For any label y E {l...., nj. denote its one-hot encoding by ey. Let 1 E R" be all-1 vector, we train on the new encoding -0.1 1 + By. That is, we change the "1" to 0.9, and the "0" to 00.1, suggested Novak al. (2018).
0.35
0.35
0.25
0.25
0.15
0.15
Finite
Fintte DEQ
Fintte Dopth
W-0.6 bm0.4
W-0.8, bm0.2
0.38
0.36
0.38
0.32
0.37
0.28
Finbe Depth NTK Finkia
0.35
rinte Depth NTK
0.26
10
20
50
60
3.
20
40
50
60
Depth meural metaworicc
Figure Finite depth NTK vs. finite depth iteration of DEQ-NTK. is initialized with oy and of in the title. For DEQ-NTK we set of: of 0.1 in the title, and the actual of 0.1. All models are trained on 1000 CIFAR-10 data and tested on 100 test data for 20 random draws. The error bar represents the 95% confidence interval (CI). As expected, as the depth increases, the performance NTKs drop, eventually their 95% CI becomes a singleton, yet the performance of DEQs stabilize. Also note with larger of the freezing of NTK takes more depths to happen.
experiments, the NTK
Result. On MNIST data, we test the performance DEQ-NTK with of = 0.25, of 0.25, o
2 = 0.5 and achieve 98.6% test accuracy. The results listed in table 1.
On CIFAR-10, 10, we trained DEQ-NTK with three different sets random initializations. These initializations not fine-tuned, yet
still see they comparable, finite-depth NTK with carefully chosen regularization. For CDEQ-NTK, we compute the kernel matrix 2000 training data and tested 10000 samples. See the results in table 2.
superior, to the
7

---

Under review as a conference paper at ICLR 2021
MNIST
Method DEQ-NTK
Model size
Acc.
98.6% %
Neural ODE
84K
98.2%
MON DEQ
84K
98.2%
Table 1: Performance of DEQ-NTK on MNIST dataset, compared to neural ODE (Chen et al., 2018b) and monotone operator DEQ, see these results from Winston & Kolter (2020).
We should emphasize that calculation of NTK requires a huge amount of computing resource, even for shallow networks. On the other hand, method provides 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, only need to calculate (k), whereas the actual CNTK needs to calculate both (()) and (1h)
CIFAR-10
Method
Parameters
Acc.
DEQ-NTK
di = 0.25, o 0.25, o 0.5
59.08%
DEQ-NTK
oy = 0.6, of. -0.4, 0f =
59.77% %
DEQ-NTK
O = 0.8, où = -.2, = 0
59.43%
NTK with ZCA regularization
Ow o = 0.01
59.7%
DEQ-CNTK with 2000 training data
Tw 0.65, of 0.35
37.49%
Table 2: Performance of DEQ-NTK and CDEQ-NTK on CIFAR-10 10 dataset, see Lee et al. (2020). for NTK with ZCA regularization.
With 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" our setting. Therefore the NTK starts to become linearly independent unique solution. To cireumvent this issue, we add a regularization term e9(2.2) where the size of the training data. Such regularization is known to guarantee uniform stability (Bousquet & Elisseeff, 2002), and still interpolates data in the classification sense (training accuracy is 100%).
the depth increases, and its kernel regression does not have
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 comparable 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 ohyo of in the fixed point computation, and how they affect generalizations of DEQ- NTKs.
REFERENCES
Luis B Almeida. learning rule for asynchronous perceptrons with feedback in combinatorial environment. In Artificial neural networks: concept learning, pp. 102-111. 1990.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ Salakhutdinov, and Ruosong Wang- On exact computation with infinitely wide neural net. In Advances in Neural Information Processing Systems, pp. 8141-8150, 2019.
Shaojie Bai. Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, pp. 690-701, 2019.
8

---

Under review as a conference paper at ICLR 2021
Shaojie Bai, Vladlen Koltun, and J Zico Kolter. Multiscale deep equilibrium models. arXiv preprint arXiv:2006.08663 2020.
Olivier Bousquet and André 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 mean field theory of runs: Gating enables signal propagation in recurrent neural networks. arXiv preprint arXiv: 05394, 2018a.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David 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 dual view expressivity. In Advances In Neural Information Processing Systems, pp. 2253-2261, 2016.
Adrià Barriga-Alonse Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional net- works as shallow gaussian processes. arXiv preprint arXiv: 008.08887 2018.
Dar Gilboa, Bo Chang, Minmin Chen, Greg Yang, Samuel S Schoenholz, Ed H Chi, and Jeffrey Pennington. Dynamical isometry and mean field theory of Istms 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. 2019.
Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks gener- alize better than deep feedforward networks?-a neural tangent kernel perspective. arXiv preprint 20fi:22002.06662. 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. In Advances in neural information processing systems, pp. 8571- - 8580, 2018
Arthur Jacot, Franck Gabriel, and Clément Hongler. Freeze and chaos for dnns: an ntk view of batch normalization, checkerboard and boundary effects. arXiv preprint arXiv:1907.05711 2019.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint D:NNv://11..01110 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 Matthews, Mark Rowland, Jiri Hron, Richard Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. arXiv preprint arXiv:1804.1122 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: 2018.
Fernando Pineda. Generalization of back propagation to recurrent and higher order neural net- works. In Neural information processing systems, pp. 602-611, 1988.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponen- tial expressivity in deep neural networks through transient chaos. Advances in neural informa- fion processing systems, pp. 3360-3368, 2016.

---

Under review as a conference paper at ICLR 2021
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. the ex- pressive 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. (rXiv:2006.08591, 2020.
Monotone operator equilibrium networks. arXiv preprint
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Penning- ton. Dynamical isometry and mean field theory of cnns: How train 10,000-layer vanilla convolutional neural networks. arXiv preprint arXiv:/806.05399, 2018.
Greg Yang. Tensor programs i: Wide feedforward or recurrent neural networks any architecture gaussian processes. arXiv preprint arXiv:1910.1244: 2019.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006.14542 14548. 2020.
A DETAILS OF SECTION 3
In this section, give the detailed derivation of DEQ-NTK. There two terms that different a.(6.2)
from NTK: (x) y) and the extra EE (
a((0.w)
in the kernel.
80
8U
Let us restate the FCNN-IJ here:
Let d be the input dimension. v.y E Ed be pair of inputs, N₁ be the width of the 4t hidden layers. Let No d and NL+1 = Define the FCNN-IJ with L layers as follows:
dh (1- + Nn
oh +
oh
(x) = V
(b) ( N
(()) h E [4]
V Nn
(x) (ff (x))
where W(A) N.N (h)) E R.Nxxd xd the internal weights, and ((h) E R.N, the bias terms. These parameters are chosen using the NTK initialization. Let us pick ow, ou ab IR arbitrarily in this section.
Proof of theorem 1. First we note that
[gnntt pmm)
E
Σo ([ a (( 1=1
N
of
N
+
Σ
x y + 0
N
N 1=1
S (b+
(x, v) 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.
10

---

ař over E [L] we conclude the assertion. Under review as a conference paper at ICLR 2021
Recall the splitting:
<
(0((0,)) as(0,y)
(x, Eg
08
@f(0,.)) (f(
@0((0.z) x)
0f(0.y)
@f(0, a)
0f(0,y) v)
Eg
EB
+ Eg
OW
ow
OU
au
Öb
The following equation has been proven in many places:
L++
L++
L+1
L++1
Σ
ow
E (1.0)-((..AC.))
[o(u)o(v)]
II (()) (x,
3
of II h/mh
(
(x,
For instance, see Arora et al. (2019). So we only need to deal with the second term Write f = fa( and ff = fo(y), by chain rule, we have
(f(0.))
0f(0.) 00
Eg
80
80
0f of OU()) (U(N
of QYTN) of of
0fC) OU
ofth aU()
@fth) ff OU(N aU()) >
os os afth) of(0)
L+1
+ T y
II (2,0)) k=
where the last line the existing conclusion that convergence almost surely holds when N + by law of large numbers. Finally, summing
ai (yyक. af
26
l1 $000 v),, this
We now proceed to explain more about the fixed point convergence in theorem 1. Let us first show the limit converges.
Use the same notations and settings in theorem and theorem 2. (2D y) in eq. (6) converges absolutely if oy <
Lemma 1.
Proof. Since we pick x, E gd-1, and by DEQ-NTK initialization. we always have (hh) v) < for x # Let [(h (x. v). by eq. (5) and eq. (11), if ofy < 1, then there exists e such that 1(1) y)
1 for all finite number of # y gd-1, and large enough h. This is because lim (N) (x., y) シ
(.. (z. y) E" (x,
Use comparison test,
L+1 II M=
I+t1 L+1-h
lim
() (z,w)
s(() (x, w) <
lim
Since e < 1, the geometric sum converges absolutely, hence O Oz y) converges absolutely o and the limit exists.
< 1
11

---

Under review as a conference paper at ICLR 2021
B DETAILS OF SECTION 4
We first explain the choice of S in the CDEQ-NTK initialization. In the original CNTK paper (Arora al., 2019), the normalization is simply 1/4². However, due to the zero-padding, 1/9² does normalize all Σ (2),2)) ij,
expected: only the variances that away from the corners normalized to 1, but the ones near the corner are not. [S]s is simply the number non-zero entries
X
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:
< @f((.2) x)
0f(0.
(10) = Es
00
00
(8f(0x))
0f(0, u)
0f v)
Eg
+ EE
OW
aw
OU
aU
Omit the input symbols 2, y. let
(b)
ow
E
[o(u)a(o)).
[s)s
[Shes (u. (O0.A 15.5)))
As shown in Arora et al. (2019), have
@fo(x) 0fe (.)(.)) Tr RR0 o (KU-NI (...-))((.... k(N o (R)-I)...))
DW(h)' OW(h)
Write f= so(x) and f = (o Following the same step, by chain rule, we have
0f af (0)
k000 @ oz ( + Tr k00 o kU- o O
(K(O))....
(K)
Rewrite the above two equations in recursive form, we can calculate the L-depth iteration CDEQ-
NTK by:
For the first layer e (I, (00) v)-
For h I,...,L let
(RRIN) (I, @(h11) (0 (x.v) + k(h) (α, ) (22)
(@(h) (xr, wi] ig.813
Tr
For h L, let
((4) = (4) (x, 2 ((4-1) (x, k(b) (x,z)
(23)
The final kernel value is (r)) y)).
Using eq. (22) and eq. (23), can find the following recursive relation:
(1++1) (((+)) o v)) + ((h+1) (2.
(24)
12

---

Under review as a conference paper at ICLR 2021
At this point, need to show that (x, & lim K (2. (z. 3) and (2. - lim (( (ar, v) exist, Let us first agree that for all h E (1.. (P) x (Q)
the diagonal entries of A "' are all ones. Indeed, these diagonal entries are I's at h = by initial- ization. Note that iterating eqs. (17) (20) solve for (2, is equivalent iterating f: PRQXXXX IR PXQXPXQ
p(h+1) = f(p(h)) & L [51ts|5ks Ro (P((i)) ,P(0) K((0)
(25)
where
1 (plt) 2 (118)) + -
cos
(p()) (P (((or)))
pot)
cos
(o(P
Ro 15,0g))))) oow
+ Ki (26)
is applied to p(h) entrywise.
Due to CDEQ-NTK initialization, if (0)) h. This is true by the definition of S. sj.33 - i [P], E [Q]. then plh)
fj.éj for all iterations
Now if we can show f is contraction, then Ex y) A lime (0)) (x, y) exists, hence K and K* also exist. We should keep the readers aware that PPXQxxx should be careful with the metric spaces. We want every entry of (1 (1r, y) to converge, since this tensor has finitely many entries, this is equivalent to say its 8RD norm (imagine flattenning this tensor into vector) converges. So we can equip the domain an co-domain of f with 8rD norm (though these are finite-din diiie-dimensional spaces so we can really equip them with any norm, but picking e norm makes the proof easy).
PxQxxxx so we
Now have
L B RS 60
If we flatten the four-way tensor p(h) into co0
vector, then L. can be represented by (P x Q Px (P Q P x Q) dimensional matrix, whose (kl, - entry in the ("")"""") J-th row is if (kl, KT) E Dis.ays and otherwise. other words, the ed norm of the (ij, (') row represents the number of non-zero entries in Dija'y', but by the CDEQ-NTK initialization, the row l norm divided by [S] [S] is at most 1! Using the fact that |[C|e=>->>> is the maximum ll norm of the row, and the fact R, is contraction (proven in theorem 2), we conclude that f indeed contraction.
With the same spirit, can also show that eq. (23) is contraction if oiy < 1, hence eq. (21) is indeed the unique fixed point. This finishes the proof.
B.1 COMPUTATION OF CDEQ-NTK
One may wish directly compute fixed point (or more precisely, a fixed tensor) of ((L) PxQxxxx PxQ 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 (P x Q x P Q) (P x Q P x Q)-dimensional matrix that represents L. Even if we exploit the fact that only entries on the same "diagonal" depend on each other, at least P x Q P x Q, which 32 for CIFAR- 10 data.
Moreover, this system is nonlinear. Therefore we cannot compute the fixed point " by root-finding efficiently. Instead, approximate it using finite depth iterations, and we observe that in experi- ments they typically converge to 10-6 accuracy in 8 within 15 iterations.
13


Reviews:
[{"review_id": "msfYlcZaJFs", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper considers the problem of deriving NTK for DEQ models and shows that DEQ models have non-degenerate NTKs even in the infinite depth limit. It also provides several experimental comparisons on the performance of the obtained DEQ-NTK and CDEQ-NTK with NTK and CNTK, respectively. \n\nStrengths: \n\n1) The derivation of NTK for DEQ models seems to be reasonable and correct, and the resulting DEQ models have non-degenerate NTKs. \n\n2) By using the rooting-finding ability of the DEQ models, these derived NTK kernels can be computed for both fully-connected and convolutional variants.\n\nWeaknesses: \n\n1) The experiment design has some serious flaws. For example, the experiments include comparisons only with NTK and CNTK. It would be natural to also include comparisons with DEQ models, as DEQ-NTK (or CDEQ-NTK) can be viewed as augmented model of DEQ (or CDEQ). Thus, it would be interesting to see whether the new models can indeed achieve better performance. Otherwise, it would be no point to use NTK on DEQs. \n\n2) Particularly, the paper misses comparisons with several important DEQs models, like MON DEQ and Single stream DEQ,  on CIFAR-10 and MNIST.\n \n          https://arxiv.org/abs/2006.08591​ Monotone operator equilibrium networks\n          https://arxiv.org/abs/2006.08656​ Multiscale Deep Equilibrium Models\n          https://arxiv.org/abs/1909.01377​ Deep equilibrium models \n\n   These DEQ models achieve much better experimental results than the reported ones in this paper. For example, \n        CIFAR-10:\n           MON DEQ: single conv 74.1% \n          Single stream DEQ: around 82.2% \n     MNIST:\n         MON DEQ: single conv 99.2%\n\n  Missing these comparisons intentionally or unintentionally significantly weakens the paper. \n\n3) This paper also misses the comparison between CNTK  and CDEQ-NTK on  CIFAR-10. \n\n     CDEQ-NTK with 2000 training set result: 37.49% --- reported in this paper\n     CNTK (vanilla) with 2000 training set: 40.94% (Depth 3), 42.54%(Depth 4), 43.43%(Depth 6), 43.42%(Depth 11), 42.53%(Depth 21).     \n      https://arxiv.org/abs/1904.11955​ On Exact Computation with an Infinitely Wide Neural Net\n\nIt would strengthen the paper if it can show that CDEQ-NTK achieves better performance than CNTK when the number of layers increases. This is also what the paper tries to claim.\n\n4)  The DEQ-NTK in the experiment is derived for FCNN. From the first experiment, one can only see that DEQ-NTK achieves better performance than NTK for FCNN when the depth is large.  Since it makes more sense to apply DEQ to a complicated neural network (like transformer or trellis net) than to an FCNN, it leaves a doubt whether the proposed DEQ-NTK has any  practical importance. Thus, it would be better if the paper conducts experiments for the more complicated networks.\n\n5) It is not surprising to see that when the depth increases, DEQ-NTK remains stable while NTK does not, as this seems to be enabled more by the mechanism of the DEQ model.\n\n6) Related work is not sufficiently discussed. For example, the freezing of NTK. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "mpxa544gMyn", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nOverall I vote for rejecting. My concerns are as follows:\n\nThe paper lacks related literature. First, the motivation of considering deep equilibrium models is unclear to me. The authors should provide some further literature review. The advantage of using such a model in practice should be explained. Second, related proof techniques in the existing literature needs to be discussed. \n\nThe result is expectable and the proof techniques are not novel. The main theorem (Theorem 1) is the simple extension of the existing results on neural tangent kernel. The following theorem (Theorem 2) is the consequence of the main theorem under some specified initialization.\n\nThe theorems in the paper are lack of explanation. More discussion is needed to explain and extend the results in the paper.\n\nThe experiment part is not well-organized. More description is needed to improve the results.\n\nThe paper has some grammar mistakes and misuse of words. The paper needs to be revised carefully. To name a few:\nAbstract: DEQ model....DEQ models have...\nSection 3, 1st paragraph: we simplify fully-connected DEQs as DEQs.\nSection 3, 1st paragraph: In section 3.1, we show the NTK of the approximated DEQ using finite depth iteration...\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "gkRs5ZcZNYf", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nThis 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. \nThe experimental part is not very detailed, and it is not clear what is the take-home message from it. \nPros: This is quite interesting, the technique is nice. \nCons: the scope is somehow limited to a class of models which is not very much used; (as the authors say), the role of the normalizations factors that appear is not super clear; no very surprising phenomena, somehow. \nOverall, I think that these results are mathematically interesting and could lead in principle to practically useful insights, although this is not realized at this point. There could be some presentation effort in terms of the sizes of the formulae: what do we really need to know from them?", "extra_scores": {}, "extra_text": {}}, {"review_id": "0lkQHux4YQ5", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe theorems derived in this paper are incremental given existing studies on NTKs. The calculations are very similar to existing ones except the network structures considered in the current paper are slightly different. The infinite-depth limit of the kernels now indeed depends on the input, but the result is not surprising as the input is injected in each layer. \n\nThis paper also lacks a proper introduction to many concepts. For example, it is hard to understand what does DEQ-NTK really means in the introduction. Also the term NTK at first refers to the general concept in (1), but later seems to specifically refer to the NTK of fully-connected neural networks without input injection. Section 2 presents some background, but it gives many pieces of related works without a clear structure. For instance, I don't see how many concepts like \"weakly-trained\", \"fully-trained\", \"edge-of-chaos\" are relevant to the current paper. \n\nIn the experiments, the choices of parameters seem to be very arbitrary. The authors do not provide systemic guidance on how they tune those parameters, while the performance improvement is minor. Since one major motivation for studying NTKs is the relation to the actual neural networks, some proper comparisons or comments should be included. The experiment section in the current form is not very convincing. \n\nGiven the above concerns, I don't think this paper is suitable for publication in ICLR. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper combines recently emerging NTK theory and kernels with DEQ models. In particular the authors use the root-finding capability of DEQ models to compute the corresponding NTK of DEQ models for fully connected and convolutional variants. The reviewers raised various concerns including lack of experimental details, incremental theoretical results(which the authors agree with but postulate that this is a practical paper), lack of proper literature review, explaining how it applies in practical scenarios and grammatical mistakes. Some of these concerns were addressed during the response period but none of the reviewers were fully satisfied with the author's response. While I think there are interesting ideas in this paper I agree with the reviewers that a substantial revision is required and therefore recommend rejection.
Decision: Reject