Fast Neural Kernel Embeddings for General Activations
Authors: Insu Han
Amir Zandieh
Jaehoon Lee
Roman Novak
Lechao Xiao
Amin Karbasi
Venue: NeurIPS 2022 Accept (2022)

Abstract: Infinite width limit has shed light on generalization and optimization aspects of deep learning by establishing connections between neural networks and kernel methods. Despite their importance, the utility of these kernel methods was limited in large-scale learning settings due to their (super-)quadratic runtime and memory complexities. Moreover, most prior works on neural kernels have focused on the ReLU activation, mainly due to its popularity but also due to the difficulty of computing such kernels for general activations. In this work, we overcome such difficulties by providing methods to work with general activations. First, we compile and expand the list of activation functions admitting exact dual activation expressions to compute neural kernels. When the exact computation is unknown, we present methods to effectively approximate them. We propose a fast sketching method that approximates any multi-layered Neural Network Gaussian Process (NNGP) kernel and Neural Tangent Kernel (NTK) matrices for a wide range of activation functions, going beyond the commonly analyzed ReLU activation. This is done by showing how to approximate the neural kernels using the truncated Hermite expansion of any desired activation functions. While most prior works require data points on the unit sphere, our methods do not suffer from such limitations and are applicable to any dataset of points in $\mathbb{R}^d$. Furthermore, we provide a subspace embedding for NNGP and NTK matrices with near input-sparsity runtime and near-optimal target dimension which applies to any \emph{homogeneous} dual activation functions with rapidly convergent Taylor expansion. Empirically, with respect to exact convolutional NTK (CNTK) computation, our method achieves $106\times$ speedup for approximate CNTK of a 5-layer Myrtle network on CIFAR-10 dataset. 

Fast Neural Kernel Embeddings for General Activations
Insu Han'
Amir Zandieh² Lechao Xiao
Jaehoon Lee Amin Karbasi
Roman Novak
Yale University
2 "Max-Planck-Instiee für Informatik
Google Research
Abstract
Infinite width limit has shed light on generalization and optimization aspects of deep learning by establishing connections between neural networks and kernel methods. Despite their importance, the utility of these kernel methods was limited in large-scale learning settings due their (super- (super-)quadrall )quadratic runtime and memory complexities. Moreover, most prior works on neural kernels have focused on the ReLU activation, mainly due its popularity but also due the difficulty of computing such kernels for general activations. In this work, we overcome such difficulties by providing methods work with general activations. First, we compile and expand the list of activation functions admitting exact dual activation expressions to compute neural kernels. When the exact computation unknown, present methods to effectively approximate them. We propose fast sketching method that approximates any multi-layered Neural Network Gaussian Process (NNGP) kernel and Neural Tangent Kernel (NTK) matrices for a wide range of activation functions, going beyond the commonly analyzed ReLU activation. This is done by showing how to approximate the neural kernels using the truncated Hermite expansion of any desired activation functions. While most prior works require data points the unit sphere, applicable to any dataset of points in Rd. Furthermore, we provide a subspace embedding for NNGP and NTK matrices with near input-sparsity runtime and near-optimal target dimension which applies to any homogeneous dual activation functions with rapidly convergent Taylor expansion. Empirically, with respect to exact convolutional NTK (CNTK) computation, our method achieves 106x speedup for approximate CNTK of 5-layer Myrtle network on CIFAR-IO dataset.
methods do not suffer from such limitations and are
Introduction
Infinite width limit has enabled fundamental understandings of deep neural networks by establishing a correspondeece kernel methods. In this limit, the network's function prior is Gaussian process [1- 3] and under gradient descent training with squared loss, the network behaves as a linearized function [4. 5]. Underlying these limit, core object is neural kernel which encapsulates architectural inductive prior in its functional form
[6]. The kernel describing gradient descent dynamics, the Neural Tangent Kernel (NTK) [4], and Neural Network Gaussian Process (NNGP) [2] kernel have been extensively studied [7-12] since they were initially identified. In particular, the infinite width theory has shed light on powerful abilities of deep neural networks including optimization [13-16], generalization [17-19], regularization [20-22] and robustness [23, 24]. Beyond theoretical findings, has been extensively reported that neural kernels can enhance practical applications including small data classification/reresssi tasks [25], neural architect search [26, 27]. dataset distillation [28, 29]. federated learning [30], meta learning [31], generalization attack [32], just to name few.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Despite those powerful advantages, there is still gap between practice and theory in the utility of these kernel methods. First, the NNGP and NTK can be exactly computed recursively [2, 4] however. the explicit forms are only known when the corresponding neural networks contain few set of activation functions such as ReI.U or Error functions. While ReLU activation is the default choice for many deep learning applications, recently different activation functions have shown to work well in various domains of machine learning. For example, GeLU [33] has been widely used in Transformer based natural language processing settings [34-36] and sinusoidal activation functions work well for implicit neural representation (e.g. NeRF) [37, 38]. Moreover, Xie et al. [39] showed that smooth activation functions could improve robustness compared to ReLU-based models. To enable better theoretical understanding on the role of these activation functions in these domain, expanding the infinite width limit tool set to general activation function is important step forward.
Secondly, even if the exact neural kernel computation is explicitly known, it requires significantly huge amount of computing resources. For example, it will take order of few 100 to 1,000 GPU hours to compute the exact NTK of depth 10 convolutional neural networks with pooling on 60,000 CIFAR-10 10 dataset. High compute requirement is often too expensive to perform extensive studies or use in a practical setting. While Novak et al. [40] have sped up Monte Carlo estimation of the NTK, random sampling remains impractical due to still high kernel computation cost, and cubic (in the training set size) inference cost. Recently, Zandieh et al. [41] proposed an efficient method to approximate the NTK computation via sketching algorithms. Their algorithm can approximate the neural kernels with ReLU activation orders of magnitude faster than the exact one. But it remains unclear how sketching algorithms are extended to other activations.
In this work, we fill this gap by showing that neural kernel for arbitrary smooth activation can be expressed in a form of series expansion. We first focus on how to express kernel function of neural network with a single hidden layer. Under the infinite width limit, this kernel converges to a static function, so-called a dual kernel, and is determined by activation in the network. This is key block to compute the NNGP and NTK of deeper architectures. We establish an explicit expression of dual kemel by expanding activation with the Hermite polynomial basis, and combining it with the fact that Hermite polynomials can play a role of random features of monomial kernels. As a result, our dual kernel formulation relies coefficients of series expansion of the activation. In addition, derive dual kernel expression of the first-order derivative of activation. The NTK can be computed by combining these kernel computations. To the best of our knowledge, our work is the first to study the computation of the NTK for general activations. Furthermore, we provide a subspace embedding for NNGP and NTK matrices with near input-sparsity runtime and near-optimal target dimension. As activation functions play an important role in modern neural network architectures, we hope our work could empower researchers to explore properties of activations in more principled way. Our main contributions are summarized as follows:
also
Building blocks for infinite-widhh neural kernel computations: We derive an explicit expression of the dual kernel for a polynomial activation, which can be building block for infinite-wine width neural kernel computations. For non-polynomial activation, we suggest to use its truncated Hermite expansion and analyze an error bound of the dual kernel.
Compiling and expanding dual activation Table 1: We compile various known dual kernel for point-wise activations providing pointers to the original work and expand the set further. We hope our work also serve as an easy reference for various analytic expressions. We emphasize that while many prior references lack required computation for NTK, this work is comprehensive in covering both NNGP/NTK transformations for various activations where analytic computation is possible.
Dual kernels of both activation and its derivative are essential for the NTK computation. Since our formulation requires coefficients of Taylor series of the activation, it is applicable to the dual kernel of derivative of the activation. In addition, we propose how to automatically compute the dual kernel of the derivative without knowing the activation. This approach is useful to characterize the NTK for kernel functions whose activation function is unavailable, e.g., normalized Gaussian, or whose dual kernel of the derivative is unavailable, e.g., GeLU and ELU.
NTK computation:
We analyze a pointwise error bound of approximated dual kernel via truncated Hermite expansion of the activation with a finite degree. The estimation error can decay polynomially faster in the degree. Furthermore, due to specific decomposition of our kernel formulation, we accelerate the NTK approximation by sketching techniques, similar to [41]. We also propose new sketching method for the Convolutional NTK with homogeneous activations
Kernel approximation:

