On the Neural Tangent Kernel of Equilibrium Models

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 .

Summary: 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.