Fast Neural Kernel Embeddings for General Activations
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 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 homogeneous dual activation functions with rapidly convergent Taylor expansion.
Empirically, with respect to exact convolutional NTK (CNTK) computation, our method achieves 106 × speedup for approximate CNTK of a 5-layer Myrtle network on CIFAR-10 dataset.
Introduction
Infinite width limit has enabled fundamental understandings of deep neural networks by establishing a correspondence to kernel methods.
In this limit, the network’s function prior is a Gaussian process [1– 3] and under gradient descent training with squared loss, the network behaves as a linearized function [4, 5].
Underlying these limit, a core object is a 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, it has been extensively reported that neural kernels can enhance practical applications including small data classification/regression tasks [25], neural architect search [26, 27], dataset distillation [28, 29], federated learning [30], meta learning [31], generalization attack [32], just to name a few.
Despite those powerful advantages, there is still a 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 a few set of activation functions such as ReLU 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 an 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 ,
CIFAR-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 a 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 a key block to compute the NNGP and NTK of deeper architectures.
We establish an explicit expression of dual kernel 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 on coefficients of series expansion of the activation.
In addition, we also 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 a more principled way.
Our main contributions are summarized as follows:
• Building blocks for infinite-width neural kernel computations : We derive an explicit expression of the dual kernel for a polynomial activation, which can be a building block for infinite-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.
• NTK computation :
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.
• Kernel
approximation : 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 a new sketching method for the Convolutional NTK with homogeneous activations
[Table]
and analyze both a 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 open-source NNGP and NTK for new activations within the Neural Tangents library [42] at github.com/google/neural-tangents , and the sketching algorithm at github.com/insuhan/ntk_activations .
1.1 Related Work
Neural kernels (NTK, NNGP) can be computed using the recursive formula [2–5].
A prerequisite for these kernels is computing a 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 a dual kernel of erf( t ) and Gaussian.
Cho and Saul [44] derived dual kernels for the rectified monomials, i.e., t q 1 { t ≥ 0 } , this function is equal to arc-cosine kernels where ReLU activation is a special case when q = 1 .
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 to obtain a 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 , | t | , sinusoidal.
Tsuchida et al.
[48] studied the dual kernels of both Gaussian
Error Linear Unit (GeLU) [33] and Exponential Linear Unit (ELU) [33].
For activation that does not admit a closed-form expression, Lee et al.
[2] numerically computed dual activation by doing interpolation on predetermined grid of variances and covariances.
Table 1 summarizes activations whose dual kernels were priorly known, as well as expanding (in this work) the set to previously unknown expressions.
Recently, Simon et al.
[49] discovered that NTK of fully-connected neural network with any depth can be converted into that of a 1 hidden-layer neural network by modifying activation function.
However, their method is limited to the normalized input data and fully-connected networks.
Preliminaries
Notations.
We denote the identity matrix of dimension d by I d .
For a scalar function f , we write f ( k ) to denote its k -th derivative.
We use 1 E to denote the indicator of event E .
For a smooth function σ : R → R , we use σ ( k ) to denote its k -th derivative and define ∥ σ ∥ 2 N (0 ,ν 2 ) :
= E t ∼N (0 ,ν 2 ) [ | σ ( t ) | 2 ] for some ν ∈ R and simply write ∥ σ ∥ N (0 , 1) :
=
∥ σ ∥ N .
For scalar functions f, g we use f ◦ g to denote the composition of these functions and f ◦ q to denote the q times self-composition of f , e.g., f ◦ 3 ( x ) =
f
( f ( f ( x ))) .
Given a positive semidefinite matrix K and λ > 0 , the statistical dimension of K with regularizer λ is defined as s λ ( K ) :
= tr ( K ( K + λ I ) − 1 ) .
We use nnz( x ) to denote the number of nonzero entries in x .
Given x ∈ R m and y ∈ R n , we define x ⊗ y :=
[ x 1 y 1 , x 2 y 1 , . . .
x m y 1 , x 1 y 2 , . . .
x m y 2 , . . .
x m y n ] and x ⊗ p as the p -fold self-tensoring of x .
We also define ⊕ as the direct sum between vectors.
Hermite polynomials.
The Probabilist’s Hermite polynomials of degree ℓ ≥ 0 is defined as
[Equation]
The polynomials { h ℓ } ℓ ≥ 0 form a set of orthogonal basis for the space of square-integrable functions in R with respect to the normal measure N (0 , 1) , i.e., the L 2 space of functions L 2 ( R , N ) :=
{ f :
R → R | ∥ σ ∥ 2 N < ∞} .
Particularly, it holds that E t ∼N (0 , 1) [ h ℓ ( t ) h m ( t )] =
ℓ ! · 1 { ℓ = m } .
Thus, any function f ∈ L 2 ( R , N ) has a unique Hermite expansion in the sense of ∥ f − (cid:80) ∞ t =
c j h j ∥ N = 0 and coefficient c j can be computed as c j = E t ∼N (0 , 1) [ f ( t )
h j ( t )]
/j ! .
Infinite width neural kernels.
Given an activation σ : R → R satisfying that ∥ σ ∥ N = 1 , consider a fully-connected L -layered neural network f : R d → R for L ≥ 2 defined as
[Equation]
l − 1 (cid:17) , z 0 = x (2) where W := vec (cid:0) w ( L ) , ∪ L − 1 ℓ =1 W ( ℓ ) (cid:1) for w ( L ) ∈ R d L − 1 , W ( ℓ ) ∈
R d ℓ × d ℓ − 1 , d 0 :
= d, d l := m for l > 0 is a collection of learnable parameters, m is the width of the network, and σ ( · ) is applied point-wisely.
In the infinite width limit, i.e., m → ∞ , when all elements of W are initialized by i.i.d.
random samples from N (0 , 1) and optimized via gradient descent on 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 Θ ( L ) σ ( x, y ) :
=
plim m →∞ ⟨∇ W f σ ( x ; W ) , ∇ W f σ ( y ; W ) ⟩ (convergence in probability to a 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 Σ ( L ) σ ( x, y ) :
=
plim m →∞ ⟨ f σ ( x ; W ) , f σ ( y ; W ) ⟩ , known as the Neural Network Gaussian Process (NNGP) kernel.
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 ∈ R d , let K (0) σ ( x, y ) :
= ⟨ x, y ⟩ and for every layer
h
= 1 , . . .
, L , recursively define kernel functions K ( h ) σ , .
K ( h ) σ : R d × R d → R as: K ( h ) σ ( x, y ) :
= E ( u,v ) ∼N (
[Equation]
where the covariance matrix is Λ ( h ) σ :=
[Equation]
The depth- L NNGP kernel is K ( L ) σ ( x, y ) and the depth- L NTK Θ ( L ) σ can be recursively computed
as Θ (0) σ ( x, y ) := ⟨ x, y ⟩ and
[Equation]
At the core of the expression for Θ ( L ) σ , 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 Λ ( ℓ ) σ are equal to one, was previously studied in [46].
We extend this to encompass general symmetric covariance matrices in the following definition.
Definition 1 (Dual Activation and Dual Kernel) .
For a smooth σ : R → R , we define the Dual
Kernel of σ as K σ : R d × R d → R defined as
[Equation]
Equation (5) only depends on bivariate Gaussian random variables ⟨ w, x ⟩ , ⟨ w, y ⟩ where E [ ⟨ w, x ⟩ 2 ] =
∥ x ∥ 22 , E [ ⟨ w, y ⟩ 2 ] = ∥
y ∥ 22 and E [ ⟨ w, x ⟩·⟨ w, y ⟩ ] = ⟨ x, y ⟩ .
Hence one can look at the dual kernel from a different perspective by choosing a proper covariance matrix.
To this end, let Λ a,b,c :=
(cid:104) a 2 abc abc b 2 (cid:105) for every a, b ∈ R + and c ∈ [ − 1 , 1] and the Dual Activation of σ with respect to Λ a,b,c is the function k σ : R + × R + × [ − 1 , 1] → R defined as k σ ( a, b, c ) :
= E ( u,v ) ∼N (0 , Λ a,b,c ) [ σ ( u ) σ ( v )] .
With these definitions in place, the following relationship between dual kernel and activation holds
[Equation]
Observe that K σ ( x, y ) corresponds to the NNGP kernel of a 1 -hidden layer neural network with activation σ .
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 to convolutional neural networks called CNN-GP [7, 8] and CNTK [9].
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 a 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 σ : R → R and any x, y ∈ S
d − 1 , the dual kernel is equal to K σ ( x, y )
= (cid:80) ∞
j =
c 2 j j ! · ⟨ x, y ⟩ j . where { c j } j ≥ 0 are coefficients of Hermite expansion of σ .
We now proceed to generalize this result from S d − 1 to entire R d \ { 0 } .
First we remark that it can be naturally extended to the dual kernel of q -homogeneous activation functions, i.e., σ ( at ) = | a | q σ ( t ) for every a, t ∈ R , on the entire R d \ { 0 } .
For every x, y ∈ R d \ { 0 } , the corresponding dual kernel is
[Equation]
As examples, (leaky) ReLU and rectified polynomials fall into this activation class.
Now suppose that σ is not homogeneous.
In particular, we first consider a polynomial activation σ ( t ) =
(cid:80) qj
=
a j t j with coefficients { a j } qj =0 .
Recall that K σ ( x, y ) can be obtained by taking the expectation of σ ( ⟨ w, x ⟩ ) σ ( ⟨ w, y ⟩ ) over w ∼ N (0 , I d ) for every x, y ∈ R d \ { 0 } .
To make use of Daniely et al.
[46]’s result, we factorize the input into its radial and angular part and rewrite the activation by expressing monomials in the Hermite polynomial basis.
Formally, let us write monomials in the Hermite basis as t i = (cid:80) iℓ =
µ i,ℓ h ℓ ( t ) for some coefficients { µ j,i } ji =0 .
Then
[Equation]
Then, we can derive the dual kernel of polynomial activation.
We further relax a condition on the activation and propose the result below.
Theorem 1.
For a polynomial (cid:101) σ ( t ) = (cid:80) qj =0 a j t j , the dual kernel of (cid:101) σ ( · ) , as per Definition 1, is
[Equation]
where r (cid:101) σ,ℓ ( t ) :
= (cid:80) ⌊ q − ℓ 2 ⌋ i = a ℓ +2 i ( ℓ +2 i )!
2 i · i ! ·√ ℓ !
t 2 i + ℓ .
Moreover, if an activation function σ : R → R satisfies ∥ σ ∥ 2 N (0 ,ν 2 ) < ∞ and ∥ σ − (cid:101) σ ∥ 2 N (0 ,ν 2 ) ≤ ε for some ε > 0 and ν ≥ 1 , then for every x, y ∈ R d such that ∥ x ∥ 2 , ∥ y ∥ 2 ∈ (0 , ν ] the following holds
[Equation]
The proof of Theorem 1 is provided in Appendix B.2.
For non-polynomial activations, one can consider approximating σ with its Hermite or Taylor expansion and then apply Theorem 1.
Examples can be found in Appendix B.2.
For activation functions that do not have a 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 i = 0 , . . .
, k , σ ( i ) is absolutely continuous and lim t →±∞ e − t 2 4 σ ( i ) ( νt )
=
and moreover ∥ σ ∥ 2 N (0 ,ν 2 ) < ∞ and (cid:13)(cid:13) σ ( k ) (cid:13)(cid:13) 2 N (0 ,ν 2 ) < ∞ .
Consider the Hermite expansion coefficients { c j } j ≥ 0 of function σ ( νt ) and denote (cid:101) σ ( t ) :
= (cid:80) qj =
c j h j ( t/ν ) .
Given x, y ∈ R d with ∥ x ∥ 2 , ∥ y ∥ 2 ∈ (0 , ν ] , | K σ ( x, y ) − K (cid:101) σ ( x,
[Equation]
| ≤
[Equation] [Equation]
[Equation]
where K σ ( · , · ) and K (cid:101) σ ( · , · ) are dual kernels corresponding to σ ( · ) and (cid:101) σ ( · ) in Definition 1, respectively.
Moreover, for the ReLU activation σ ( t ) = max( t, 0) , it holds that
[Equation]
[Equation]
The proof of Theorem 2 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 ( 1 √ kq k − 1 ) 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 σ and σ ′ or their polynomial approximations are calculated, one can compute (approximate) NNGP and NTK using Theorem 1 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 a normalized
Gaussian kernel defined as
[Equation]
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 K G is non-trivial.
From the dual kernel perspective, the activation should be 1 -homogeneous
and its Hermite series expansion is of form (cid:80) ∞
j =0 ± 1 j !
h j ( 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 σ : R → R which satisfies | σ ( t ) | ≤
C 1 exp (cid:16)
[Equation]
, | σ ′
[Equation]
t ) | ≤ C 2 exp
[Equation]
t 2 4 .
[Equation]
2 (cid:17) , ∥ σ ∥ 2 N (0 ,ν 2 ) < ∞ and ∥ σ ′′ ∥ 2 N (0 ,ν 2 ) < ∞ for some ν ≥ 1 and constants C 1 , C 2 , the following holds for any x, y ∈ R d with ∥ x ∥ 2 , ∥ y ∥ 2 ∈ (0 , ν ] and |⟨ x, y ⟩|
[Equation]
x ∥
[Equation]
∥ y ∥ 2 :
[Equation]
Additionally, if ∂∂c k σ ( · , · , c ) is continuous at c = ± 1 then Equation (14) holds for x, y such that |⟨ x, y ⟩| =
∥ x ∥ 2 ∥ y ∥ 2 .
The proof of 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 ∥ x ∥ 2 = ∥
y ∥ 2 .
Applying Theorem 3 to Equation (13) provides that .
K G ( x,
[Equation]
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 et al.
[48], k σ ′ was not derived explicitly.
Theorem 3 provides a simple way to compute k σ ′ (given in Equation (126)) via automatic differentiation, without requiring to take the expectation under multivariate Gaussian distribution or computing derivatives by hand.
This is implemented in stax_extensions.
Elementwise 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 2 d Gaussian distribution as numerical integration.
This can be efficiently done by Gauss-Hermite quadrature
[Equation]
where ( x i , w i ) , correspond to i -th root of degree q Hermite polynomial
h
i ( x ) and associated weights [55] w i =
[Equation]
See Appendix E for the derivation of the quadrature formula.
For smooth activation functions errors will quickly go down as q increases by Theorem 2.
We use this method to compute approximate (non-sketched) kernels for general activation functions in Figureand Figure 4.
It is implemented as stax_extensions.
ElementwiseNumerical in our code.
Approximating Neural Kernels via Sketching
Although using our Theorem 1, Theorem 2, 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 2 ( d + Lq 2 )) runtime and Ω( n 2 ) memory for datasets with n points in R d .
This becomes even more expensive for CNTK, where its runtime can be Ω(( nd 1 d 2 ) 2 ( c + Lq 2 )) 2 for n of images with size d 1 × d 2 × c .
To avoid quadratic complexities, we adopt a fast and efficient feature map construction via randomized sketching [41] for both NTK and NNGP, i.e.,
[Equation]
The previous approach was only applicable for the ReLU activation but we establish more general scheme based on our new results for dual kernel approximation.
Subspace embedding for homogeneous dual kernels.
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.
More specifically, we call a dual kernel K σ homogeneous if there exists a positive definite dot-product kernel function κ : [ − 1 , 1] →
[ − 1 , 1] such that,
[Equation]
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) is equal to the following for any positive integer L ,
[Equation]
where κ ◦ L denoted the L -fold composition of function κ .
Furthermore, if κ has a derivative κ ′ : [ − 1 , 1] → [ − 1 , 1] , using Theorem 3, there exists a depth- L NTK for this dual kernel, equal to
[Equation]
where we use the convention that κ ◦ 0 ( t )
=
t .
Therefore, if κ ( · ) can be tightly approximated by a low-degree polynomial, then the NNGP and NTK functions can also be tightly approximated by lowdegree polynomials.
Thus, by applying P OLY S KETCH , which is a 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 P OLY S KETCH see Appendix A.
We provide the details of this procedure in Algorithm 1 and prove the correctness and runtime of our procedure in Theorem 4.
Theorem 4 (Homogeneous NTK Embedding) .
Suppose that the dual kernel K σ is homogeneous as per Equation (17) .
Also suppose (cid:101) κ ( t ) is a degree- q polynomial with non-negative coefficients that satisfies ( 1 ) max t ∈ [ − 1 , 1] | (cid:101) κ ( t ) − κ ( t ) | ≤ 1 poly( n ) and max t ∈ [ − 1 , 1] | (cid:101) κ ′ ( t ) − κ ′ ( t ) | ≤ 1 poly( n ) , ( 2 ) max | t |≤ 1+ 1 poly( n ) | (cid:101) κ
( t + γ ) − (cid:101) κ ( t ) | ≤ 1 poly( n ) and max | t |≤ 1+ 1 poly( n ) | (cid:101) κ ′ ( t + γ ) − (cid:101) κ ′ ( t )
| ≤ 1 poly( n ) for any |
γ | ≤ 1 poly( n ) .
Then for any integer L ≥ 1 , any ε, λ ≥ 1 poly( n ) , and any dataset X ∈ R d × n with ∥ X ∥ F ≤ poly ( n ) , if K ntk ∈ R n × n is the depth- L NTK kernel matrix on this dataset, there exists m =
O (cid:16) s λ ( K ntk )
ε· poly (cid:0) q L , log n (cid:1)(cid:17)
such that the output ψ ( L ) ( X ) ∈ R m × n of Algorithmsatisfies with probability at least 1 − 1 poly
[Equation]
n )
[Equation]
Moreover, the runtime of Algorithm 1 is O (cid:0) poly (cid:0) q L , log n (cid:1) · ε − 2 · ( s λ ( K ntk ) · n + nnz( X )) (cid:1) .
We prove this theorem in Appendix C. As an example, let us apply Theorem 4 on the normalized Gaussian kernel K G defined in Equation (13), which is homogeneous.
The dot-product factor corresponding to this dual kernel is κ ( t ) =
exp( t − 1) .
The truncated Taylor series of this
[Figure]
function is (cid:101) κ ( t ) =
(cid:80) qj =
t j e · j ! .
If q = Ω(log n ) then it can be verified that the polynomial (cid:101) κ ( t ) satisfies the preconditions of Theorem 4.
Therefore, one can invoke Algorithm 1 to get a subspace embedding for the NTK kernel matrix corresponding to the normalized Gaussian dual kernel
K G in O (cid:16) ε − 2 · ( s λ ( K ntk ) · n + nnz( X )) · poly (cid:16) log L n (cid:17)(cid:17) time and with a target dimension of m =
O (cid:16) ε − 2 · s λ ( K ntk ) · poly (cid:16) log L n (cid:17)(cid:17) .
For any constant number of layers, L , this runtime and target dimension is is optimal up to poly (log n ) factors.
The implementation of our sketching algorithm is available at github.com/insuhan/ntk_activations .
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 matrix.
We use ReLU , Abs (i.e., σ ( t ) =
| t | ), sin , Gaussian, erf and GeLU activations and approximate them by their Hermite expansion where degree changes from q = 1 to 20 .
We randomly generate n = 1 , 000 of 256 -dimensional inputs where each entry is i.i.d.
drawn from
N (0 , 1 / √ 256) .
We also compare our approach to the Monte Carlo estimation of dual kernel, i.e., K σ ( x, y )
≈ 1 m (cid:80) mi =1 σ ( ⟨ w i , x ⟩ ) σ ( ⟨ w i , y ⟩ ) where { w i } mi =1 are i.i.d.
standard Gaussian vectors.
In Figure 1, we plot relative errors of the Frobenius norm of 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 outperforms the Monte Carlo method for all activations we used.
In particular, sin and Gaussian are well approximated because they are smooth and norms of their derivatives are bounded with respect to the normal measure.
[Figure]
We also benchmark the proposed CNTK approximating via sketching algorithm.
We perform CIFAR-10 classification [57] by solving the ridge regression problem.
The image classes are converted into 10 -dimensional one-hot vectors and inputs are pre-processed with regularized ZCA [54, 58].
We report the best test accuracy among 20 choices of ridge parameters in { 10 − 10+ 1219 i | i = 0 , 1 , . .
We extract
CNTK features of a 5 -layer convolutional neural network (known as Myrtle5 [54]) without pooling by setting degree q = 8 and explore feature dimension m = { 2 9 , . . .
, 2 14 } and homogeneous dual kernels including ReLU, ABReLU activations as well as deep normalized Gaussian kernels with 2 scaling factors.
See Appendix G for more details.
In 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 is because the coefficients decay of the normalized Gaussian is faster than that of the ReLU, which leads to a lower approximation error of sketching algorithm.
We also perform comparison among different activation functions in neural kernels in Appendix E.
Speedup.
We observe that the exact CNTK of Myrtle-5 constructs a kernel matrix of size 60 , 000 × 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 106 × speedup.
If we use less training data to construct 20 , 000 × 20 , 000 kernel matrix, the accuracy is about 77% accuracy and the runtime is 16.8 GPU hours in which our approximation is still 12 × faster without loss of accuracy.
We believe such acceleration through our methods open the door to using neural kernels in a wide range of research domains.
Discussion
In this work, we introduced methods to efficiently compute neural kernels for general activations.
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 a 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 Zandieh 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 (N0001419-1-2406), and the AI Institute for Learning-Enabled Optimization at Scale (TILOS).
We thank
Timothy Nguyen and Jeffrey Pennington for discussions and feedback on the project.
A Sketching Preliminaries
The P OLY S KETCH algorithm is a norm-preserving dimensionality reduction that can be applied to the tensor product of multiple vectors very quickly [56], i.e., for any v 1 , . .
. , v q ∈ R d , there exists a randomized mapping Q q : R d q → R m which satisfies that
[Equation]
with high probability and Q q ( v 1 ⊗ · · · ⊗ v p ) can be computed very fast.
Here, sketching dimension m is a trade-off parameter between runtime and accuracy.
Algorithm 2 describes the pseudo-code of
P OLY S KETCH and Theorem 5 summarizes Theorems 1.2 and 1.3 of [56] which guarantees spectral approximation of output of P OLY S KETCH .
, d } m 2: let y ← [ x 1 s 1 , x 2 s 2 , . . .
, x d s d ] 3: compute z ← FFT( y ) 4: return 1 √ m
[ z b 1 , . .
z b m ]
Theorem 5 (P OLY S KETCH ) .
For every integers p, d ≥ 1 and every ε, δ > 0 , there exists a distribution on random matrices
Q p ∈ R m × d p , called degree p P OLY S KETCH such that (1) for some m =
O (cid:0) pε 2 log 3 1 εδ (cid:1) and any y ∈ R d p , Pr (cid:2)
∥
Q p y ∥
22 ∈ (1 ± ε ) ∥
y ∥ 22 (cid:3) ≥ 1 − δ ; (2) for any x ∈ R d , the total time to compute Q p x ⊗ p is O (cid:16) pm log m + p 3 / 2 ε log 1 δ nnz( x ) (cid:17) ; (3) for any collection of vectors v 1 , . .
. , v p ∈ R d , the time to compute Q p ( v 1 ⊗ · · · ⊗ v p ) is bounded by O (cid:16) pm log m + p 3 / 2 ε d log 1 δ (cid:17) ; (4) for any λ > 0 and any matrix A ∈ R d p × n , where the statistical dimension of A ⊤ A is s λ , there exists some m =
O (cid:16) p 4 s λ ε 2 log
3 nεδ (cid:17) such that,
[Equation]
B Proofs B.1 Properties of Hermite Polynomials
We first introduce that Hermite polynomials can be used as the random feature of monomial kernels for inputs on the unit sphere, which will be used in our analysis.
Proposition 1.
For x, y ∈ S d − 1 , it holds that
[Equation]
Proof of Proposition 1: Let a := ⟨ w, x ⟩ , b := ⟨ w, y ⟩ then E w [ a ] =
E w [ b ] =
and Cov( a, b ) =
E [ ab ] = ⟨ x, y ⟩ .
Hence, we have that
[Equation]
where
[Equation] [Equation]
We introduce Proposition 11.31 in O’Donnell [60]:
[Equation]
This completes the proof of Proposition 1.
B.2 Proof of Theorem
Theorem 1.
For a polynomial (cid:101) σ ( t ) = (cid:80) qj =0 a j t j , the dual kernel of (cid:101) σ ( · ) , as per Definition 1, is
[Equation]
where r (cid:101) σ,ℓ ( t ) :
= (cid:80) ⌊ q − ℓ 2 ⌋ i = a ℓ +2 i ( ℓ +2 i )!
2 i · i ! ·√ ℓ !
t 2 i + ℓ .
Moreover, if an activation function σ : R → R satisfies ∥ σ ∥ 2 N (0 ,ν 2 ) < ∞ and ∥ σ − (cid:101) σ ∥ 2 N (0 ,ν 2 ) ≤ ε for some ε > 0 and ν ≥ 1 , then for every x, y ∈ R d such that ∥ x ∥ 2 , ∥ y ∥ 2 ∈ (0 , ν ] the following holds
[Equation]
Proof of Theorem 1: Due to homogeneity of the inner-product
[Equation]
we can write
[Equation]
Note that monomial t i of degree i ≥ 0 can be explicitly written in the Hermite basis as t i
= (cid:80) iℓ =
µ i,ℓ h ℓ ( t ) where
[Equation]
Plugging this into Equation (26) and re-arranging terms
[Equation]
we obtain that
[Equation]
Applying Equation (28) to the definition of the dual kernel K (cid:101) σ ( x
[Equation]
y ) given in Equation (5) and taking the expectation over w gives
[Equation]
a j µ j,m ∥ y ∥
j
[Equation]

 h ℓ (cid:18) ⟨ w, x ⟩ ∥ x ∥
[Equation]
(cid:19) h m (cid:18)
⟨ w
[Equation]
y ⟩ ∥
y
∥(cid:19) =
[Equation]
(cid:88) ℓ
[Equation] [Equation]
(cid:88)
[Equation]
a
[Equation]
(cid:20) h
[Equation]
(cid:18)
[Equation] [Equation]
x ∥
[Equation]
(cid:19) h m (cid:18)
⟨ w
[Equation]
y ⟩ ∥
y ∥
[Equation]
(cid:19)(cid:21) =
[Equation]
i
[Equation]
(cid:33)
[Equation]
ℓ !
(cid:18)
[Equation]
x
[Equation] [Equation]
where Equation (29) follows from Proposition 1.
Now using Equation (27)
[Equation]
we have
[Equation]
Therefore
[Equation]
we obtain that
[Equation]
We finish off the proof of Theorem 1 by bounding the error | K σ ( x, y ) − K (cid:101) σ ( x, y ) | .
We use Equation (5) along with the assumption that both x, y ̸ = 0 to write,
[Equation]
K
[Equation]
where the inequality above follows from the triangle inequality
[Equation]
Now we bound each of Equation (30) and Equation (31) separately.
First let us bound Equation (30) using Cauchy–Schwarz inequality as follows,
[Equation]
where the first line follows from the Cauchy–Schwarz inequality and the third line above follows from the assumption that ∥ x ∥ 2 , ∥ y ∥ 2 ̸ = 0 and the last line follows from the precondition of Theorem 1.
Similarly, we can bound Equation (31), as follows,
[Equation]
where the fourth line above follows from the assumption that ∥ x ∥ 2 , ∥ y ∥ 2 ̸ = 0 , the fifth line above follows from the AM-GM inequality along with the the preconditions of Theorem 1, and the last equality above follows from the preconditions of Theorem 1.
Now by plugging Equation (32) and Equation (33) back into Equation (30) and Equation (31) we find that,
[Equation]
This completes the proof of Theorem 1.
Examples for Taylor expansion.
Observe that σ ( t ) =
sin( t ) is analytic and has a Taylor expansion with coefficients
a
[Equation]
+
[Equation]
=
[Equation] [Equation]
· 1 ( ℓ is odd ) .
By invoking Theorem 1 we have,
[Equation]
Therefore,
[Equation]
Similarly, we can derive that K cos ( x, y ) =
e − ∥ x ∥ 22+ ∥ y ∥ 22cosh( ⟨ x, y ⟩ ) that corresponds to Table 1.
B.3 Proof of Theorem
Theorem 2.
Given σ : R → R , suppose that there exists an integer k ≥ 2 and some ν ≥ 1 such that for every i = 0 , . . .
, k , σ ( i ) is absolutely continuous and lim t →±∞ e − t 2 4 σ ( i ) ( νt )
=
and moreover ∥ σ ∥ 2 N (0 ,ν 2 ) < ∞ and (cid:13)(cid:13) σ ( k ) (cid:13)(cid:13) 2 N (0 ,ν 2 ) < ∞ .
Consider the Hermite expansion coefficients { c j } j ≥ 0 of function σ ( νt ) and denote (cid:101) σ ( t ) :
= (cid:80) qj =
c j h j ( t/ν ) .
Given x, y ∈ R d with ∥ x ∥ 2 , ∥ y ∥ 2 ∈ (0 , ν ] , |
[Equation]
where K σ ( · , · ) and K (cid:101) σ ( · , · ) are dual kernels corresponding to σ ( · ) and (cid:101) σ ( · ) in Definition 1, respectively.
Moreover, for the ReLU activation σ ( t ) = max( t, 0) , it holds that
[Equation]
In order to prove this theorem we first need to establish the following bound on the decay rate of the
Hermite expansion coefficients of smooth functions,
Lemma 1.
Suppose that there exists an integer k ≥
such that for every i = 0 , . . . , k , σ ( i ) ( t ) are absolutely continuous in R and
lim t →±∞ e − t 2 4 σ ( i ) ( t ) = 0 .
Assume that ∥ σ ∥ 2 N < ∞ and (cid:13)(cid:13) σ ( k ) (cid:13)(cid:13) 2 N < ∞ .
Let { c j } ∞ j =
be the Hermite expansion coefficients of this function such that
[Equation]
∞
[Equation]
=
c j h j (cid:13)(cid:13)(cid:13) N = 0 .
Then, for any integer j ≥ k :
[Equation]
The proof of Lemma 1 is provided in Appendix B.4.
Proof of Theorem 2: First, because of the precondition of Theorem 2 about ∥ σ ∥ 2 N (0 ,ν 2 ) =
E t ∼N (0 ,ν 2 ) (cid:104) | σ ( t ) | 2 (cid:105) =
E t ∼N (0 , 1) (cid:104) | σ ( νt ) | 2 (cid:105) < ∞ , the function σ ( νt ) is an L 2 function with respect to the normal measure N (0 , 1) on the real line.
Therefore, because the Hermite polynomials { h j } ∞ j =0 provide an orthogonal basis for L 2 function with respect to normal measure N (0 , 1) , σ ( νt ) converges to its Hermite expansion, i.e., E t ∼N
[Equation]
= 0 .
We obtain an error bound on the dual kernel by invoking Theorem 1.
To do so, we need to first upper bound
E t ∼N
[Equation] [Equation]
| σ ( t )
[Equation]
) | 2 (cid:105) , as follows
[Equation]
where the second line above follows from the fact that h j ’s are orthogonal with respect to the normal measure N (0 , 1) .
The third line follows from the fact that ∥
h j ∥ 2 N = j ! .
We now proceed to upper bound the term in Equation (37), using the bound on the Hermite expansion coefficients we proved in Lemma 1.
We apply this lemma to the function σ ( νt ) whose Hermite expansion coefficients are { c i } ∞ i =0 .
By precondition of Theorem 2 we have (cid:13)(cid:13) σ
[Equation]
k )
[Equation]
N = E t ∼N (0 ,
[Equation]
) (cid:104)(cid:12)(cid:12)
[Equation]
[Equation] [Equation]
t ) (cid:12)(cid:12) 2 (cid:105) <
∞ .
This implies that,
[Equation]
Furthermore, the precondition of Theorem 2 about lim t →±∞ e − t 2 4 σ ( i ) ( νt ) = 0 implies the following,
[Equation]
Therefore, the preconditions of Lemma 1 are satisfied and by invoking this lemma we have the following inequality for any integer j ≥ k ,
[Equation]
Plugging the above inequality into Equation (36) into Equation (37)
[Equation]
gives
[Equation]
Thus we can now invoke Theorem 1 with
[Table]
=
[Table]
to find that [Equation] [Table] [Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Now we prove the second statement of the theorem about the ReLU activation σ ( t )
=
max( t, 0) .
It is easy to check that for this function
[Equation]
Furthermore for any j ≥ 0 , the Hermite coefficients of σ ( νt ) are
[Equation]
Using integration-by-parts and the fact that h ′ j ( t ) =
jh j − 1 ( t ) for all j ≥ 1 , we get that
[Equation]
∞
[Equation]
Therefore
[Equation]
E t
[Equation]
By invoking Theorem 1, using Equation (39) and Equation (40), we have
[Equation]
This completes the proof of Theorem
[Equation]
B.4 Proof of Lemma
Lemma 1.
Suppose that there exists an integer k ≥
such that for every i = 0 , . . . , k , σ ( i ) ( t ) are absolutely continuous in R and
lim t →±∞ e − t 2 4 σ ( i ) ( t ) = 0 .
Assume that ∥ σ ∥ 2 N < ∞ and (cid:13)(cid:13) σ ( k ) (cid:13)(cid:13) 2 N < ∞ .
Let { c j } ∞ j =
be the Hermite expansion coefficients of this function such that
[Equation]
N = 0 .
Then, for any integer j ≥ k :
[Equation]
Proof of Lemma 1: The proof can be obtained by slightly modifying Theorem 3.1 in [61].
The precondition ∥
[Equation]
< ∞ implies that σ is an L 2 -function with respect to measure e − t 2 2 on real line.
Because Hermite polynomials { h j } ∞ j =0 form an orthogonal basis for the
Hilbert space of L 2 -functions with respect to normal measure N (0 , 1) , σ ( t ) converges to its Hermite expansion, i.e., (cid:80) ∞
j =
c j h j ( t ) .
The j -th coefficient in this expansion is
[Equation]
[Equation]
Using the Rodrigues’ expression of Hermite polynomials in Equation (1) and integration-by-parts, we have, (cid:90) ∞
[Equation]
where the last line above follows from the Rodrigues’ expression of degree j − 1 Hermite polynomial in Equation (1).
Therefore, using 22.14.17 in [55] 3 , the first term in Equation (42) is 0 and by 3 Equation 22.
applying the above repeatedly we have
[Equation]
where the second last inequality comes from Cauchy-Schwarz inequality and the last one holds from that E t ∼N (0 , 1) (cid:104) | h ℓ ( t ) | 2 (cid:105) =
ℓ ! and the assumption.
This completes the proof of Lemma
[Equation]
B.5 Proof of Theorem
Theorem 3.
Given a differentiable activation function σ : R → R which satisfies | σ ( t ) | ≤
[Equation]
N (0 ,ν 2 ) < ∞ and ∥ σ ′′ ∥ 2 N (0 ,ν 2 ) < ∞ for some ν ≥ 1 and constants C 1 , C 2 , the following holds for any x, y ∈ R d with ∥ x ∥ 2 , ∥ y ∥ 2 ∈ (0 , ν ] and |⟨ x, y ⟩| < ∥ x ∥ 2 ∥ y ∥ 2 :
[Equation]
Additionally, if ∂∂c k σ ( · , · , c ) is continuous at c = ± 1 then Equation (14) holds for x, y such that |
[Equation]
y
[Equation]
=
[Equation]
x ∥ 2 ∥ y ∥ 2 .
Note that our assumption on σ can be weakened to be: there exists ϵ > 0 and constants C 1 and Csuch that
[Equation]
Proof of Theorem 3: Recall that the dual activation is defined as
[Equation]
where for a, b ∈ R ≥ 0 and c ∈ [ − 1 , 1]
[Equation]
Using a whitening transformation, we introduce the standard i.i.d.
Gaussian random variables w 1 , w 2 ∼ N (0 , 1) that satisfy
[Equation]
(cid:21) .
Thus, by denoting w = (cid:20) ww
[Equation]
(cid:21) , the dual activation can be written as
[Equation]
Using Equation (46), we can calculate ∂∂c k σ ( · , · , c ) if the derivative can be interchangeable with the expectation.
To this end, we use the “measure theory” statement of Leibniz integral rule .
Lemma 2 (Measure theory statement of Leibniz integral rule, Theorem 6.28 of [62]) .
Let µ be a probability distribution with support Ω , let I ⊂ R be a nontrivial open interval, also let f : Ω
× I → R be a map with the following properties:
For any x ∈ I , E w ∼ µ [ | f ( w, x ) | ] < ∞ .
For almost all w ∈ Ω , the map x → f ( w, x ) is differentiable with derivative ∂∂x f ( w, x ) .
There is a map h : Ω → R such that E w ∼ µ [ | h ( w ) | ] < ∞ and (cid:12)(cid:12) ∂∂x f ( · , x ) (cid:12)(cid:12) ≤ h .
Then, for any x ∈ I , E w ∼ µ (cid:2)(cid:12)(cid:12) ∂∂x f ( w, x ) (cid:12)(cid:12)(cid:3) < ∞ and the function F : x → E w ∼ µ [ f ( w, x )] is differentiable with derivative F ′ ( x ) =
E w ∼ µ (cid:2) ∂∂x f
( w, x ) (cid:3) .
To invoke Lemma 2 on the above expression of the dual kernel, we let I be the open interval ( − 1+ ε, 1 − ε ) for an arbitrarily small ε > 0 and f : R 2 × I
→ R be the function defined as f ( w, c ) :
= σ ( aw 1 ) · σ ( bcw 1 + b √ 1 − c 2 w 2 ) for some fixed values of a, b .
With these notations in place, we proceed to check if the preconditions of Lemma 2 are satisfied.
To verify the first precondition, we need to show that for any c ∈ I , E w ∼N (0 , I 2 ) (cid:2)(cid:12)(cid:12) σ ( aw 1 ) · σ ( bcw 1 + b √ 1 − c 2 w 2 ) (cid:12)(cid:12)(cid:3) < ∞ .
We find that,
E w
[Equation]
where the first line above follows from Cauchy–Schwarz inequality.
The second line above follows from the fact that w 1 and w 2 are independent copies of the normal random variable N (0 , 1) , thus the random variable bcw 1 + b √ 1 − c 2 w 2 is indeed bγ for a normal γ ∼ N (0 , 1) .
Now using the preconditions of Theorem 3, for any a, b ∈ (0 , ν ] we have
[Equation]
Also in case a = 0 or b =
we have E w 1 ∼N (0 , 1) (cid:104) | σ ( aw 1 ) | 2 (cid:105) =
| σ (0) | 2 < ∞ , therefore, above inequalities along with Equation (47) proves the first precondition of Lemma 2 for any a, b ∈ [0 , ν ] .
To verify that the second precondition of Lemma 2 holds, we show that for almost all w 1 , w 2 ∈ R the map c → σ ( aw 1 ) · σ ( bcw 1 + b √ 1 − c 2 w 2 ) is differentiable.
This holds true because of the assumption of Theorem 3 on about the activation σ ( · ) being differentiable.
The derivative of this map is in fact
[Equation]
Finally
[Equation]
we check the third precondition of Lemma 2.
Since | c | <
[Equation]
there is an ε > 0 such that c ∈ ( − 1 + ε, 1 − ε ) .
We have,
[Equation]
where the second inequality follows from the preconditions of Theorem 3 about the upper bounds on σ ( · ) and σ ′ ( · ) , the third one follows from ( cw 1 + √ 1 − c 2 w 2 ) 2 ≤ w 21 + w 22 , and the fourth one follows from a, b ≤ ν .
Now it is easy to check that this upper bound function satisfies E w ∼N (0 , I 2 ) [ | h ( w ) | ] < ∞ .
Therefore, we can invoke Lemma 2 to calculate the derivative of the dual kernel k σ ( a, b, c ) with respect to c as follows,
[Equation]
Next we compute Equations (50) and (51) by using Stein’s lemma,
Lemma 3 (Stein’s Lemma) .
For a differentiable function ϕ : R → R with E x ∼N (0 , 1) [ | ϕ ′ ( x ) | ] < ∞ ,
[Equation] [Equation]
Applying Stein’s Lemma to Equation (50) gives,
[Equation]
Applying Stein’s Lemma to Equation (51) gives,
[Equation]
[Equation]
Here we show that the term b 2 c E w (cid:2) σ ( aw 1 ) · σ ′′ (cid:0) bcw 1 + b √ 1 − c 2 w 2 (cid:1)(cid:3) in Equations (52) and (53) has a bounded value as follows,
[Equation]
where the first line above follows from Cauchy–Schwarz inequality.
The second line above follows from the fact that w 1 and w 2 are independent copies of the normal random variable N (0 , 1) , thus the random variable bcw 1 + b √ 1 − c 2 w 2 is indeed bγ for a normal γ ∼ N (0 , 1) .
Therefore, in order for the expectation b 2 c E w (cid:2) σ ( aw 1 ) · σ ′′ (cid:0) bcw 1 + b √ 1 − c 2 w 2 (cid:1)(cid:3) to make sense, it is enough to have b 2 E w 1 (cid:104) | σ ( aw 1 ) | 2 (cid:105) < ∞ and b 2 E γ (cid:104) | σ ′′ ( bγ )
| 2 (cid:105) < ∞ .
Note that the dual activation is symmetric with respect to swapping a and b , in the sense that k σ ( a, b, c ) =
k σ ( b, a, c ) .
Thus, we can without loss of generality assume that b ≤ a .
With this assumption b 2 E w 1 (cid:104) | σ ( aw 1 ) | 2 (cid:105) ≤ a 2 E w 1 (cid:104) | σ ( aw 1 ) | 2 (cid:105) .
Now, by recalling Equation (48), we have a 2 E w 1 (cid:104) | σ ( aw 1 ) | 2 (cid:105) ≤ aν · ∥ σ ∥ 2 N (0 ,ν 2 ) < ∞ and b 2 E γ (cid:104) | σ ′′ ( bγ )
| 2 (cid:105) ≤ bν · ∥ σ
′′ ∥ 2 N (0 ,ν 2 ) < ∞ .
Summing Equations (52) and (53) and dividing the sum by ab give that
[Equation]
Finally, plugging in the values a = ∥ x ∥ 2 , b = ∥
y ∥ 2 and c = ⟨ x,y ⟩ ∥ x
∥ 2 ∥ y ∥ 2 such that a, b ≤ ν and | c | <and using Equation (6) result in Equation (14).
Now suppose that the map c → ∂∂c k σ ( · , · , c ) is continuous at c = ± 1 .
Since we consider that σ ′′ ( · ) exists, σ ′ ( · ) is continuous almost everywhere.
Using these properties, we claim that the right-hand side in Equation (54) is continuous in c because for every c ′ ∈ [ − 1 , 1] it holds that
[Equation]
The above equality holds from the dominated convergence theorem (see Corollary 6.26 in [62]) with the dominated function obtained as
[Equation]
where the first inequality follows from the preconditions of Theorem 3 and the second one follows from ( cw 1 + √ 1 − c 2 w 2 ) 2 ≤ w 21 + w 22 , and the third one follows from a, b ≤ ν .
And it is easy to check that
E w ∼N (0 , I 2 ) [ | h ′ ( w ) | ] < ∞ .
Hence both sides of Equation (54) are continuous at c = ±and taking lim c →± 1 in both sides of Equation (54) gives that Equation (14) holds for x,
y such that |⟨ x, y ⟩| =
∥ x ∥ 2 ∥ y ∥ 2 .
This concludes the proof of Theorem 3.
Examples.
For σ ( t ) = sin( t ) , the corresponding dual kernel is known to be
[Equation]
[Equation] [Equation]
Applying Theorem 3 to k sin
[Equation]
which is equivalent to k cos ( a, b, c ) (see Table 2 for detailed derivations).
For σ ( t ) = erf( t ) , the corresponding dual kernel is known as
[Equation]
[Equation]
Again, applying Theorem 3 to k erf provides that
[Equation]
One can check that this matches the dual kernel of (erf( t )) ′
= 2 √ π e − t 2 from Table 2.
In addition, Theorem 3 holds for the ReLU activation because
[Equation]
which is equivalent to the dual kernel of ReLU ′ ( t )
=
Step( t ) .
This theorem is used in Elementwise in our codebase to automatically derive the NTK given only the NNGP function.
C Proof of Theorem
Theorem 4 (Homogeneous NTK Embedding) .
Suppose that the dual kernel K σ is homogeneous as per Equation (17) .
Also suppose (cid:101) κ ( t ) is a degree- q polynomial with non-negative coefficients that satisfies ( 1 ) max t ∈ [ − 1 , 1] | (cid:101) κ ( t ) − κ ( t ) | ≤ 1 poly( n ) and max t ∈ [ − 1 , 1] | (cid:101) κ ′ ( t ) − κ ′ ( t ) | ≤
[Equation]
poly( n
[Equation]
, ( 2 ) max | t |≤ 1+ 1 poly
[Equation]
n
[Equation]
κ
[Equation]
γ ) − (cid:101) κ
[Equation]
t )
[Equation]
≤ 1 poly( n ) and max | t |≤ 1+
[Equation]
|
[Equation]
for any | γ | ≤ 1 poly( n ) .
Then for any integer L ≥ 1 , any ε, λ ≥ 1 poly( n ) , and any dataset X ∈ R d × n with ∥ X ∥ F ≤ poly ( n ) , if K ntk ∈ R n × n is the depth- L NTK kernel matrix on this dataset, there exists m =
O (cid:16) s λ ( K ntk )
ε· poly (cid:0) q L , log n (cid:1)(cid:17)
such that the output ψ ( L ) ( X ) ∈ R m × n of Algorithmsatisfies with probability at least 1 −
[Equation]
Moreover, the runtime of Algorithm 1 is O (cid:0) poly (cid:0) q L , log n (cid:1) · ε − 2 · ( s λ ( K ntk ) · n + nnz( X )) (cid:1) .
Proof of Theorem 4: We start the proof by showing that the polynomial R ( L ) ( t ) defined as
[Equation]
tightly approximates the following function at every point t ∈ [ − 1 , 1]
[Equation]
Specifically, we prove that
[Equation]
In order to prove Equation (62), we first show that for every h = 0 , 1 , 2 , . . .
L the following holds
[Equation]
The proof of the above is by induction on h .
For h = 0 by convention κ ◦ h ( t ) =
(cid:101) κ ◦ h ( t ) =
t , which proves the base of induction.
For the inductive step suppose that max t ∈ [ − 1 , 1] (cid:12)(cid:12) κ ◦ h − 1 ( t ) − (cid:101) κ
[Equation]
h − 1 ( t ) (cid:12)(cid:12) ≤ 1 poly( n ) holds for some h ≥ 1 .
Using this inductive hypothesis along with preconditions of Theorem 4, for any t ∈ [ − 1 , 1] we can write,
[Equation]
where the first line above follows from triangle inequality.
The second line above follows from precondition ( 2 ) of the theorem.
The third line follows from precondition ( 1 ) of the theorem.
Therefore max t ∈ [ − 1 , 1] (cid:12)(cid:12) κ ◦ h ( t ) − (cid:101) κ ◦ h ( t ) (cid:12)(cid:12) ≤ 1 poly( n ) for any h = 0 , 1 , . . .
L .
Moreover, by preconditions of the theorem, we can show in a similar fashion that
[Equation]
These inequalities are sufficient to prove
Equation (62).
Now, let us define the kernel (cid:101) Θ ( L ) σ as
[Equation]
The depth- L NTK kernel, as we showed in Equation (19), is
[Equation]
Using Equation (62), for any x, y ∈ R d , we have,
[Equation]
For any dataset X = [ x 1 , x 2 , . . .
, x n ] ∈ R d × n , we let (cid:102) K ntk ∈ R n × n be the kernel matrix corresponding to the kernel function (cid:101) Θ ( L ) σ and X , i.e., [ (cid:102) K ntk ] i,j = (cid:101)
Θ ( L ) σ ( x i , x j ) to have that
[Equation]
where the third line above follows from the assumption of the theorem about ∥ X ∥ F ≤ poly ( n ) and ε, λ ≥ 1 poly( n ) .
Therefore, in order to prove the desired subspace embedding guarantee of Theorem 4, it suffices to prove that with probability at least 1 − 1 poly( n ) , the following holds
[Equation]
From now on we focus on proving the above inequality.
If we let R ( L ) ( t ) = (cid:80) pi =
c j t j be the polynomial defined in line 3 of Algorithm 1 then we have that
[Equation]
where
D = diag
[Equation]
∥
[Equation]
Note that each of the term (cid:0) Y ⊗ j D (cid:1) ⊤ Y ⊗ j D
= D ( Y ⊗ j ) ⊤ Y ⊗ j D is a positive definite Gram matrix.
Also, from the fact that coefficients c
j are positive and by Courant-Fischer’s min-max theorem, the statistical dimension of the Gram matrix c j · (cid:0) Y ⊗ j D (cid:1) ⊤ Y ⊗ j D for every j ≥ 0 is upper bounded by the statistical dimension of the kernel matrix (cid:102) K ntk .
More specifically, for any µ > 0 and every j = 0 , 1 , . . .
p , we have
[Equation]
Now let µ := λp +1 and note that from the definition of statistical dimension it follows that s µ
[Equation]
≤
[Equation]
p + 1) s λ
[Equation]
K ntk
[Equation]
The sketch matrix Q j defined in line 4 of the algorithm has m = Ω (cid:0) ε
[Equation]
s λ ( K ntk ) · poly (cid:0) q L , log n (cid:1)(cid:1)
=
Ω (cid:0) ε −
[Equation]
s µ ( K ntk ) · poly (cid:0) q L , log n (cid:1)(cid:1) rows.
Therefore, by Theorem 5, the following holds for every j = 0 , 1 , . . .
p with probability at least 1 − 1 poly( n ) , c
[Equation]
D
[Equation]
⊤ Y ⊗ j D +
[Equation]
By union bound over p +1 =
O ( q L ) =
o (poly ( n )) events, the above inequality holds simultaneously for all j with high probability in n .
Thus, by summing up the above inequality over all j and using the fact that µ = λp +1 we find that,
[Equation]
This proves the theorem because the output of the algorithm satisfies ψ ( L ) ( X ) ⊤ ψ ( L ) ( X )
=
(cid:80)
[Equation]
Q j Y ⊗ j (cid:1)
⊤ Q j Y ⊗ j D j .
The runtime bound follows immediately from Theorem 5.
D Convolutional Neural Tangent Kernel
In this section, we design and analyze an efficient oblivious sketch for the Convolutional Neural
Tangent Kernel (CNTK), which is the kernel function corresponding to a CNN with infinite number of channels.
Arora et al.
[9] gave dynamic programming (DP) based solutions for computing two variants of CNTK; one is the vanilla version which performs no pooling, and the other performs
Global Average Pooling (GAP) on its top layer.
For conciseness, we focus mainly on the CNTK with
GAP, which also exhibits superior empirical performance [9].
However, we remark that the vanilla
CNTK has a very similar structure and hence our techniques can be applied to it, as well.
We start by restating the DP approach proposed in [9] for computing the L -layered CNTK with an arbitrary activation function σ , convolutional filters of size q × q and GAP.
Consider two input images y, z ∈ R d 1 ×
d 2 × c where c is the number of channels ( c = 3 for the standard color image).
For every i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] , define
[Equation]
For every h ∈ [ L ] , every i, i ′ ∈ [ d 1 ] and j, j ′ ∈ [ d 2 ] , define
[Equation]
where the covariance matrix is
[Equation]
For
every h ∈ [ L
[Equation]
every i
i ′
∈ [ d
and
j j ′ ∈
[ d 2 ] , define
[Equation]
Let Π (0) ( x, y ) :
= 0 and for every h ∈
[ L − 1] , every i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] , define Π
[Equation]
i
[Equation]
and also Π ( L ) ( y, z ) :
= Π ( L − 1) ( y, z ) ⊙ .
Γ ( L
[Equation] [Equation]
The final CNTK expressions is defined as:
[Equation]
The above procedure for exact computation of the depth- L CNTK value Θ ( L ) cntk ( y, z ) takes Ω (cid:0) ( d 1 d 2 ) 2 ( c + L ) (cid:1) runtime, which is extremely slow particularly due to its quadratic dependence on the number of pixels of input images d 1 d 2 .
Fortunately, we are able to show that the CNTK for homogeneous dual kernels, as per Equation (17), is a highly structured object that can be fully characterized in terms of tensoring and composition of the dot-product factor of dual kernels, and exploiting this special structure is key in designing efficient sketching methods for the CNTK.
D.1 CNTK for Homogeneous Dual Kernels
In this section we show that the CNTK function corresponding to any homogeneous dual kernel, i.e.,
K
[Equation]
x,y
[Equation]
x
[Equation] [Equation]
y
[Equation]
2 (cid:17) for some κ : [ − 1 , 1] → [ − 1 , 1] , takes a simple form which enables us to devise efficient sketching algorithms for the CNTK.
Unlike the fully-connected NTK, the CNTK is not a simple dot-product kernel function.
The key reason being that CNTK works by partitioning its input images into patches and locally transforming the patches at each layer, as opposed to the NTK which operates on the entire input vectors.
The depth- L CNTK corresponding to homogeneous dual kernels can be fully characterized in terms of tensoring and composition of the dot-product kernel κ and its derivative κ ′ .
Definition 2 (CNTK for Homogeneous Dual Kernels) .
For every positive integers q, L , the L -layered
CNTK for a homogeneous dual kernel, as per Equation (17), and convolutional filter size of q × q is defined as follows
For x ∈ R d 1 × d 2 × c , every i ∈ [ d 1 ] and j ∈ [ d 2 ] let N (0) i,j ( x ) :
[Equation] [Equation]
x i
[Equation]
j,l | 2 , and for every h ≥ 1 , recursively define,
[Equation]
For every h ∈ [ h ] , every i, i ′ ∈ [ d 1 ] and j, j ′ ∈ [ d 2 ] , define Γ
[Equation]
) i,
[Equation]
,i ′ ,
[Equation]
where A =
[Equation]
For every h ∈ [ L ] , every i, i ′ ∈ [ d 1 ] and j, j ′ ∈ [ d 2 ] , define .
[Equation]
i
[Equation]
Let Π (0) ( y, z ) := 0 and for every h ∈
[ L − 1] , every i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] , define
Π (
[Equation]
i
[Equation] [Equation]
Furthermore, define
[Equation]
The final CNTK expressions for ReLU activation is
[Equation]
We now describes some of the basic properties of the functions Γ ( h ) ( y, z ) , .
Γ ( h ) ( y, z ) , and Π ( h ) ( y, z ) defined in Equation (70), in the following lemma,
Lemma 4 (Properties of Γ ( h ) ( y, z ) , ˙Γ ( h ) ( y, z ) , and Π ( h ) ( y, z ) ) .
Suppose that the dot-product kernel κ ( · ) in Equation (17) and its derivative satisfy κ (1) = κ ′ (1) = 1 .
For every images y, z ∈ R d 1 × d 2 × c , every integer h ≥ 0 and every i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] the following properties are satisfied by functions Γ ( h ) , .
Γ ( h ) , Π ( h ) and N ( h ) defined in Equation (70) , Equation (71) , Equation (72) and
Equation (73) , and Equation (69) of Definition 2:
Cauchy–Schwarz: (cid:12) (cid:12) (cid:12) Γ ( h ) i,j,i ′ ,j ′ ( y, z ) (cid:12) (cid:12) (cid:12) ≤
[Equation]
) i,j ( y ) ·
N ( h ) i ′ ,j ′ ( z
[Equation]
, and (cid:12) (cid:12) (cid:12) ˙
Γ ( h ) i,j,
[Equation] [Equation]
′ ( y, z ) (cid:12) (cid:12) (cid:12) ≤ 1 q 2 , and Π ( h ) i,j,i ′ ,j ′
( y, z ) ≤ (cid:113) Π ( h ) i,j,i,j ( y, y )
· Π ( h ) i ′ ,j ′ ,i ′ ,j ′ ( z, z ) .
Norm value: Γ ( h ) i,j,i,j ( y, y ) =
N ( h ) i,j ( y ) q≥ 0 , and ˙Γ ( h ) i,j,i,j ( y, y )
= 1 q 2 ≥ 0 , and Π ( h ) i,j,i,j ( y, y ) =
(cid:40) h · N ( h +1)
i,j ( y ) if h < L L − 1 q 2 · N ( L ) i,j ( y ) if h = L .
The properties stated in the above lemma can be straightforwardly proved using induction.
D.2 CNTK Sketch for Homogeneous Dual Kernels
Our sketching method relies on approximating the dot-product kernel function κ ( · ) and its derivative κ ′ ( · ) with low-degree polynomials via Taylor expansion, and then applying P OLY S KETCH to the resulting polynomial kernels.
Our sketch computes the features for each pixel of the input image, by tensor product of the sketches for function κ ( · ) at consecutive layers, which in turn can be sketched efficiently by P OLY S KETCH .
Additionally, the features of pixels that lie in the same patch get locally combined at each layer via direct sum operation.
This precisely corresponds to the convolution operation in neural networks.
We start by presenting our CNTK Sketch algorithm in Algorithm 4 and prove the correctness and runtime of our procedure in Theorem 6.
Theorem 6 (Correctness and Runtime of Algorithm 4) .
Suppose that the dual kernel K σ is homogeneous as per Equation (17) also assume that κ (1) = κ ′ (1) = 1 .
Fix some ε > 0 and
L ∈ Z > 0 and suppose that (cid:101) κ ( t ) and (cid:101) κ ′ ( t ) are degree- p polynomials with non-negative coefficients that satisfies ( 1 ) max t ∈ [ − 1 , 1] | (cid:101)
κ ( t ) − κ ( t ) |
= O (cid:0) εL (cid:1) and max t ∈ [ − 1 , 1] | (cid:101) κ ′ ( t ) − κ ′ ( t ) | =
O (cid:0) εL (cid:1) , ( 2 ) max | t |≤ 1+ O ( ε ) | (cid:101) κ
( t + γ ) − (cid:101) κ ( t )
| ≤ O ( γ ) and max | t |≤ 1+ O ( ε ) | (cid:101) κ ′ ( t + γ ) − (cid:101) κ ′ ( t ) | ≤ O ( γ ) for any | γ | ≤ O ( ε ) .
If m = Ω (cid:16) L 4 p ε 2 · log 3 n (cid:17) and m ′
=
Ω (cid:16) L 2 ε 2 · log 3 n (cid:17) , then for any y, z ∈ R d 1 × d 2 × c , the output of Algorithm 4 satisfies
[Equation] [Equation]
Furthermore, for every image x ∈ R d 1 × d 2 × c , Ψ ( L ) cntk ( x ) ∈ R m ′ can be computed in time
O (cid:0)
Lp 2 m log m · d 1 d 2 (cid:1) .
Proof.
The correctness proof is by induction on the value of h = 0 , 1 , 2 , . . .
L .
More formally, consider the following invariants for every iteration h = 0 , 1 , 2 , . . .
L of the algorithm:
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
j,ℓ ( x ) 8:
Let Q 2 be a degree-2 P OLY S KETCH with target dimension m ′
9: if h =
L then 10: for every i ∈ [ d 1 ] , j ∈ [ d 2 ] compute ψ
[Equation]
x ) ←
[Equation]
11: else 12: for every i ∈ [ d 1 ] and j ∈ [ d 2 ] compute
[Equation]
←
[Equation] [Equation]
13: return Ψ ( L ) cntk ( y, z ) :
=
[Equation]
P 1 ( h ) :
Simultaneously for all i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] : (cid:12)(cid:12)(cid:12)(cid:68) ϕ
( h ) i,j ( y ) , ϕ ( h ) i ′ ,j ′ ( z ) (cid:69) − Γ
( h ) i,j,i ′ ,j ′ ( y, z ) (cid:12)(cid:12)(cid:12) ≤ ( h + 1) · ε 60 L 2 · (cid:113) N ( h ) i,j ( y ) ·
N ( h ) i ′ ,j ′ ( z ) q, (cid:12)(cid:12)(cid:12)(cid:12)(cid:13)(cid:13)(cid:13)
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
ψ ( h ) i,j ( y ) (cid:13)(cid:13)(cid:13) 2 2 − Π
( h )
i,j,i,j (
[Table]
, y ) (cid:12)(cid:12)(cid:12)(cid:12) ≤ ε 10 ·
[Equation]
[Equation] [Equation]
[Table] [Table]
ψ
[Table]
h ) i ′ ,j ′ (
[Table]
) (cid:13)(cid:13)(cid:13)
[Table]
− Π (
[Table]
) i ′ ,j ′ ,i ′ ,j ′ ( z, z
[Table]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
ε 10 ·
[Equation]
We prove that probabilities Pr[ P 1 (0)] and Pr[ P 2 (0) | P 1 (0)] are both greater than 1 − 1 poly( n ) .
Additionally, for every h = 1 , 2 , . . .
L , we prove that the conditional probabilities Pr[ P 1 ( h ) | P 1 ( h − 1)] and
Pr[ P 2 ( h ) | P 2 ( h − 1) , P 1 ( h ) , P 1 ( h − 1)] are greater than 1 − 1 poly( n ) .
These invariants immediately give the correctness proof.
The base of induction corresponds to h = 0 .
By line 3 of the algorithm, ϕ (0) i,j ( y ) =
y i,j, : and ϕ (0) i ′ ,j ′ ( z ) =
z i ′ ,j ′ , : , therefore, by using Equation (70), it trivially holds that Pr[ P 1 (0)] = 1 ≥ 1 − 1 poly( n ) .
Moreover, by line 3, we have that ψ (0) i,j ( y )
=
and ψ (0) i ′ ,j ′ ( z ) = 0 , thus, by Equation (72), it trivially holds that Pr[ P 2 (0) | P 1 (0)]
= 1 ≥ 1 − 1 poly( n ) .
This completes the base of induction.
We now proceed to prove the inductive step .
By assuming the inductive hypothesis for h − 1 , we prove that statements P 1 ( h ) and P 2 ( h ) hold.
More precisely, first we condition on the statement
P 1 ( h − 1) being true for some h ≥ 1 , and then prove that P 1 ( h ) holds with probability at least 1 − 1 poly( n ) .
Next we show that conditioned on statements P 2 ( h − 1) , P 1 ( h ) , P 1 ( h − 1) being true,
P 2 ( h ) holds with probability at least 1 − 1 poly( n ) .
This will complete the induction.
First, by conditioning on the inductive hypothesis P 1 ( h − 1) and using the definition of µ ( h ) i,j ( · ) in line 5 of the algorithm and applying Cauchy–Schwarz inequality and invoking Lemma 4 we find that,
[Equation]
where the last line follows from Equation (69)
[Equation]
Furthermore, if we let the collection of vectors (cid:110)
Z ( h ) i,j,ℓ ( y ) (cid:111) p ℓ
=0 and (cid:110)
Z ( h )
i,j,ℓ ( z ) (cid:111) p ℓ =0 be defined as per line 5 of the algorithm
[Equation]
then by Theorem 5 and union bound, the following inequalities hold, with probability at least 1 − 1 poly( n ) , simultaneously for all ℓ = 0 , 1 , 2 , . . .
p , all i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] :
[Equation]
Therefore, by Cauchy–Schwarz inequality, we find that with probability at least 1 − 1 poly( n ) , the following holds simultaneously for all i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] :
[Equation]
77) where B :=
[Equation]
By conditioning on the inductive hypothesis P 1 ( h − 1) and using Lemma 4 we have, (cid:12)(cid:12)(cid:12)(cid:12)(cid:13)(cid:13)(cid:13) µ ( h ) i,j ( y ) (cid:13)(cid:13)(cid:13) 2 2 − 1 (cid:12)(cid:12)(cid:12)(cid:12) ≤ h · ε 60 L 2 , and (cid:12)(cid:12)(cid:12)(cid:12)(cid:13)(cid:13)(cid:13) µ ( h ) i ′ ,j ′ ( z ) (cid:13)(cid:13)(cid:13) 2 2 − 1 (cid:12)(cid:12)(cid:12)(cid:12) ≤ h · ε 60 L 2 .
Therefore, the precondition of the theorem implies that (cid:12)(cid:12)(cid:12)(cid:101) κ (cid:16)
∥ µ ( h ) i,j ( y ) ∥ 22 (cid:17) − (cid:101) κ (1) (cid:12)(cid:12)(cid:12) ≤ h · ε 60 L 2 and (cid:12)(cid:12)(cid:12)(cid:101) κ (cid:16)
∥ µ ( h ) i ′ ,j ′ ( z ) ∥ 22 (cid:17) − (cid:101) κ (1) (cid:12)(cid:12)(cid:12) ≤ h · ε 60 L 2 .
Consequently, because (cid:101) κ (1) ≤ 1 . 01 κ (1) = 1 . 01 , we find that
[Equation]
By plugging this into Equation (77) we find that the following holds simultaneously for all i, i ′ ∈
[ d 1 ] and all j, j ′ ∈ [ d 2 ] , with probability at least 1 − 1 poly( n ) , (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:68)
[Equation]
We recall that A :=
[Equation]
Note that by Lemma 4 and Equation (69), − 1 ≤ A ≤ 1 .
Hence, using the precondition of the theorem and Equation (75) to find that,
[Equation]
By incorporating the above inequality into Equation (78) using triangle inequality we find that, with probability at least 1 −
[Equation]
( n ) , the following holds simultaneously for all i, i ′
∈
[ d 1 ] and all j, j ′ ∈ [ d 2 ] :
[Equation]
Additionally, since − 1 ≤ A ≤
[Equation]
using the preconditions of the theorem we can conclude that | (cid:101) κ ( A ) − κ ( A ) | ≤ ε 76 L 2 .
By combining the above inequality with Equation (79) via triangle inequality and using the fact that, by Equation (70)
[Equation]
we get the following inequality, with probability at least 1 −
[Equation]
Similarly, we can prove that with probability at least 1 − 1 poly( n ) the following hold, simultaneously for all i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] ,
[Equation]
ϕ ( h ) i,j ( y )
[Equation]
2 2 −
[Equation]
( h ) i,j,i,j
[Equation]
y, y )
[Equation] [Equation]
60 L
[Equation]
ϕ ( h ) i ′ ,j
[Equation]
( z
[Equation]
Γ ( h ) i
[Equation]
,j
[Equation]
,i
[Equation]
( z, z
[Equation]
h
[Equation]
1) ε 60 L· N ( h
[Equation]
i
[Equation]
j
[Equation]
z ) q.
This is sufficient to prove the inductive step for statement P 1 ( h ) , i.e., Pr[ P 1 ( h ) | P 1 ( h − 1)] ≥ 1 −
Now we prove the inductive step for statement P 2 ( h ) .
That is, we prove that conditioned on
P 2 ( h − 1) , P 1 ( h ) , and P 1 ( h − 1) , P 2 ( h ) holds with probability at least 1 − 1 poly( n ) .
First note that using the definition of .
ϕ ( h ) i,j ( y ) , .
ϕ ( h ) i ′ ,j ′ ( z ) in line 7 of the algorithm and Equation (76),we find that with probability at least 1 − 1 poly( n ) , the following holds simultaneously for all i, i ′
∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] :
[Equation]
(80) where (cid:98) B := 1 q 2 · (cid:114) (cid:101) κ ′ (cid:16) ∥ µ ( h
[Equation]
i
[Equation]
y
[Equation]
∥
[Equation]
· (cid:101) κ
[Equation]
(cid:16)
∥ µ ( h )
i
[Equation]
j
[Equation]
( z ) ∥ 22 (cid:17) .
By conditioning on the inductive hypothesis
P 1 ( h − 1) and using Lemma 4 we have, (cid:12)(cid:12)(cid:12)(cid:12)(cid:13)(cid:13)(cid:13)
[Equation]
ε 60 L 2 and (cid:12)(cid:12)(cid:12)(cid:12)(cid:13)(cid:13)(cid:13) µ ( h ) i ′ ,j ′ ( z ) (cid:13)(cid:13)(cid:13) 2 2 − 1 (cid:12)(cid:12)(cid:12)(cid:12) ≤ h · ε 60 L 2 .
Therefore, the precondition of the theorem implies that (cid:12)(cid:12)(cid:12)(cid:101) κ ′ (cid:16) ∥ µ ( h ) i,j ( y ) ∥ 22 (cid:17) − (cid:101) κ ′ (1) (cid:12)(cid:12)(cid:12) ≤ h · ε 20 Land (cid:12)(cid:12)(cid:12)(cid:101) κ ′ (cid:16)
∥ µ ( h ) i ′ ,j ′ ( z ) ∥ 22 (cid:17) − (cid:101) κ ′ (1) (cid:12)(cid:12)(cid:12) ≤ h · ε 20 L 2 .
Consequently, because (cid:101) κ ′ (1) ≤ 1 . 01 κ ′ (1) = 1 . 01 , we find that
[Equation]
By plugging this into Equation (80) we get the following, with probability at least 1 − O (cid:0) δL (cid:1) ,
[Equation]
Furthermore, we can use the precondition of the theorem to find that Equation (75) implies the following,
[Equation]
By incorporating the above inequality into Equation (81) using triangle inequality, we find that, with probability at least 1 − 1 poly( n ) , the following holds simultaneously for all i, i ′ ∈
[ d 1 ] and all j, j ′ ∈ [ d 2 ] :
[Equation]
ε
[Equation]
Since − 1 ≤ A ≤ 1 , we can use the precondition of the theorem to conclude | (cid:101) κ ′ ( A ) − κ ′ ( A ) | ≤
[Equation]
By combining this inequality with Equation (82) via triangle inequality and using the fact that .
Γ ( h ) i,j,i ′ ,j ′ ( y, z ) = 1 q 2 · κ ′ ( A ) , we get the following bound simultaneously for all i, i ′ ∈
[ d 1 ] and all j, j ′ ∈ [ d 2 ] , with probability at least 1 −
[Equation]
Similarly we can prove that with probability at least 1 − 1 poly( n ) , the following hold simultaneously for all i, i ′ ∈
[ d 1 ] and all j, j ′ ∈ [ d 2 ] ,
[Equation]
2 2 − .
Γ h )
i,j,i,j ( y, y ) (cid:12)(cid:12)(cid:12)(cid:12) ≤ 1 q 2 · ε 8 L, and (cid:12)(cid:12)(cid:12)(cid:12)(cid:13)(cid:13)(cid:13) .
ϕ ( h ) i ′ ,j ′ ( z ) (cid:13)(cid:13)(cid:13) 2 2 − .
Γ ( h ) i ′ ,j ′ ,i ′ ,j ′ ( z, z ) (cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
We will use Equation (83) and Equation (84) to prove the inductive step for P 2 ( h ) .
Next, we consider two cases for the value of h .
When h < L , the vectors ψ ( h ) i,j ( y ) , ψ ( h ) i ′ ,j ′ ( z ) are defined in line 10 and when h
=
L , these vectors are defined differently in line 12.
First we consider the case of h < L .
If we let f i,j := ψ ( h − 1) i,j ( y ) ⊗ .
ϕ ( h ) i,j ( y ) and g i ′ ,
j ′ :
= ψ ( h − 1) i ′ ,j ′ ( z ) ⊗ .
ϕ ( h ) i ′ ,j ′ ( z ) and η ( h ) i,j ( y ) :
= (cid:0) Q 2 · f i,j (cid:1)
⊕ ϕ
( h ) i,j ( y ) and η ( h ) i ′ ,j ′ ( z ) :
= (cid:0) Q 2 · g i ′ ,j ′ (cid:1)
⊕ ϕ
( h ) i ′ ,j ′ ( z ) , then by Theorem 5 and union bound, with probability at least 1 − 1 poly( n ) , we have the following inequalities simultaneously for all i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] :
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
Now we bound the term
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Table]
using Equation (85), Equation (84)
[Table]
and Lemma 4 along with inductive hypotheses P 2 ( h − 1) .
With probability at least 1 − 1 poly
[Equation]
) the following holds simultaneously for all i, i ′
∈
[ d 1 ] and all j, j ′ ∈ [ d 2 ] :
[Equation]
where the last line above follows from Lemma 4 together with the fact that ˙Γ ( h ) i,j,i,j ( y, y )
=
˙Γ ( h ) i ′ ,j ′ ,i ′ ,j ′ ( z, z ) =
1 q 2 .
By combining the above with inductive hypotheses P 1 ( h ) , P 2 ( h − 1) and Equation (83) via triangle inequality and invoking Lemma 4 we get that the following holds simultaneously for all i, i ′ ∈
[ d 1 ] and all j, j ′ ∈ [ d 2 ] , with probability at least 1 −
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
By plugging the above bound into the definition of in line 10 of the algorithm using triangle inequality and using Equation (72) we get the following with probability at least 1 − 1 poly( n ) :
[Equation]
,j ′ ( z ) .
Similarly, we can prove that with probability at least 1 − 1 poly( n ) the following hold simultaneously for all i, i ′ ∈
[ d 1 ] and all j, j ′ ∈ [ d 2 ] ,
[Equation]
This is sufficient to prove the inductive step for statement P 2 ( h ) , in the case of h < L , i.e.,
Pr[ P 2 ( h ) | P 2 ( h − 1) , P 1 ( h ) , P 1 ( h − 1)] ≥ 1 − 1 poly( n ) .
Now we prove the inductive step for P 2 ( h ) in the case of h = L .
Similar to before, if we let f i,j :=
ψ ( L − 1) i,j ( y ) ⊗ ˙ ϕ
( L ) i,j ( y ) and g i ′ ,j ′ :
= ψ ( L − 1) i ′ ,j ′ ( z ) ⊗ ˙ ϕ
( L ) i ′ ,j ′ ( z ) , then by (12), we have ψ ( L ) i,j ( y )
= (cid:0) Q 2 · f i,j (cid:1) and
ψ ( L ) i ′ ,j ′ ( z ) =
(cid:0) Q 2 · g i ′ ,j ′ (cid:1) .
Thus by Theorem 5 and union bound, we find that, with probability at least 1 − 1 poly( n ) , the following inequality holds simultaneously for all i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] :
[Equation]
Therefore, using (84) and Lemma 4 along with inductive hypotheses P 2 ( L − 1) , with probability at least 1 − 1 poly( n ) , the following holds simultaneously for all i, i ′
∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] ,
[Equation]
q
[Equation]
By combining the above with inductive hypotheses P 1 ( L ) , P 2 ( L − 1) and Equation (83) via triangle inequality and invoking Lemma 4 and also using the definition of Π ( L ) ( y, z ) given in Equation (73), we get that the following holds, simultaneously for all i, i ′ ∈
[ d 1 ] and j, j ′ ∈ [ d 2 ] , with probability at least 1 − 1 poly( n ) ,
[Equation]
[Equation]
This proves the inductive step for statement P 2 ( h ) , in the case of h = L , i.e., Pr[ P 2 ( L )
| P 2 ( L − 1) , P 1 ( L ) , P 1 ( L − 1)] ≥ 1 − 1 poly( n ) .
The induction is complete and hence the correctness of
Algorithm 4 is proved by union bounding over all h = 0 , 1 , 2 , . . .
L .
The runtime of the algorithm immediately follows by invoking Theorem 5 because computing vector
Z ( h ) i,j,ℓ ( x ) for every i, j, ℓ and h = 1 , 2 , . . .
L dominates the runtime of this algorithm.
[Figure]
As an example, let us invoke Algorithm 4 and Theorem 6 on the CNTK with GAP corresponding to the normalized Gaussian dual kernel K G , defined per Equation (13).
Note that the dot-product factor corresponding to this dual kernel is κ ( t ) =
exp( t − 1) .
The truncated Taylor series of this function is (cid:101) κ ( t )
= (cid:80) pj =
t j e · j ! and the truncated Taylor series expansion of the derivative of this function is (cid:101) κ ′ ( t ) =
(cid:80) pj =
t j e · j ! .
If p = Ω(log n ) then it can be verified that polynomials (cid:101) κ ( t ) , (cid:101) κ ′ ( t ) satisfy the preconditions of Theorem 6.
Therefore, by Theorem 6, we can sketch the
CNTK kernel using O (cid:16) L 4 ε 2 · d 1 d 2 · poly (log n ) (cid:17) running time.
Also the target dimension of the sketch is m ′
=
O (cid:16) L 2 ε 2 log 3 n (cid:17) .
So the runtime of our Algorithm 4 is only linear in the number of image pixels d 1 d 2 , which is in stark contrast to quadratic scaling of the exact CNTK computation [9].
In fact, using our CNTK sketching method, the kernel regression can be solved approximately in time O (
cid:16) L
[Equation]
n
[Equation]
poly (log n ) (cid:17) , which is significantly faster than the exact kernel regression which takes Ω (cid:0) L ( d 1 d 2 · n ) 2 (cid:1) when the number of pixels d 1 d 2 or the training set size n are large.
E Gauss-Hermite Quadrature Derivation
Here we provide more details on Section 3.3.
Utilizing the whitening transformation of covariance Λ used in the proof of Theorem 3 in Appendix B.5 the dual activation function can be expressed as
[Equation]
[Equation]
Here ( x i , w i ) , correspond to roots of q -th degree (Physicist’s)
Hermite polynomial H q ( x ) and associated weights [55]
[Equation]
[Equation]
where the conversion between physicist’s to probabilist’s convention H n ( x ) = 2 n 2 h n ( √ 2 x )
The roots are obtained by Golub-Welsch algorithm [63] and can be found in scientific computing package
[Figure]
n t
i
[Figure]
i
[Figure]
i d _
[Figure]
i
[Figure]
activation functions.
The of ABReLU such as Scipy [64]’s scipy.special.roots_hermite function.
For alternative parameterization for multivariate Gauss-Hermite quadrature, refer to notes by Jäckel [65].
For activation function where exact dual activation is known, one can measure the error from the quadrture.
In Figure 3, we compute errors for ReLU , Abs (i.e., σ ( t ) =
| t | ), sin , Gaussian, erf and
GeLU activations.
For non-smooth activation ( ReLU , Abs), approximation error decays as powerlaw like where as for smooth activation the error decays exponentially as one increases Hermite polynomial degree q .
We utilize this method as well as our expanded dual activation Table 1 to compare performance of various activation functions on CIFAR-10 dataset.
In Figure 4, we study three architectures; 1 hidden layer fully connected network (FC1, equivalent to pure dual activation kernel), 8 layer convolutional network with vectorization (CV8), and Myrtle5 network.
We compared classification performance on subset of CIFAR-10.
In each plot activation function is sorted by NTK’s classification performance.
One notable observation is that normalized Gaussian shows consistently best performance across architecture.
Also note that smooth activations computed with Gauss-Hermite quadrature (denoted by ∗ ) shows almost identical performance when analytic form is available (e.g. GeLU, Erf, RBF(1/2)).
Notable outlier is FC1 NTK with ReLU, however we expect that non-smooth activation may be approximated poorly.
It’s also interesting to observer sigmoid-like activations (Sigmoid, Tanh, Erf) performs poorly across the board whereas ReLU-like activations (Normalized Gaussian, ABReLU,
ReLU, GeLU, RBF) are among high performant group.
This is implemented as ElementwiseNumerical in our code.
F Table of dual activation functions
We describe dual kernel functions of several activations and their derivatives in Table 2.
One can generalize duel kernels of affine transformations of these activations.
Specifically, if (cid:101) σ ( t ) =
A · σ ( Bt ) + C for some A, B, C ∈ R then
[Equation]
which follows from that
[Equation]
Below we provide detailed expressions of omitted dual kernel formulations in the table.
Activation
[Table]
Implemented as
Rectifiedmonomials [44]
[Table]
[Table] [Table]
[Table] [Table]
F.1 Rectified monomials
Cho and Saul [44] proposed closed-form expressions of dual kernel functions for rectified activations, i.e., σ ( t ) =
t n · 1 { t ≥ 0 } for n ≥ 0 , as
[Equation]
where for θ = cos − 1 ( c ) ∈ [0 , π ]
[Equation]
For n = 0 and
[Equation]
Applying Theorem 3 provides that
[Equation]
These are implemented in our code as RectifiedMonomial (with a special case of Sign for convenience).
F.2 ABReLU, Leaky ReLU, Abs
ABReLU activation function is given by
[Equation]
The dual kernel functions can be obtained by extension of [44] which is worked out in [50, 51]
[Equation]
and
[Equation]
A special case of ABReLU covers leaky ReLU [66] ( B = 1 ), that is,
[Equation]
and the corresponding dual kernel functions are
[Equation]
and
[Equation] [Equation]
Another special case is the absolute value function (Abs)
( A = − 1 , B = 1 ), that is,
[Equation]
) = |
[Equation]
| ,
[Equation]
107) and the corresponding dual kernel functions are
[Equation]
and
[Equation] [Equation]
These are respectively implemented as ABRelu , LeakyRelu , and Abs in [42].
F.3 Sinusoidal and RBF
A generalized sinusoidal activation is given by
[Equation]
The corresponding dual kernels are
[Equation]
[Equation]
Note that the generalized sinusoidal activation with a = √ 2 , b = √ 2 A , and c = π 4 gives that
[Equation]
[Equation]
which corresponds to (translation invariant) the Gaussian RBF kernel:
[Equation]
for some x, y ∈ R d .
Moreover, one could consider mixture of activation functions as discussed in Louart et al.
[47], Adlam et al.
[67] of 50% cos and 50% sin which also leads to stationary kernel
[Equation] [Equation]
In order to obtain stationary kernel with respect to inputs, one only needs to insert these transformation at the first layer of the network as highlighted in implicit neural representation (e.g. NeRF) [37,
F.4 Error function
The error function is given by
[Equation] [Equation]
Following [43] and applying Theorem 3, we get
[Equation]
An affine transformation of the error function could behave similar to sigmoid activation function with range (0 , 1) , that is,
[Equation]
The corresponding dual kernels can be obtained by applying affine transformation to that of error function as discussed in Equation (93).
The error function is implemented in [42] as Erf , and we release Sigmoid_like in our code.
F.5 Gaussian function
Consider Gaussian function
[Equation]
One can obtain k σ [43],
[Equation]
(121) and using Theorem 3 obtain
[Equation]
[Equation]
Note that Gaussian function itself can be obtained as derivative of Affine Erf thus could use Theoremwith Affine Erf.
This function is implemented as Gaussian in our code.
F.6 GeLU
The Gaussian Error Linear Unit (GeLU) [33] is defined as
[Equation]
where erf( · ) is the Gauss error function.
For efficiency, sometimes approximate formulation
[Equation]
is used.
We note that GeLU activation function is becoming popular in recent language models such as BERT [34], ALBERT [68], RoBERTa [69] and GPT [35, 36].
The corresponding dual kernel is studied in Tsuchida et al.
[Equation]
Using Theorem 3, we have
[Equation]
This is implemented in our code as Gelu .
F.7 Monomials
Consider monomials
[Equation]
The dual activation function is given in terms of Hypergeometric function 2 F 1 .
For even power n ∈ 2 Z
[Equation]
For odd power n ∈ 2 Z +k
[Equation]
The first five k σ n s are
[Equation]
Note that dual activation functions of monomials are also obtained from Theorem 1 by choosing c n = 1 , c n − 1 = · · · = c 0 = 0 .
Moreover, obtaining k σ ′ n is simple either by σ n ( t ) ′ =
nt n − 1 or applying Theorem 3 to above expressions on k σ n .
These are implemented in our code as Monomial .
F.8 Gabor
Let us consider a simple version of localized oscillatory activation function given by
[Equation]
The dual actiavtion of Gabor function can be expressed as
[Equation]
and that of derivative of Gabor function can be obtained using Theorem 3 as k Gabor
[Equation]
[Equation] [Equation]
This is implemented in our code as Gabor .
F.9 ELU
For Exponential Linear Unit (ELU) [70]
[Equation]
The k σ ( a, b, c ) is computed in Tsuchida et al.
[48] and we refer to the original paper for the expression.
Note that k σ ′ for ELU has not been computed but Theorem 3 allows to simply obtain it using expression in Tsuchida et al.
G Additional Experiment: Kernel Informed Activation
We explore an activation informed by the normalized Gaussian kernel that achieves the best performance among neural kernels [54].
Although the exact activation is unknown, one can conduct a reverse engineering to find a proper activation whose dual kernel is known and close to the normalized
Gaussian.
In particular, we focus on the ABReLU activation and recall that its dual kernel is
[Equation]
for some A, B ∈ R .
Observe that k ABReLU is also homogeneous as like the normalized Gaussian, i.e.
, k σ ( a, b, c ) =
ab · κ σ ( c ) for c ∈ [ − 1 , 1] .
We find two slope variables A, B by fitting κ σ at extreme points, i.e., κ ABReLU ( c ) =
exp( c − 1) for c = ± 1 .
This turns into a quadratic equation and gives us
[Equation]
which is illustrated in Figure 2 ( left ).
We train a 5 -layer ConvNet (known as Myrtle-5 [54]) ofwidth for CIFAR-10 classification.
Similar to the CNTK experiment in Section 5, we convert image classes into 10 -dimensional one-hot vectors and pre-process CIFAR-10 images with regularized
ZCA [54, 58].
We use the SGD optimizer with initial learning rate 0 .
1 , Nesterov momentum with factor 0 .
9 and ℓ 2 regularizer 0 . 0005 .
The batch size is set to 64 .
The network is trained by minimizing the mean-squared-error (MSE) loss and we report the best test accuracy for 200 epochs.
Interestingly, the ABReLU can achieve the highest test accuracy compared to ReLU, GeLU, Erf and parameterized
ReLU (PReLU) activations.
This supports a connection between infinite width neural kernels and finite width networks in aspect of activation.
[Figure]