---

Table I: Activation functions and references for their dual kernels. More detailed expressions are provided in Appendix F.
Reference for the NNGP
Reference for the NTK
Activation
(tt)
Rectified monomials
(9 (120)
[44]
[44]
Error function
erf(t)
[43]
[5]
ABReLU (Leaky ReLU)
-Amin(t, 0) max(t, 0)
[42, 50, 51]
[42, 50, 51]
Exponential
exp(At)
[46, 52]
[46, 52]
Hermite polynomials
(,)
[46]
This work
Sinusoidal
sin(A! B)
[45, 53]
This work
Gaussian
exp (-AF)
[43]
This work
GeLU
1/2 1 erf
[48]
This work
ELU
step(t)t+ + step -t) (e' -1)
[48]
This work
Normalized Gaussian
Unknown
[54]
This work
RBF
v2sin( 221 7)
[45]
This work
Gahor
exp( exp(-f)sin(t
This work
This work
Monomial
19
This work
This work
Polynomial
Σ)= a,t
This work
This work
and analyze both pointwise error bound and its runtime in Appendix D.2. Notably, our sketching method's runtime scales only linearly in the number of pixels of the input images, while the exact CNTK computation scales quadratically in the number of pixels.
Implementation: We sourc- NNGP and NTK for new activations within the Neural Tangents library [42] and sketching algorithm at https: //t::////iii /githt.....//.../ com comm/inuuaan/nee lorikacctiviiiin.
1.1
Related Work
Neural kernels (NTK, NNGP) can be computed using the recursive formula [2-5]. A prerequisite for these kernels is computing static kernel function which is defined as the expectation of some function of (non-linear) activation in neural network over the standard normal distribution. Williams [43] studied this dual kernel erf(t) and Gaussian. Cho and Saul [44] derived dual kernels for the rectified monomials, i.c., ta [e20). this function is equal to arc-cosine kernels where ReLU activation a special case when q Rahimi and Recht [45] showed that sinusoidal activations, e.g., sin or cos, can result in the Gaussian RBF kernel function using the Fourier transform. Daniely et al. [46] proposed a method obtain dual kernel if activation can be expanded by Hermite polynomials. However, inputs of the resulting kernels are restricted to be on the unit sphere. Louart et al. [47] analyzed asymptotic properties of dual kernel with random matrix theory and show closed-form formula of such as erf, 14]. sinusoidal. Tsuchida et al. [48] studied the dual kernels of both Gaussian Error Linear Unit (GeL.U) [33] and Exponential I. Linear Unit (ELU) [33]. For activation that does not admit closed-fome expression, Lee al. [2] numerically computed dual activation by doing interpolation on predetermined grid of variances and covariances. Table summarizes activations whose dual kernels were priorly known, well as expanding (in this work) the set previously unknown expressions. Recently, Simon et al. [49] discovered that NTK fully-connected neural network with any depth can be converted into that of 1 hidden-layer neural network by modifying activation function. However, their method is limited to the normalized input data and fully- calyycconnected networks.
2 Preliminaries
Notations. We denote the identity matrix of dimension by Ia- For scalar function f, we write (k)) to denote its k-th derivative. We use 1e to denote the indicator of event E. For smooth function F + F. use (A) to denote its k--th derivative and define |o for some
(o-N( (l))) || For scalar functions f.g we use f o
and simply write ||σ]|/(o.1) to denote the composition of these functions and re to denote the y times self-composition of f. e.g., (2:) [()([))))) Given positive semidefinite matrix K and X > 0, the statistical dimension of K with regularizer À is defined as 8x (x() == tr(K( (KK K XI) -1). We use nnz(x)
to denote the number of nonzero entries in Given E 1 and y 1 we define X y

---

[113), #221, ..... 11, ; 4 1122, Lmm22 Lm and 80 as the p-fold self-tensoring of z. We also
define as the direct sum between vectors.
Hermite polynomials.
The Probabilist's 's Hermite polynomials of degree E 2 is defined as
de
L8/2] il(E 21)! 2
(-1)²
he(1)= = (-1)fe
= 8!
(1)
di
The polynomials {he the form a set of orthogonal basis for the space of square- inteeraalee functions in IR with respect to the normal measure (N, 1), i.e., the L space of functions L (R,.V) == {f
R - R |o w oo}. Particularly, it holds that (0,1((0,11 [he(t) Imm [t)] = !! [e=m) Thus, any
function f L (B, has a unique Hermite expansion in the sense of Il. Ei= cyhy 0 and
coefficient ej be computed cj EL-A((.1)) [f(t) (t)
Infinite width neural kernels. a fully-connected L-layered neural network f: IRd + R for L 2 2 defined as'
Given an activation - satisfying that ||a||x 1, consider
(x; W) = w (w'",,,00 (11,,211) de ze ( 26-11 1/ (w(92)-1/
-/V 20
(2)
Rdixd do d, d; == m for I > 0 is a collection of learnable parameters, m is the width of the network, and c(-) is applied point-wisely, In the infinite width limit, i.c., m - oc, when all elements of W are initialized by ii.d. random samples from (0, 1) and optimized via gradient descent the least-square loss with an infinitesimal learning rate, the prediction of trained network becomes identical to that of its first order Taylor approximation at W. Hence, inference with such ultra-wide network is equivalent to kernel regression with a static kernel, the so-called Neural Tangent Kernel (NTK), defined as (x, 3) == plim,,-+o (V (z; W), Wwf W)) (convergence in probability to constant). In addition, at initialization the output of an infinitely wide network is equivalent to a sample from a Gaussian process with mean zero and covariance x(L) (I, == plimm (Ja(x; W), (y: W)), known as the Neural Network Gaussian Process (NNGP) kernel.
where W: vee ( W(O) for w(L)) E RR
Recursive expression for NNGP and NTK. Several previous works [2-5] have shown that the NNGP and NTK can be expressed using the following recursive procedure:
For every x, y E Rdd let K (a,, 8) == (x.y) and for every layer h 1,...,, L, recursively
define kernel functions k!h) ((hh) Rd x Rd R as:
kh(h) (I,y) =
E
(o))
K- (I,
E (1.))~/10015))))
[o'(u)o((v)), (3)
where the covariance matrix !!! = 2g (x,
K-M (x, v)
E 2x22
K (h--l) (y.r)
K!th- 1) (v.2)
2. The depth-L kernel is k. (zz y) and the depth NTK () as ())) (z. v) : (x. y) and
can be recursively computed
th) (x. = 0 1h-11) (x. (x, + K (x (z.g).
(4)
At the core of the expression for (.7)) there is the expectation term over 2-dimensional Gaussian distribution in Equation (3). This expectation term for the case where both diagonal entries of the covariance matrix (2) are equal to one, was previously studied in [46]. We extend this to encompass general symmetric covariance matrices in the following definition.
'Throughout the paper, we consider scllar-valued networks without biases for simplicity. but this can be extended to vector-valued networks with biases. We also assume ||o llw which does change results.

---

Definition 1 (Dual Activation and Dual Kernel), For a smooth a R + I, we define the Dual Kernel of as Kg R. F defined as
(z) == E
(ow ((..)
for every v. y E R
(5)
Equation (5) only depends on bivariate Gaussian random variables (we, æ)) (w, y) where E( xx) || B[ y) = lla/12 and E[(w, æ) (w, y>1 = (a. y). Hence one can look at the dual kernel from a different perspective by choosing a proper covariance matrix. To this end. Aa,be :=
a² abe
obe 8
for every a, E R+ and C [-1, 1] and the Dual Activation of a with respect to Aa,b.e the function ko IR+ x R2 [-1,1] - lR defined as h b, e) == (0..))) [=)))(())]
With these definitions in place, the following relationship between dual kernel and activation holds
(ar, y)
Ka y) = lio |101|2- lulll
(6)
| 1ull2
Observe that Kg (z. y) corresponds to the NNGP kernel of 1-hidden layer neural network with activation a. For some specific activations, e.g., ReLU, Error function, closed form expressions for their dual activations are known (see Table 1). Hence, one can compute the NTK analytically when dual kernels of the activation and its derivative have a closed form expression. The above also holds for kernels corresponding convolutional neural networks called CNN-GP [7, 8] and CNTK [9].
3 NNGP and NTK for Smooth Activations
In this section, we focus on the NNGP and NTK for a wide range of smooth activation functions. We first show that a series expansion for the dual kernel can be obtained from that of the activation function, which is key to NNGP kernel computation. By applying this result to the derivative of the activation function, we can also compute the NTK for the same activation.
3.1 Dual Kernel Computation
Daniely et al. [46] proved that for absolutely continuous a R - IR and any x, y E Sd-1 the dual kernel is equal to K2 (xr = Eド= c j! (z. y)² where {cr bio are coefficients of Hermite expansion of We now proceed generalize this result from g to entire F {0}. First we remark that it can be naturally extended to the dual kernel of g-homogeneous activation functions, i.e., (at) = ia|" (t)) for every a, & E R, on the entire Rd (0). For every x, y Rd {0), the corresponding dual kernel is
(x, v)
Kz v) = | 1
Egj jl
1=0 || ll 11yll,
As examples, (leaky) ReLU and rectified polynomials fall into this activation class.
Now suppose that a is not homogencous. In particular, we first consider a polynomial activation (t) = Σi a,ti with coefficients {az {= Recall that K, (z, y) can be obtained by taking the expectation of a ((w,, of Daniely et al. [46]'s result, we factorize the input into its radial and angular part and rewrite the activation by expressing monomials the Hermite polynomial basis. Formally, let us write
u))
NO, Ia) for every x, y Rdd {0). To make
monomials in the Hermite basis as Ei (t)) for some coefficients {u Then
9
Σ P9. || 11/2 ag hi ((u:
a(wm, xx))
Σ
aj
1001
W,
||01||
|1
|
(8)
1=0
Then, we can derive the dual kernel of polynomial activation. We further relax condition on the activation and propose the result below.
5

---

Theorem 1. For polynomial o(t) = EE a,th, the dual kernel of (.), as per Definition 1, is (x, (1.1(1==112) (a.e(llvll2) (TEEll llvll2
9
Kz(I,y) 1)
(9)
t=
where (t)e == E. = ae221(+22)
gi+t Moreover, if an activation function o : R - R satisfies
2i
1o < xc and l|o ã
(i( & for some & > 0 and v 2 1, then.for every n, E R2 such
that ||=||2 llvll2 E (0, v) the following holds
v2 @6 [6|aal (oa, 46
[Ko (z. y) - Ka v)1 S
(10)
||r||2 llv|l2
The proof of Theorem 1 is provided in Appendix B.2. For non-polynomial activations, one can consider approximating a with its Hermite Taylor expansion and then apply Theorem 1. Examples can be found in Appendix B.2. For activation functions that do not have Taylor expansion but are k-th order differentiable, we show that, using their Hermite expansion, one can obtain a good approximation to the corresponding dual kernel.
Theorem 2. Given R - R, suppose that there exists an integer k > 2 and some 1 such that for every = .k. g(0) is absolutely continuous and lim e T (vvt = 0 and moreover ||a [1a lix( < oc and || l[Ni0,u2) < <0.. Consider the Hermite expansion coefficients {cy)220 of function o(v)) and denote ā(1) = Σÿ- (5==0g);((1/)) Given x, E R.d with ||=/|l2 11vll2 E (0,
5 (|o(()) A((o.w2) max (101 w | (N
(K v) Kz(
(11)
V|aal2 lull2 k qh-1
where ()(-.)) and (x( ) are dual kernels corresponding to a()) and (.) in Definition 1, respec- rively. Moreover, for the ReLU activation (t) max(t, 0), it holds that
2
Ka Ka v) Ka( S
(12)
|| 11v/l22
The proof of Theorem is provided in Appendix B.3. Observe that when the activation is k-th order differentiable and the norms of its derivative and inputs are bounded then the approximation error decreases with O -
rate. In Section 5, we empirically evaluate the dual kernel of various
activations using Hermite expansion and verify that smooth activations (e.g., Gaussian or sinusoidal) provides much lower approximation errors than non-smooth ones (e.g., ReLU).
3.2 NNGP and NTK Computations
Once dual kernels of a and a' or their polynomial approximations are calculated, one can compute (approximate) NNGP and NTK using Theorem or Theorem 2 and the recursion in Equation (4). However, there are scenarios where we are only given the dual kernel and the corresponding activation or derivative of the activation is unknown to us. For example, Shankar et al. [54] devised normalized Gaussian kernel defined as
Ko(æ, y) = || Ilyll2 exp ( (Ir v)
-1)
(13)
| llvll2
and reported that NNGP with this dual kernel performs better than the ReLU NTK by showing promising results on various tasks. Note that, recovering the activation from KG is non-trivial. From the dual kernel perspective, the activation should be 1-homogeneous and its Hermite series expansion is of form ΣØ= =1 hy (t) and it is generally unknown how to choose the sign pattern on coefficients of this series that would satisfy homogeneity constraint, Instead of trying to recover the activation from dual kernel, we show how to directly derive the dual kernel of derivative of activation without knowing the activation.

---

Theorem 3. Given a differentiable activation function o R + R which satisfies la(t)) S
C Crexp exp Ti o o'( I exex TIVE |o < ac and lo"Il < oo for some
2 1 and constants C1 C2. the following holds for any 2, y E R with || llvlla (0, v/ and (x, v) < |||||| l1yll2
1
a
K (x,z)= =
to (112l12 1
(14)
||>|| llvll2 de
Additionally, if 2ko (c) ) continuous +1 then Equation (14) holds for x. such that (xr, v) = ||0|| 17ll2.
The proof Theorem 3 is provided in Appendix B.5, Our result is more general compared to [49] where the previous work assumes that the Hermite expansion of given activation should converge and
||ж||| = ||vl22- Applying Theorem 3 to Equation (13) provides that Ka(2, y) = exp
(a.ay)
hence one can compute the NTK function even if the corresponding activation is unknown. In the previous work [54], only "NNGP" performances of the normalized Gaussian kernel were reported.
Moreover, with Theorem 3, only the knowledge of dual activation suffices to compute both NNGP and NTK. For example, while dual activation (thus NNGP) of GeLU was known in Tsuchida al. [48], lig was not derived explicitly. Theorem provides a simple way to compute lig (given in Equation (126)) via automatic differentiation, without requiring to take the expectation under multivariate Gaussian distribution computing derivatives by hand. This is implemented in stax_extens ons El ementwise in our code supplement. Our method allows to omit the enitre effort, lines of code, and potential mistakes in deriving and implementing the NTK.
3.3 Gauss-Hermite Quadrature
One simple approach to obtain dual activation function for general activation functions without closed form expressions is to evaluate the expectation of under the 2d Gaussian distribution as numerical integration. This can be efficiently done by Gauss-Hermitt quadrature
1 ka)
ΣΣ wi wawl [o(V2ax,) a oVVbbcr, + π
- -2bV1-
(15)
where (x;; wi), correspond i-th root of degree q Hermite polynomial hx(()) and associated
weights [55] Wt (h,, (V See Appendix E for the derivation of the quadrature formula.
glv/
For smooth activation functions will quickly go down as q increases by Theorem 2. We use this method to compute approximate (non-sketched) kernels for general activation functions Figure and Figure 4. It is implemented stax_ extensions... E.l.entwiseeNmmrree.. in code.
Approximating Neural Kernels via Sketching
Although using our Theorem 1, Theorem and Theorem 3, one can analytically compute NTK for general activation functions, computing all entries in the NTK kernel matrix requires massive amount of resources, i.e., ((n + Lq²)) runtime and 2(n²) memory for datasets with n points in Rd. This becomes more expensive for CNTK, where its runtime can be Ω (((dn d2 )2))²( ( + L)² 2 for of images with size dj x d2 x c. To avoid quadratic complexities, we adopt a fast and efficient feature map construction via randomized sketching [41] for both NTK and NNGP, i.c..
( ( (v) (1)). (x))), = (z)
(16)
The previous approach was only applicable for the ReLU activation but we establish more general scheme based results for dual kemel approximation.
This is assuming Hermite expansion degree g, when exact expression is known q² is constant.
7

---

Algorithm 1 Subspace Embedding of Homogeneous NNGP and NTK
1: input: x E F depth L, sketching dimension m. polynomial k(t) E3=oaf with oj R+
2: calculate the polynomial (()) = (t) = Σj= === byte with coefficients by E R__
3: calculate the polynomial (t) = En= Koh (t) II
1-1 W o (t) = Σy ejt with coeffi-
cients Cj E R., and degree p g0(1.)
4: for f = O.....P.. let Q' E mxd be degree- POLYSKETCH (See Appendix A) P, p.u' u¹ 5: for every E = ~=0,.... O.....,pp
Q |irlaa
6: construct (( (x) ||ا|| θ) Vojud and wbt (x) ||ا||
D Voju
7: return (z) (NNGP embedding), (x) (NTK embedding)
Subspace embedding for homogeneous dual kernels. We provide subspace embedding for NNGP and NTK matrices with near input-sparsity runtime and near-optimal target dimension which applies to any homogeneous dual activation functions with rapidly convergent Taylor expansion. More specifically, we call a dual kernel K, homogeneous there exists a positive definite dot-product kernel function K [-1,1] -
[-1, 1] such that,
(xz, y)
K₂ (1 v) | l1ull2 is
(17)
[10/22 |1yll2
For such homogeneous dual kernels, the NTK and NNGP take a similar homogeneous form. In fact, one can show by induction that when the dual kernel is in form of Equation (17), the depth-L NNGP function defined in Equation (3) equal the following for any positive integer L.
K K (z. y) = ||>|| llvll2 x (
(or, v)
(18)
[alll 11lll2
where = denoted the L-fold composition of function k. Furthermore, if has derivative n' [-1,]] [-,11]+ [-11, 1). using Theorem 3, there exists depth-T. L NTK for this dual kernel, equal to
L-1
(x, = 1|=||| llvil2
(t)
II i=h
K' o (t)
(19)
h=
1
where we use the convention that x: (t) = t. Therefore, if k(-) can be tightly approximated by low-degree polynomial, then the NNGP and NTK functions can also be tightly approximated by low- degree polynomials. Thus, by applying POLYSKETCHS which is a norm- norm-preserving dimensionality reduction that can be applied to the tensor product of multiple vectors very quickly [56], to the polynomial approximations to these kernels, we can spectrally approximate the NNGP and NTK kernel matrices. For details on POLYSKETCH see Appendix A. We provide the details of this procedure in Algorithm and prove the correctness and runtime of our procedure in Theorem 4.
Theorem 4 (Homogeneous NTK Embedding). Suppose that the dual kernel K₂ homogeneous as per Equation (17). Also suppose k(t)) is a degree-q polynomial with non-negative coefficients that satisfies (1) max [k(t) x(t) S poly(n) and max (R'(t) (t) K6' (t) S (oly(n)' (2)
max
(a(t + 7) R(()]= S poly(n; and max
[t((+ (t) poly(n)
for any hy| polytny. Then for any integer L 2 any e, Z poly(xj' and any dataset X E Rddx
with |X|] 1F I poly (n), if Kitk E R=xx is the depth-L NTK kernel matrix on this dataset, there exists m = O (a((ain)) poly (gt, (g=,log)) such that the output wi(l) (X) Rmx of Algorithm I
satisfies with probability at least polyin)
(1 (K ntk AI < (T) (X) ( + XI (1 (1+ (Kuk + AL,)
(20)
Moreover, the runtime of Algorithm is O (poly (q log n) 6- (s> (ss(Kntk) Kntk) n nnz (X))).
We prove this theorem in Appendix C. As example, let us apply Theorem 4 on the normal- ized Gaussian kernel Kg defined in Equation (13), which is homogeneous. The dro-product fac- exp(t 1). The truncated Taylor series of this
for corresponding to this dual kernel is s(t)
8

---

ReLU
Abs
Sin(d)
eriit
GeLU
16mk
Permit
10
100
10
10-1
10-2
16
10
10
10
10'0
10
10-4
1.0 Elapsed time
1.0 Flapsed
3.0
1.2 Elapsad
0.5
Elapsed
Elapsad time
Elapsed (9ec)
ReLU
Abs
Sintt)
orttt)
GeLU
10-1
Harmie
10-
10
10-1
2012
10
10
10-11
Degroe
Degree
Degroe
Degree
Degree
Degree
Figure 1: Relative errors of dual kernel approximations via the truncated Hermite expansion and Monte Carlo estimation under synthetic dataset with = 1,000, 256.
2(loggn) n) then it can be verified that the polynomial R(t) satisfies the preconditions of Theorem 4. Therefore, one can invoke Algorithm to get sub- space embedding for the NTK kernel matrix corresponding to the normalized Gaussian dual kernel KG in O E=- (sx(
function R(t) = O e.j.- If q
nnz(X) poly (oog" ))) time and with target dimension of 2 sx (Knk) poly (log" mn) )- For any constant number of layers, L, this runtime and target dimension is is optimal up to poly (log factors. The implementation of our sketching algorithm is available at https:://// /eti::/gittttt //gtttub..... homm/nsuhaan/nen saapanntkkkacccae thaahhn//kk.cciivttioonss activatiovi.
Experiments
In this section, we perform experiments with the proposed neural kernels based on our dual kernel approximation. All experiments run using a single A100 GPU machine.
Kernel approximation. We first benchmark our algorithm to approximate the dual kernel ma- trix. We use ReLU, Abs (i.e., (t) = [t), sin, Gaussian, erf and GeLU activations and ap- proximate them by their Hermite expansion where degree changes from = to 20. We ran- domly generate n 1,000 of 256-dimensional inputs where each entry is i.i.d. drawn from 1/0.1/V256) We also compare our approach to Monte Carlo estimation of dual kernel, i.e. K2 (x. R -E=,a(( (.), (o) (w,, u)) where ) are i.i.d. standard Gaussian vectors. In Figure 1, we plot relative errors of the Frobenius norm kernel approximations in terms of wall-clock times (top) and polynomial degree (bottom). We run 10 independent trials and evaluate the average approxmation errors. We observe that our approximation with Hermite expansion outper- forms the Monte Carlo method for all activations approximated because they are smooth and norms of their derivatives are bounded with respect to the normal measure.
((w,,
used. In particular, sin and Gaussian are well
Performance on CIFAR-10 classification.
We also benchmark the proposed CNTK ap- proximating via sketching algorithm. We per-
Myrtle5
Functions
1.0
form CIFAR-10 classification [57] by solving the ridge regression problem. The image classes converted into 10-dimensional one-hot vec- tors and inputs are pre-processed with regular-
Accuracy
0.75
0.70
0.65
expL - 3))
0.4
0.50
expíć 1)
ABRALLI
0.2
ized ZCA [54, 58]. We report the best test ac- curacy among 20 choices of ridge parameters
0.55
Relli
0.0
2 2
2
in 100-10+ 1
= 0, 1. ... 19}. We extract CNTK features of 5-layer convolutional neural network (known as Myrtle5 [54]) without pool-
Feature Dimension
Figure 2: Test accuracy of CIFAR-10
ing by setting degree q = 8 and explore feature dimension mn = {2", (2",..., 224) and homogeneous dual kernels including ReLU, ABReLU activations well as deep normalized Gaussian kernels with 2 scaling factors. See Appendix for more details. Figure 2, the test accuracy of neural
kernels (left) and the corresponding their dual activations (right) are plotted. The dual activation of

---

ABReLU is very similar to the normalized Gaussian without scaling and their test performances are also comparable. We observe that the scaled normalized Gaussian shows the best performance which achieves 78. 13% while the ReLU CNTK features [41] shows 75.56% with the same runtime. This because the coefficients decay of the normalized Gaussian is faster than that of the ReLUU, which leads to a lower approximation error of sketching algorithm. We also perform comparison among different activation functiins in neural kernels in Appendix E.
Speedup. We observe that the exact CNTK of Myrtle-5 constructs a kernel matrix of size 60,000 x 60,000 and achieves 86-87% test accuracy. However, this requires approximately 151 GPU hours. Under the same setting, our CNTK features for the normalized Gaussian kernel take about 1.4 GPU hours, i.e. a 106x speedup. If we use less training data to construct 20,000 x 20,000 kernel matrix, the accuracy is about 77% accuracy and the runtime is 16.8 GPU hours in which our approximation is still 12x faster without loss of accuracy. We believe such acceleration through methods open the door to using neural kernels in a wide range of research domains.
6 Discussion
In this work, we introduced methods to efficiently compute neural kernels for general activations. As activation functions play an important role modern neural network architectures, we hope our work could empower researchers explore properties of activations in more principled way. We are excited with sketching method's compute efficiency by orders of magnitude on highly performant neural kernels to open up applications in dataset distillation [29] or uncertainty critical problems [59] such as autonomous driving, healthcare and science.
Acknowledgements
Amir Zandich was supported by the Swiss NSF grant No. P2ELP2 195140. Insu Han and Amin Karbasi acknowledge funding in direct support of this work from NSF (IIS- 1845032), ONR (N00014- 19-1-2406), and the AI Institute for Gearning-Enablee Optimization Scale (TILOS). We thank Timothy Nguyen and Jeffrey Pennington for discussions and feedback on the project.
References
[1] Radford M. Neal. Priors for infinite networks. University of Toronto, 1994.
[2] Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep Neural Networks as Gaussian Processes. In International Conference on Learning Representations (ICLR), 2018.
[3] Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard Turner, and Zoubin Ghahramani. Gaussian Process Behaviour in Wide Deep Neural Networks. In International Conference on Learning Representations (ICLR), 2018.
[4] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Neural Information Processing Systems (NeurIPS), 2018 8.
[5] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Neural Information Processing Systems (NeurIPS), 2019.
[6] Lechao Xiao. Eigenspace restructuring: principle of space and frequency in neural networks. In Conference on Learning Theory (COLT), 2022.
[7] Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A. Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian Deep Convolutional Net- works with Many Channels are Gaussian Processes. In International Conference on Learning Representations (ICLR), 2019.
10

---

[8] Adria Garriga-Alone Alonso. Laurence Aitchison, and Carl Edward Rasmussen. Deep convolutional networks as shallow Gaussian processes. In International Conference on Learning Representa- lions (ICIR), 2019.
[9] 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 Neural Information Processing Systems (NeurIPS), 2019.
[10] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: NNGP and NTK for deep attention networks. In International Conference on Machine Learning (ICML), 2020.
[11] Simon S Du, Kangcheng Hou, Russ R Salakhutdinov, Barnabas Poczos, Ruosong Wang, and Keyulu Xu. Graph Neural Tangent Kemel: Fusing Graph Neural Networks with Graph Kernels. In Neural Information Processing Systems (NeurIPS). 2019.
[12] Greg Yang. Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation. arXiv preprint 1rXiv:190..04760, 2019.
[13] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In International Conference on Machine Learning 2018.
(ICMI.,, 2018.
[14] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over-parameeterizatoo. In International Conference on Machine Learning (ICML), 2019.
[15] Mert Pilanci and Tolga Ergen. Neural networks are convex regularizers: Exact polynomial-time -ime convex optimization formulations for two-layer networks. In International Conference on Machine Learning (ICML), 2020.
[16] Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and generalization in deep neural networks. In International Conference on Machine Learning (ICML), 2020.
[17] Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. The role of pve-parramttrization in generalization of neural networks. In International Conference on Learning Representations (ICLR), 2019.
[18] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning (ICML), 2019.
[19] Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Neural Information Processing Systems (NeurIPS), 2019.
[20] Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Neural Information Processing Systems (NeurIPS). 2019.
[21] Wei Hu, Zhiyuan Li, and Dingli Yu. Simple and effective regularization methods for training on noisily labeled data with generalization guarantee. In International Conference on Learning Representations (ICLR), 2020.
[22] Arthur Jacot, Berfin Simsek, Francesco Spadaro, Clément Hongler, and Franck Gabriel. Implicit regularization of random feature models. In International Conference on Machine Learning (ICML), 2020.
[23] Elvis Dohmatob. Fundamental tradeoffs between memorization and robustness in random features and neural tangent regimes. arXiv preprint arXiv:2/06.02630, 2021.
[24] Hamed Hassani and Adel Javanmard. The curse of overparametrization in adversarial train- ing: Precise analysis of robust generalization for random features regression. arXiv preprint 2200:2200. 2022.
11

---

[25] Sanjeev Arora, Simon S Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the Power of Infinitely Wide Deep Nets on Small-data Tasks. In International Conference Learning Representations (ICLR), 2019.
[26] Daniel S Park, Jaehoon Lee, Daiyi Peng. Yuan Cao, and Jascha Sohl-Dickstein. Towards nngp-guided neural architecture search. arXiv preprint arXiv:201 2020.
[27] Wuyang Chen, Xinyu Gong, and Zhangyang Wang. Neural Architecture Search on ImageNet in Four GPU Hours: A Theoretically Inspired Perspective. International Conference on Learning Representations (ICLR), 2021.
[28] Timothy Nguyen, Zhourong Chen, and Jaehoon Lee. Dataset Meta-Learning from Kernel Bidge-Regression. In International Conference Learning Representations (ICLR), 2021.
[29] Timothy Nguyen, Roman Novak, Lechao Xiao, and Jaehoon Lee. Dataset distillation with infinitely wide convolutional networks. In Neural Information Processing Systems (NeurIPS), 2021.
[30] Baihe Huang, Xiaoxiao Li, Zhao Song, and Xin Yang. Fl-ntk: A neural tangent kernel-based framework for federated learning analysis. In International Conference on Machine Learning (ICML), 2021.
[31] Yufan Zhou, Zhenyi Wang, Jiayi Xian, Changyou Chen, and Jinhui Xu. Meta-learning with neural tangent kernels. International Conference on Learning Representations (ICLR), 2021.
[32] Chia-Hung Yuan and Shan-Hung Wu. Neural tangent generalization attacks. In International Conference Machine Learning (ICML), 2021.
[33] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint /606. 15, 2016.
[34] Jacob Devlin, King-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Conference of the North American Association for Computational Linguistics (NAACL), 2018.
[35] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl Blog, 1(8):9, 2019.
[36] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Neural Information Processing Systems (NeurIPS), 2020.
[37] Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Neural Information Processing Systems (NeurIPS), 2020.
(38] Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Pridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Neural Information Processing Systems (NeurIPS), 2020.
[39] Cihang Xie, Mingxing Tan, Boqing Gong, Alan Yuille, and Quoc V Le. Smooth adversarial training. arXiv preprint arXiv:2006. 14536, 2020.
[40] Roman Novak, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Fast Finite Width Neural Tangent Kernel. In International Conference on Machine Learning, 2022. URL https: / ctctluub..ommoooggggee wttpe//enurrllaaaae tangents.
[41] Amir Zandieh, Insu Han, Haim Avron, Neta Shoham, Chaewon Kim, and Jinwoo Shin. Scaling Neural Tangent Kernels via Sketching and Random Features. In Neural Information Processing Systems (NeurIPS), 2021.
12

---

[42] Roman Novak, Lechao Xiao, Jiri Hron, Jachoon Lee, Alexander A. Alemi, Jascha Sohl- Dickstein, and Samuel S. Schoenholz. Neural Tangents: Fast and Easy Infinite Neural Networks in Python. In International Conference on Learning Representations (ICLR), 2020. URL https://gi htt::////gtub...... com/ wttccooo//o/eeeeecoo -eaual-tangents..
[43] Christopher Williams, Computing with infinite networks. In Neural Information Processing Systems (NeurIPS), 1996.
[44] Youngmin Cho and Lawrence Saul. Kernel methods for deep learning. In Neural Information Processing Systems (NeurIPS), 2009.
[45] Ali Rahimi and Benjamin Recht. Random Features for Large-Scale Kernel Machines. In Neural Information Processing Systems (NeurIPS), 2009.
[46] Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and dual view on expressivity. In Neural Information Processing Systems (NeurIPS), 2016.
[47] Cosme Louart, Zhenyu Liao. Romain Couillet, et al. A random matrix approach to neural networks. The Annals of Applied Probability, 2018.
[48] Russell Tsuchida, Tim Pearce, Chris van der Heide, Fred Roosta, and Marcus Gallagher. Avoiding Kernel Fixed Points: Computing with ELU and GELU Infinite Networks. Conference on Artificiil Intelligence (AAAI), 2021.
[49] James B Simon, Sajant Anand, and Michael R DeWeese. Reverse Engineering the Neural Tangent Kernel. arXiv preprint arXiv:2106.03186, 2021.
[50] Russell Tsuchida, Fred Roosta, and Marcus Gallagher. Invariance of weight distributions in rectified MLPs. In International Conference on Learning Representations (ICLR), 2018.
[51] Russell Tsuchida, Fred Roosta, and Marcus Gallagher. Richer priors for infinitely wide multi- layer perceptrons. arXiv preprint arXiv: 1eXXv:1911..1227. 2019.
[52] Julien Mairal, Piotr Koniusz, Zaid Harchaoui, and Cordelia Schmid. Convolutional kernel networks. Neural Information Processing Systems (NeurIPS), 2014.
[53] Tim Pearce, Russell Tsuchida, Mohamed Zaki, Alexandra Brintrup, and Andy Neely. Expressive priors in Bayesian neural networks: Kernel combinations and periodic functions. In Conference on Uncertainty in Artificial Intelligence (UAI), 2020.
[54] Vaishaal Shankar, Alex Fang, Wenshuo Guo, Sara Fridovich-Keil Jonathan Ragan-Kelley, Ludwig Schmidt, and Benjamin Recht. Neural kernels without tangents. In International Conference Machine Learning (ICML), 2020.
[55] Milton Abramowitz, Irene A Stegun, and Robert H Romer. Handbook of mathematical functions with formulas, graphs, and mathematical tables, 1988.
[56] Thomas Ahle, Michael Kapralov, Jakob BT Knudsen, Rasmus Pagh, Ameya Velingker, David P Woodruff, and Amir Zandieh. Oblivious sketching of high-degree polynomial kemels. In Symposium on Discrete Algorithms (SODA), 2020.
[57] Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
[58] Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Neural Information Processing Systems (NeurIPS), 2020.
[59] Ben Adlam, Jachoon Lee, Lechao Xiao, Jeffrey Pennington, and Jasper Snock. Exploring the Uncertainty Properties of Neural Networks' Implicit Priors in the infinite-Widtt Limit. In International Conference on Learning Representations (ICLR), 2021.
[60] Ryan O'Donelll. Analysis of boolean functions. Cambridge University Press, 2014.
13

---

[61] Shuhuang Xiang. Asymptotics on Laguerre or Hermite polynomial expansions and their applications in Gauss quadrature. Journal of Mathematical Analysis and Applications, 2012.
[62] Achim Klenke, Probability theory: Media, 2013.
comprehensive course. Springer Science Business
[63] Gene H Golub and John H Welsch. Calculation of Gauss quadrature rules. Mathematics of computation, 1969.
[64] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Couraapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C Carey, ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E, A. Quintero, Charles Harris, Anne M. Archibald, Antônio Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-277 2020. doi: 10. 5592-019-0666-2
[65] Peter Jäckel. A note on multivariate Gauss-Hermice quadrature. London: ABN-Amro. Re, 2005.
[66] Andrew Maas, Awni Hamnun, Andrew Y Ng, et al. Rectifier nonlinearities improve neural network acoustic models. In International Conference Learning Representations (ICLR), 2013.
[67] Ben Adlam, Jake Levinson, and Jeffrey Pennington. A random matrix perspective on mixtures of nonlinearities in high dimensions. In Conference on Artificial Intelligence and Statistics (AISTATS), 2022.
[68] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. International Conference on Learning Representations (ICI.R, 2020.
[69] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint 19Xiv:1907.11699 2019.
[70] Djork-Aréé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep net- work learning by exponential linear units (elus). In International Conference on Learning Representations (ICLR), 2016.
14

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [No]
(c) Did you discuss any potential negative societal impacts of your work? [No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] All detailed proofs in the paper are provided in supplement.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
We provide anonymous code of some of our implementation.
(d) Did you discuss whether and how consent using/curating? [N/A]
obtained from people whose data you're
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "h0Qzos4kz9l", "reviewer": "Reviewer_LAi7", "summary": "\nIn this paper, the authors try to accelerate the computation of NTK with general activation functions. For the q-homogeneous dual kernel,  the authors extend the finite Hermite polynomials approximation technique (Daniely et al. 2016)  and provide approximation bounds. A sketching method is further proposed for the approximation of NTK.   The proposed sketching technique extends the technique (Zandieh et al.  2021) to homogeneous kernels (w.r.t more general activations) with rapidly convergent Taylor expansion. ", "questions": "1. Does the approximation of NTK with general activation guarantee the global convergence property?\n\n2. How to compute $L$-depth NTK recursively using Eq.(5) ? It is not clearly described in the context.\n\n3. In the experiments,  what is the concrete procedure to compute the time in Figure 1?  It is better to compare with orthogonal random features (Yu et al. 2016) or other fast kernel approximation techniques and present the error vs. number of features besides the error vs. time.\n\n\nYu et al. Orthogonal Random Features. NeurIPS  2016. \n\nZandieh et al. Scaling Neural Tangent Kernels via Sketching and Random Features. NeurIPS 2021.\n\nDaniely et al. Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity. NeurIPS 2016.", "limitations": "NA", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nPros.\n1. This paper focuses on NTK with general activation functions, which is an important direction for understanding deep neural networks with general activation functions. \n2. A method (Theorem 3) is proposed to compute the dual kernel of the derivative of activations without knowing the activations. \n3. The authors provide approximation bound of finite Hermite polynomials approximation to the q-homogeneous dual kernel. \n4. The sketching approximation of NTK extends (Zandieh et al.  2021) to more general homogeneous kernels. \n5. The paper is well organized and well written.\n\n\n\nCons. \n1. An important property of NTK and other prior work with ReLU is the global convergence of training. This relies on the strictly positive definite property in NTK (with ReLU) or positive-definite of the Gram matrix. This may be part of the reason why ReLU is widely used in the assumption of prior work.   This paper focuses on the approximation of NTK with general activation functions; however, it may lose the global convergence property.\n2. It seems that the fast sketching approximation technique (Sec.4) can only apply to $K_{\\sigma}$ homogeneous dual kernel. The approximation in Eq.(15) is a straightfowrad Gauss-Hermite quadrature for each entry of the kernel $K$, which is brute-force and still with high complexity. There are still some steps needed before claiming to fill the gap for general activations.\n3. In Figure 1,  how to compute the time? Is the sketching approximation technique more efficient than orthogonal random features or other fast kernel approximation techniques?\n4. It is better to state more clearly the key technique challenging/improvement compared with (Zandieh et al. 2021).  ", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "cghFxFA5RJ", "reviewer": "Reviewer_S3dK", "summary": "The Neural Tangent Kenrel offers a compelling framework to (partially) understand some theoretical aspects of neural networks, especially near initialization. In the infinite-width regime, It has been known that NTK with ReLU activation has a closed-form analytic formula which enables exact computation of the NTK in this limit. However, the exact kernel computation for general activations is intractable. This paper proposes an approximate kernel computation method that leverages some tools from functional analysis (Hermite decomposition).\n", "questions": "- There is a similar NTK decomposition in the “Spherical Harmonics“ basis that appeared in [1] and [3]. Moreover, authors of [2] show a similar extension of the decomposition on S^d to the whole of R^d (Theorem 5 in [2]). Knowing that Spherical harmonics are linked to Legendre polynomials, it would be great if the authors could discuss similarities and differences with Hermite decomposition.\n- A recent work ([4]) has showed that NTK computation (in finite-width setting) can be considerably accelerated with some computational tricks. Can the authors comment on the differences (a part from the finite versus infinite width settings of course). For instance, what is the difference between a kernel computed with [4] for width 100 and an equivalent NTK computed with the methods specified in this paper? What can you say about training/inference time?\n- In the proof of Thm3, the condition on the second derivative can be further weakened to something like $|\\sigma''(t)| \\leq C_1 \\exp(C_2 t^2 h(t))$ where $\\lim_{t \\to 0 }h(t) = 0$ and $h$ satisfies some mild conditions.\n\n\n[1] Geifman et al (2020) “On the Similarity between the Laplace and Neural Tangent Kernels”\n\n[2] Hayou et al. (2021) “Mean-field Behaviour of Neural Tangent Kernel for Deep Neural Networks”\n\n[3] Hayou et al. (2021) “Stable ResNet”\n\n[4] Novak et al. (2021) “Fast Finite Width Neural Tangent Kernel”\n\n[5] Bietti 2021 “APPROXIMATION AND LEARNING WITH DEEP CONVOLUTIONAL MODELS: A KERNEL PERSPECTIVE”", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n- The paper is very well written and flows nicely. The problem is well introduced and the results are stated in a logical sequence which serves well the purpose of the paper.\n- The theoretical results are sound and \n- The proposed methods are theoretically supported and the numerical results show some superiority compared to traditional methods used for NTK computation.\n- The proofs are well structured and easy to follow. I have one question about the proof of Thm3, see Questions section.\n\n**Weaknesses**:\n- I believe that it would benefit the paper if approximation rates are provided for general depth L networks instead of just depth 1 networks. The dependency on L would capture the growth of the approximation error as depth increases and maybe suggest when to use (or not use) one method or another. Besides, existing work such as [2] show that NTK tends to deteriorate with depth, I am keen to see what happens to the approximate NTK in this case.\n- The paper lacks a more comprehensive literature review. I believe a more in depth discussion of existing results would immensely benefit the paper. See questions section.\n- Numerical results are only reported for CIFAR10 (a part the synthetic dataset). It would be better if other (small, not necessarily ImageNet) datasets are included, this would help conclude wether activations such as the exponential or ABRelu is consistently better than ReLU. Also, results for varying depth would add value to the experiments section and make more complete.\n\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "MGDMHHQWBzm", "reviewer": "Reviewer_4Kjx", "summary": "The NTK and NNGP are powerful models for understanding and emulating the performance of neural networks in certain conditions. They are also worthy predictors in their own right, and allow one to easily predict with principled uncertainty. These models require computations involving a kernel or covariance function, which is a certain Gaussian integral representing the covariances between activations for two different inputs. Closed form expressions for kernels are known for many activations, each one having been derived by hand in previous works (see table 1). \n\nThe authors introduce a method for approximating NNGP kernels via a polynomial series (Theorem 1, 2). Furthermore, a certain expression in the NTK can be related to the NNGP via a derivative (Theorem 3), so the results readily extend to the NNTK. A random sketching method is also provided to approximate large kernel matrices of deep networks. \n\nThe paper is concluded with some experiments that show that the methods quickly (in terms of both wall-clock and number of series terms) and accurately approximate known NNGPs/NTKs. The authors also observe a 106 x speedup on a toy task compared with the exact kernel evaluation, and show reasonable performance.", "questions": "- In Theorem 2, the derivatives need to be absolutely continuous. However, the first derivative of the ReLU is not absolutely continuous (it is not even defined in a classical sense at the origin). So presumably the first half of Theorem 1 does not apply to the ReLU. How then does (12) hold?\n- What is the condition required to swap the order of the derivative and the expectation in line 672? Differentiating under the integral is valid when...? The second derivative of the ELU is unbounded, and also does not satisfy your weaker condition (44). Therefore Theorem 3 cannot be applied to the ELU. I don't think this condition is actually required. Differentiation under the integral holds in the sense of generalised functions/distributions/Schwartz functions, as long as the resulting integral is finite. This means that the derivatives will involve Heaviside, Dirac, derivative of Dirac etc.\n    - Looking at this again, your Theorem 3 is actually almost identical to part of Theorem 6 of reference [44] that you cite. Their expression for \\lambda_3 is your expression for (14), up to a normalisation. Their proof is very similar to yours, but more appropriately uses a Stein's lemma for tempered distributions i.e. Schwartz functions to differentiate under the expectation. I think it is fine to keep your Theorem 3 as is, but acknowledge that this was independently studied in [44] with a much milder condition using Schwartz functions, albeit for a different purpose. Alternatively, cite or reproduce this more general theorem.", "limitations": "I do not envisage any *direct* societal complications from this work. I hope it is not a cop-out to say that evaluation of societal impact is not necessary for this work, because the contributions are completely task, problem, data and even to some degree model agnostic. Downstream works that use the techniques introduced here should consider these impacts. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The method is generally applicable, i.e. the conditions in Theorems 1 and 2 are sufficiently mild a wide range of activations of interest. However, I am confused about the condition in Theorem 2 (see question below).\n- The authors cite a good amount of relevant literature, which is nice to see. The authors might also be interested in \"Stationary Activations for Uncertainty Calibration in Deep Learning\", which places the Matern kernel in the context of NNGP kernels.\n\nWeaknesses:\n- I did some random sampling of the proofs in the appendix. It is quite possible that I missed some errors. I have some concerns. See \"Questions\" box.\n- I found moving from Definition 1 to equation (5) a bit unnatural. In fact, the original motivation for the NNGP is as either (a) an inner product in an infinitely wide hidden layer with iid zero mean Gaussian weights or (b) the covariance of the output of a network. This is equation (5). Now (5) only depends on a bivariate gaussian (u,v)=(w.x, w.y), hence implying definition 1. Going in this direction is more natural IMO. In contrast, currently the authors go in the other direction and state \" By generalizing this observation to vectors in R^d\", without any apparent motivation for the generalisation. Note also that it is not actually a generalisation but an iff. A random vector is Gaussian iff all linear combinations are Gaussian. So while it is in principle possible to go from definition 1 to (5) without \"generalizing this observation to vectors in R^d\", it is not really required to be this abstract. I suggest the authors start with (5) and move to definition 1. That being said, there is nothing technically wrong with the current presentation apart from the word \"generalizing\".\n\n\nTypos:\n- \"NTK can be exactly computed using recursively\" removing \"using\".\n- Change \"We believe such acceleration thru our methods would open the door to using neural kernels in a wide range of research domains.\" to \"We believe such acceleration through our methods open the door to using neural kernels in a wide range of research domains.\"\n\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "8ZLgc7BlyQ", "reviewer": "Reviewer_vFLm", "summary": "This paper builds and extends the theory of infinite-width neural kernel computations. Compared with prior works, we can summarize its contribution as follows:\n(1) It explicitly computes the NNGP and NTK kernels for a wider range of activation functions as shown in table 1.\n(2) It derives an explicit expression of the dual kernel for polynomial activation, which also expends the assumption $x, y \\in S^{d-1}$ to a more general one $x, y\\in R^d$, and suggests using truncated Hermite expansion to approximate the dual kernel for non-polynomial activations.\n(3) Given the situation that sometimes we are only given the dual kernel but not the corresponding activation or derivative of the activation functions, this paper propose to compute the dual kernel of the derivative without knowing the activation as shown in theorem 3. \n(4) It uses random sketching techniques to accelerate NTK approximation, which is an extention to previous work which was only applicable for ReLU activation.", "questions": "I'm not familar with neural kernels before, so it's hard for me to provide some useful suggestions for this paper.", "limitations": "The limitations and potential negative social impact are not described since this is a theoretical work. ", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\n* This theoretical work extends the previous studies of dual kernels (NTK and NNGP), which is of high quality and might be useful as a good reference for future research in this area.\n\n* Most theorems proved in this paper are incremental, but still preserves enough significance. \n\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept