Kernel Regression with Infinite-Width Neural Networks on Millions of Examples
Authors: Ben Adlam
Jaehoon Lee
Shreyas Padhy
Zachary Nado
Jasper Snoek
Venue: Submitted to ICLR 2023 (2023)
Keywords: gaussian processes
neural tangent kernel
infinite-width neural networks

Abstract: While kernel regression remains an important practical method, its connection to neural networks as their width becomes large has initiated fresh research. These neural kernels have drastically increased performance on diverse and nonstandard data modalities but require significantly more compute, which previously limited their application to smaller datasets. We address this by massively parallelizing their computation across many GPUs. We combine this with a distributed, preconditioned conjugate gradients algorithm to enable kernel regression at a large scale (i.e. up to 5 million examples). Using this approach, we study scaling laws of several neural kernels across many orders of magnitude for the CIFAR-5m dataset. Using data augmentation to expand the original CIFAR-10 training dataset by a factor of 20, we obtain a test accuracy of 91.2\% (SotA for a pure kernel method). Finally, we explore other data modalities, obtaining results on protein and small molecule prediction tasks that are competitive with SotA methods.


Under review as a conference paper at ICLR 2023
KERNEL REGRESSION WITH INFINITE-WIDH NEURAL NETWORKS ON MILLIONS OF EXAMPLES
Anonymous authors Paper under double-blind review
ABSTRACT
While kernel regression remains an important practical method, its connection neural networks as their width becomes large has initiated fresh research. These diverse and nonstandard
neural kernels have drastically increased performance data modalities but require significantly more compute, which previously limited their application smaller datasets. We address this by massively parallelizing their computation across many GPUs. We combine this with distributed, precon- ditioned conjugate gradients algorithm to enable kernel regression at large scale (i.e. up to 5 million examples). Using this approach, we study scaling laws of sev- eral neural kernels across many orders of magnitude for the CIFAR-5m dataset. Using data augmentation to expand the original CIFAR-10 training dataset by a factor of 20, we obtaín a test accuracy of 91.2% (SotA for pure kernel method). Finally, we explore other data modalities, obtaining results protein and small molecule prediction tasks that are competitive with SotA methods.
INTRODUCTION
Kernel methods are often contrasted with deep learning, but recent advances machine learning have identified and developed exciting correspondences (Lee et al., 2018; Matthews et al., 2018; Jacot et al., 2018). While useful method in its own right, kernel regression has been used to better understand neural networks and deep learning. More specifically, if the parameters of neural net- work are treated as random variables whose distribution is set by the initialization, we can view the neural network as a random function. Then as the width of the network becomes large, the distribu- tion of this random function is a Gaussian process with specific covariance function or kernel. We refer kernels that arise from this connection with infinite-width neural networks as neural kernels. The specific kernel is determined by the architecture, inference type and other hyperparameters of the neural network.
Moreover, the connection between neural networks and Gaussian processes has generated many nigh-performance kernels for diverse or nonstandard data modalities, such as images, sequences, and graphs. This performance often comes at cost, the kernels require significantly more com- pute than standard kernels such as RBFs. For example, computing the entire kernel for the CIFAR-10 10 dataset takes less than 1 GPU minute for an RBF kernel but around 300 GPU hours for the Myrtle kernel (Shankar et al., 2020; Lee et al., 2020). However, this increase in compute significantly de- creases the test error rate from around 40% to 10%. The added demands of simply computing entries of the kernel is in addition to challenges posed from the cubic scaling in time and quadratic scaling in memory of inference for kernel regression with dataset size. Approximate inference methods frequently reduce memory requirements by recomputing kernel entries on the fly, which is infeasi- ble for these expensive kernels. Such challenges have limited our understanding of infinite- width models to small datasets. In particular, while scaling laws have been studied across many orders of magnitude for neural networks, the same is not true of their corresponding kernels. Similarly, while it is common to augment training datasets in neural networks, significantly expanding their size, the benefits for kernel methods have not been as thoroughly explored.
1.1 CONTRIBUTIONS
In this work, we address these two main computational challenges (computing large, complex ker- nels and using them for kernel regression) by parallelizing and scaling up existing algorithms to

---

Under review as a conference paper at ICLR 2023
many more machines. This enables us consider significantly larger datasets, currently up five million examples, and therefore study how performance changes as more data are added over many orders of magnitude. While similar studies have been performed for neural networks, they have been lacking for neural kernels. In addition to scaling to larger datasets, we also consider high res- olution images from the Tiny Imagenet dataset, where additional pixels also require more compute. Moreover, our approach is not restricted to image data. In fact, we obtain results for both protein sequence and small molecule datasets and demonstrate that neural kernels are a promising method for medium sized datasets from basic science.
Our contributions include:
By massively parallelizing the computation of the neural kernels, we study kernels on sig- nificantly larger (approximately two orders of magnitude) datasets;
We use a distributed, preconditioned conjugate gradients algorithm to perform inference for these kernels;
We demonstrate scaling laws across several orders of magnitude for fully-connected, CNN- Vec, and Myrtle kernels on the CIFAR-5m dataset;
We study the loss in performance incurred by approximating the linear system from infer- ence compared conjugate gradients;
Using data augmentation to expand the original (IFAR-10 training dataset by factor of 20, we obtain a test accuracy of 91.2% (SotA for a pure kernel method);
We explore other data modalities, obtaining results protein and small molecule predic- tion tasks that are competitive with SotA methods.
1.2 RELATED WORK
The contributions of this work are made possible by recent advances in large scale inference for Gaussian processes and kernel regression and advances in the understanding of the relationship between GPs and deep neural networks.
Wang et al. (2019) showed how to solve the large-scale linear systems within GPs using Conjugate Gradients (CG), which requires only matrix-vector operations with the kernel and doesn't 'it require storing the full kernel in memory. Wang et al. (2019) and Maddox et al. (2022) identify the impor- tance of using a pre-conditioner partially-pivoted Cholesky decomposition) with finite precision. We use a similar CG algorithm to solve our systems and found that the pre- conditioner was critical to convergence, particularly for the non-stationary neural kernels. Due to their eigenspectra, we unfortunately found we required high precision on CPU. In this work, we use these methods to solve even larger systems, up to million examples, but emphasize that the most computationally expensive component is computing the expressive neural kernels.
solve the system
Many methods have been developed to approximate Gaussian processes on larger datasets. Rahimi & Recht (2007) show that stationary kernels can be approximated using finite set of random ba- sis functions. Recently there has been progress in random-feature approximations for expressive, non-stationary kernels using sketching (Zandieh et al., 2021; Han et al., 2022). While this method provides efficient ways for approximating neural kernels, often there is speed-performancee trade- off (Han al., 2022), thus our work focuses on exact computation. Stochastic variational inference (SVI) (Hensman et al., 2013) is tremendously promising approach to scale GPs by optimizing a of inducing points approximate the full GP posterior. However, Wang et al. (2019) found that their SVI baseline underperformed exact GP inference in their experiments. We found that SVI was computationally infeasible due taking gradients through the neural kernels, and using subsets of data as inducing points was less-effective than the full kernel. Other approaches, such as KISS- GP (Wilson & Nickisch, 2015; Stanton et al., 2021), cleverly interpolate over a grid of inducing points, taking advantage of Toeplitz structure, to drastically reduce computational cost.
Rudi et al. (2017) develop FALKON, which uses Nystrom and random feature approximations with pre-conditioning to scale up kernel regression to millions of examples with theoretical guarantees on performance. This is extended to potentially billions of examples by Meanti et al. (2020), through distributed GPU acceleration and various optimizations. EigenPro and EigenPro2 (Ma & Belkin, 2017; 2019) accelerate and scale up kernel regression by developing pre-conditioned stochastic gra- dient descent. While we observe that neural kernels scale remarkably well in terms of predictive
2

---

Under review as a conference paper at ICLR 2023
Large
split smaller blocks for ID
Blocks computed across many parallel jobs
Miro-batch computed with several GPUs
micro-beltched
column batch size, N/m
Myrtie-NTK CIFAR-5m
row batch size, Nin
Worker (0.0)
Worker (0,m-1)
ΣK
Σ K
Classificalton Error
Mydy
Scbnet regresorr
Siect dagonal
(1 D)
Worker (1, (1,m-1)
IK
I, K
0.3
Compugats grasients
Jayesan carm machie
Worker (2.00
Worker (2,m-1)
I
M
10'
10
Worker (0-1.0)
...... Worker (0-1.111)
Σ M
K.K..
Rank Partition Size
Figure I: (Top) An upper diagonal of large kernel matrix is split into smaller, dimension 5,000x5,000 blocks for multi-threading of IO. When computing the kernel, the blocks can be com- puted independently by machines that need not be co-located. Each block is then batched further (to meet device memory limitations) and computed by the Neural Tangents library on several GPUs. (Bottom left) It is natural to split the kernel by rows across the workers, then each worker can re- ceive the vector from the host, compute the matrix product between the rows in its memory with the vector, and return this chunk of the result to the host. The host can then aggregate the results from the workers by concatenating. However, for larger systems with many workers, simply communi- cating a large vector to all workers becomes a bottleneck. Instead the kernel can also be partitioned columns. This means that each worker only needs subset of the vector entries, and found this dramatically increased communication speed. The cost of this approach is that slightly more complex aggregation required the host. (Bottom right) Classification error rate as a function of rank or partition size for four different kernel approximations. The kernel used is a 10-layer Myr- tle NTK for 1.6 million examples of the CIFAR-5m dataset. The performance of CG is shown as dashed line for comparison and exceeds all other results.
performance, using these methods could make them more computationally practical. However, found that we have required the full kernel at high precision to obtain high performance.
The correspondence between random functions from the initialization of infinite-width neural net- works and Gaussian processes, called the neural network Gaussian process (NNGP), was first noted by Neal (1994) and analytically derived by Williams (1996) for single-hidden layer, funny-connected networks. More recently this was extended to deep neural networks for fully-connected (Lee et al., 2018; Matthews al., 2018) and convolutional architectures (Novak et al., 2019; Garriga-Alonss et al., 2019). This correspondence to Gaussian processes can be extended to infinite-widd networks trained by gradient descent via the Neural Tangent Kernel (NTK) (Jacot et al., 2018) and to many other architectures, (Yang, 2019; 2020) including self-attention layers (Hron et al., 2020) and sim- ple graph neural networks (Du et al., 2019). The covariance functions or kernels associated to these Gaussian processes determined by the neural network's 's architecture and other hyperparameters from initialization, and we refer to them as neural kernels.
While these kernels are theoretically useful to better understand the functional prior of neural net- works, including their uncertainty properties (Adlam et al., 2020), they have also been used more practically. The NTK for convolutional architectures has been applied to image data (Arora al., 2019; Li et al., 2019; Shankar et al., 2020) and offers unique solutions to practical deep learning problems such data distillation (Nguyen et al., 2021a;b). The JAX-based (Bradbury et al., 2018) python library, Neural Tangents (Novak et al., 2020), has helped to enable these practical applica- tions by providing efficient computation of neural kernels on CPU, GPU and TPU.
3

---

Under review as a conference paper at ICLR 2023
2 APPROACHES TO LARGE-SCALL KERNEL METHODS
In this section, review the different approaches to large-scale kernel regression. We discuss the particular challenges introduced by neural kernels and how we addressed them. Finally, we compare the performance for the CIFAR- 5m dataset with 10-layer Myrtle NTK.
There are two main approaches to applying kernel methods to larger datasets, i.c. to solve a large linear system. First, the original linear system can be replaced by an alternative linear system that has a particular algebraic form that can be solved exactly and efficiently. Second, the original linear system can be maintained but only solved approximately. These approximate solvers are often any- time algorithms, where the current solution can be returned at anytime anyiile-typically when its residual sufficiently small.
Approximations to the linear system: Many methods fall under this category. Here we focus on four different types. First, there are the low-rank approximations, Nyström and Subset of Regressors, that can be inverted in time (r)22) where r' is the rank and n is the number of training examples (Williams & Rasmussen, 2006).
Second, there the Block Diagonal approximation that sets all kernel entries outside of some blocks along the diagonal to zero. Due to its algebraic form, this kernel is easy to invert and takes time O(r2n) where r the block or partition size. Even without parallelization, this reduces the running time by factor of the number of blocks squared.
Finally, the Bayesian Committee Machine (Tresp, 2000; Williams & Rasmussen, 2006) is an approx- imation to the likelihood of a Gaussian process that partitions the dataset, fitting separate regressors to the partitions under assumptions of independence. Here we use the posterior mean under this approximation, which be computed efficiently. Since this is transductive method, the time complexity depends on the size of the test set but is comparable to the methods above when the number of test examples equals the partition size.
One advantage of these methods is that they do not require access to all kernel entries, and so can avoid the quadratic scaling from computing and storing the entire kernel. The primary additional hyperparameters in these methods are the rank or partition size. Typically increasing these hyperpa- rameters improves performance at the cost of additional compute (see Fig. 1). However, in addition to their size, exactly which examples are included in the low-rank approximation or how the data are partitioned can affect performance. While several methods are available to select these examples more optimally, the most common approach is still to select uniformly at random, which we do here also. In addition, it was necessary to tune the jitter term in the Nyström approximation, since setting to the default value of 10-6 we used elsewhere yielded poor results.
Iterative solvers: Rather than solving the linear system to machine precision, we might be sat- isfied with a solution that had a sufficiently small residual- especially because the predictions are unlikely to change significantly. Conjugate gradients (CG) is one approach to this (Shewchuk et al., 1994; Wang et al., 2019; Meanti et al., 2020). CG iteratively finds the best solution in a subspace of dimension equal to the number of steps. This subspace is defined by sequence of vectors that are conjugate with respect to the kernel. Each step of CG takes quadratic time, and in the worst case a linear number of steps must be performed, but this is normally not necessary. To avoid this worst-case performance, pre-conditioning is essential. We used the standard partial pivoted Cholesky pre-conditioner (Wang al., 2019) and found it significantly decreased the number of steps required to reach a small residual.
A key subroutine in CG is the computation of kernel-vector products and is the only place the kernel occurs directly. This presents opportunities for parallelization. In particular, we use a host CPU to the main CG algorithm and act as a server control worker CPUs who compute kerne--vector products (see Fig. 1 (Bottom left) for details).
Additional challenges from neural kernels: Previous large-scale implementations of CG have traded off memory and time by never storing the whole kernel in memory and instead recomputing it on the fly. While this approach works well for simple kernels, such as RBFs, it is infeasible for neural kernels. In particular, computing the kernel once represents the majority of the compute com-
Note that some approaches can be interpreted in both ways.
4

---

Under review as a conference paper at ICLR 2023
pared to running CG for inference. Indeed reducing the computational burden of neural kemnels is a topic of current research and can be faster without decreasing performance (Zandich et al., 2021; Han et al., 2021), though often there is perfformee tradeoff (Han et al., 2021). The problem has not been completely solved. Moreover, there is evidence that even computing them at f oatt32 precision compared to fl f1cat64 precision degrades performance (Leee et al., 2020). Thus the ap- proach we take here to obtaining the kernel for larger datasets is to parallelize their computation across many machines (see Fig. 1 (Top) for details). Fortunately very little communication coor- dination is needed across machines, and preemption-safe routines are easy to design. The kernel can then be computed once and stored on disk in collection of sub-blocks for faster, multi-threadee IO. During CG, the kernel must be stored either completely in RAM, which is fastest but requires many machines for large kernels, or can be stored on disk with very fast IO to load chunks of the kernel into RAM.2
The spectra of neural kernels can be incredibly ill-conditioned (Lee et al., 2020), which presents challenges for inference and reinforces the need for pre-conditioning, Often the spectra are very diffuse, spanning several orders of magnitude and showing power-law decay.
Performance comparison: In Fig. 1, we compare the accuracy of the different methods for the 10 layer Myrtle NTK on 1.6 million training examples from the CIFAR-5m dataset.² We find that as the rank or partition size increases, the performance of all methods improves with Subset of Regressors and Nyström performing best. We did not increase above 50,000 as this is close to the limit of the largest linear system that can be solved on standard CPU. Note that while the gap to CG is reduced, there is still a large reduction in performance.
SCALING LAWS FOR NEURAL KERNELS
Known as Neural Scaling Laws, there has been substantial recent interest in how the performance of neural networks improves as more data are available for training and their number of parameters increases (Hestness al., 2017; Kaplan al., 2020; Rosenfeld al., 2019; Rosenfeld, 2021; Bahri et al., 2021). Frequently, these quantities increase in tandem. Empirically, it has been found that per- formance improves as a power law, for which many works seek estimate the exponent. Projecting the potential improvement in performance from scaling up neural networks is especially important in the large model era, where training such models costly.
It is natural to ask whether other models also scale as a power law, and if so, how their exponents compare to those of neural networks. This comparison is particularly natural for neural kernels. Moreover, since they are nonparametric, their capacity scales automatically as more data are added. From the language of Bahri et al. (2021), dataset scaling for neural kernels is in the resolution- limited regime, which is arguably more interesting scaling regime.
Scaling experiments for neural networks typically cover datasets over many orders of magnitude, which is challenging for kernel regression. In this section, we consider kernel regression for three different neural kernels: fully-connected neural networks, convolution neural networks without pooling operations, and the Myrtle convolutional networks. The kernels in that order increase in performance and computational burden, and all have been well tuned for the CIFAR-10 (Lee et al., 2020).
Of course, CIFAR-10 is limited in size with only 60k examples (50k train, 10k test). To enable a larger scaling study, instead we use the CIFAR-5m dataset (Nakkiran et al., 2021), which consists of samples from a generative model trained on CIFAR-10. Using CIFAR-5m, we can extend our analysis over 2 more orders of magnitude while ensuring the additional data are sampled i.idd. For generalization performance, in addition to evaluation on the i.i.d. CIFAR-5m held out data, we consider the CIFAR-10 and 10.1 test sets. Evaluation on these test sets is standard and allows for comparison against existing results.
In Fig. 2, we observe persistent dataset scaling laws as we vary architecture and evaluation data, For extremely small dataset sizes (10-100), all cases commonly show a lower slope and then transition
2 For example, our 5 million examples kernel computed f1ctt precision, storage on disk alone about 100 terahytes distributed over ~500,000 blocks.
See the next section for more details on this dataset.
5

---

Under review as a conference paper at ICLR 2023
FC3
0V8
Myrte
Classificatio Enror
0.6
0.3
02-
CFAR-10.3 Test (000)
CFAR110 Text 100D)
CEAR5m (ald
10"
10"
15"
Dataset Size
Dataset Size
Dataset Size
FC3
CV8
Myrtle
Mean Meuanruuueed Squared Brror
0.1
0.08
0.06
0.05
0.04
CIFAR-10.T, as-
D...--D... Ry -0.08
CIFAR-10-1 do 0.58 CIFAR 10, aj
CIPAR du -0
OFAR 10, do- 0.09
CIFAN-Sm, 0.09
GFAR-Sm 0.11
GFARSSI Op w0.1
10a
Dataset Size
Dataset Size
Dotaset Size
Figure 2: Dataset size scaling for neural kernels on CIFAR-5m. Using CIFAR-5m as training data, we explore dataset scaling spanning 6 orders of magnitude. Each column corresponds to differ- ent neural architectures: 3-layer fully-connected neural network (FC3), 8-layer convolutional neural network with output vectorization (CV8), and 10-layer Myrtle convolutional neural network with average pooling (Myrtle). Evaluations on three different held-out test sets, CIFAR-5m validation, CIFAR-10 test set, and CIFAR-10.J test set, are shown. Scaling exponents mean squared error are larger (thus faster improvement with more data) for more complex and computationally intensive kernels and for more in in-distribution held out data (CIFAR-5m 5m > CIFAR-10 10 CIFAR-101 10.1)....
Tuned Myrtle
Tiny ImageNet
Tuned Myrtle on Tiny ImageNet
Top-1 Validation Error
0.009
64x64, 6o 32x32, 0.04
0.9
32x32
Validation MSE
0.008
0.8
0.007
0.7
0.006
0.005
0.5
10a
10a
10a
10
10a
10a
10a
10a
10
10a
Dataset Size
Dataset Size
Figure 3: Dataset scaling for Myrtle on Tiny ImageNet. At full resolution (64x64) we achieve best accuracy of 44.7%. Higher resolution inputs result in better classification accuracy and higher scaling exponent(a.).
into a more consistent, higher slope, indicating power-law behavior across 4-5 orders of magnitude. We measure the dataset scaling exponent o&p for the test loss and observe that, as one increases the complexity of the kernel, the scaling exponent increases. Also, for the same architecture, evaluation data that are closer in distribution training data are drawn i.i. from 5m. the order of in-distribution to out-of-distribution for the test sets should be thought of as CIFAR-5m, CIFAR-10, 10, and then CIFAR-10.1..
the training data have higher scaling exponent. Since the
Tiny ImageNet: Towards ImageNet with neural kernels. With our ability to compute mas- sive kernels and use them for inference, ImageNet not far from reach. ImageNet at 1.2 million examples is actually smaller than some datasets considered sere-although augmentation beyond horizontal flips maybe challenging. However, at least two issues remain. First, the large number of output classes. Compared to the 10 classes in CIFAR-10 or CIFAR-5m, one-hot label of Ik classes expands the linear systems to be solved by a factor of 100. Second, the resolution of Ima- *This is because the generative model generate CIFAR-Sm trained the CIFAR-10 training set.

---

Under review as a conference paper at ICLR 2023
FC3
CV8
Tuned Myrtic
Classiiiation Enoor
0.4
0.2
GFFAR
0.22 0.050.1 0.2
0.02 0.050.10.2 2
0.02 0.030.1 0.2
20
Data Augmentation Factor
Data Augmentation Factor
Data Augmentation Factor
Figure 4: Neural kernel performance on CIFAR-10 with data augmentation. Using horizontal flip and RandAug, we measure scaling of neural kernels' classification eπror for 3 different archi- tectures and test sets (CIFAR-10 as well as CIFAR-10.1) Data augmentation factors smaller than 1.0 (light blue shaded region) denotes a subset of original training set, so additional i.i.d. data can be compared with augmented data. The augmentation factor 2x includes horizontal flips, whereas factors larger than 2x use RandAug combined with random crop and flips. While data augmen- tation consistently improves scaling, there a noticeable change in slope around 1.0 for FC3 and Tuned Myrtle kernels (see Appendix C). For augmentation factors of 10x and 20x, the Tuned Myr- tle kernel achieves less than 10% classification error (8.8% with 20x). (See Table for detailed comparison.)
geNet in standard deep learning architecture is x244 if not larger (c.f. EfficientNet Tan & Le (2019; 2021)). For neural kernels with pooling, such as the Myrtle kernel, compute scales quadrati- cally with the total number of pixels. Thus compared to 32x32 CIFAR-10 images, each data point requires 7 2401 times more compute.
As a step toward ImageNet, we consider the Myrtle kernel for the Tiny ImageNet dataset (Le & Yang, 2015). Tiny ImageNet consists of 200 subclasses of ImageNet with 500 images per class, i.e. 100k total training images at a lower resolution of 64x64. We experimented with further downsizing to 32 resolution. Dataset scaling results are shown in Fig. 3. Similar CIFAR-5m scaling, from a sufficiently large dataset size there is consistent power-law scaling. To the best of our knowledge, this is the first evaluation of neural kernels on Tiny Imagenet. We achieve at best classification accuracy of 44.7% on the test set. Comparing to finite networks, among very few references quoting results without data augmentation, Jeevan et al. (2022) obtained 43. % for ResNet- and 45. 53. for Cony 16, indicating our result from neural kernel is competitive to modern finite neural network architectures without data augmentation.
4 DATA AUGMENTATION
Data augmentation is widely used in deep learning applicd to image data. For example, for CIFAR- 10 ImageNet classification, most if not all models are trained using random crops and hori- zontal flips. More effective augmentation strategies such as AutoAug (Cubuk et al., 2019), Ran- dAug (Cubuk et al., 2020) have been found and used in SotA vision models.
The idea of data augmentation has a long history (Niyogi et al., 1998) including in SVMs where they called virtual examples (Schölkopf et al., 1996). However, despite the broad application of data augmentation in deep learning, recently it has been used little in kernel methods due to their limitations large datasets. For example, Li al. (2019) and Shankar al. (2020) used horizontal flips to double the training set size for convolutional kernels. In Lee et al. (2020), an ensemble of kernel predictors, where each training set was augmented randomly, were used means of data augmentation." Because of this, exactly how much of the performance gap between neural networks and kernel methods can be accounted for by data augmentation on image data remains unknown.
In this section, we use our massively parallelized computation of neural kernels and distributed CG solver to explore the effect of data augmentation in neural kernels to a regime far beyond prior work. In our experiments, we focus on the CIFAR-10 10 dataset. We call the ratio between the sizes of the original training set and the training set after augmentation the augmentation factor. Our augmen- tation strategy is to expand the 50k training set by applying various augmentations, Following Li
In effect, this performing block diagonal approximation on the augmented dataset.
7

---

Under review as a conference paper at ICLR 2023
Architecture
Method
CIFAR-10
CIFAR-10.1
CIFAR-5m
DA Ensemble (Lee et al., 2020)
62.4
Kernel
DA CG (20x, this work)
64.1
49.1
FC
CIFAR-5m (2M, this work)
69,9
54.1
78.6
DA, Small LR, no 1.2 (Loce et al., 2020)
65.3
Finite NN
DA. Large LR, L2 (Lee al. 2020)
69.4
Flip (Li et al., 2019)
70.5
DA Ensemble (I.ee et al., 2020)
73.2
Kernel
DA CG (20x, this work)
81.3
66.3
CNN-VEC
CIFAR-5m CG (5M, this work)
83.4
71.9
86.5
Small LR, 1.2 (l.ee al., 2020)
83.9
Finite NN
DA. Large LR, L2 (Lec
2020)
85.6
CNN-LAP Flip et al.. 2019)
82.2
CNN-GAP DA Ensemble (leee al., 2020)
84.8
Kemel
Syrtle0--Gaussing Flip (Shankar al., 2020)
89.8
78.3
Tuned Myrtlc DA CG (20x. this work)
91.2
79,4
CNN-Pool
Myrtle10 CIFAR-5m CG (1.6M, this work)
89.1
79.1
89.5
ResNet18 CIFAR-5m (Nakkiran al., 2021)
89.0
89.4
CNN-GAP DA, Small LR, no L2 (Lee et al., 2020)
84.4
Finite NN
CNN-GAP DA Large LR, L2 (Lee et al., 2020)
86.7
Myrtle10 DA (Shankar et al., 2020)
96.0
89.8
Table I: CIFAR-10 test accuracy for kernels and finite neural networks of the corresponding archi- tecture type.
al. (2019); Shankar et al. (2020), 2x augmentation factor is given by horizontally flipping all images. Beyond that, up to a 20x augmentation factor, we used RandAug (Cubuk et al., 2020) used in Shankar al. (2020). In particular for each image the training set, we sample one (N = 1) random augmentation op among ['FlipLR', 'Solarize', 'Color', 'Brightness', 'Contrast', Sharppees', 'Posterize', 'Equalize', 'Identity I and apply with magnitude M 2. After applying random augmentation, we apply random crop of 4 pixels well as a random flip. We do not use cutout augmentation. With the strategy, we are able to achieve an accuracy of 91.2% (see Fig. 4). To our knowledge this is the highest published accuracy for kernel method. Note that while the gap accuracy between kernels and finite neural networks is reduced, still remains (see Table 1).
SEQUENCE AND GRAPH DATA
In this section, we continue to develop neural kernels as method in their own right. We consider their performance on structured data modalities other than images. Kernels are particularly exciting for such data modalities because we lack inductive biases to take advantage of, and kernel methods such as Gaussian processes are the standard approach for guided experimental design.
First, consider a protein function prediction benchmark motivated by protein design for targeted gene therapy. There is significant interest in engineering Adeno-associaae associated virus (AAV) capsid proteins to cause the virus to integrate specific DNA into target cell (Bryant al., 2021). We perform kernel regression to predict the fitness of mutations across a 28 amino acid window of these proteins, by training and evaluating across all splits, ranging from 82k to 284k examples, of a dataset assaying VP-1 AAV protein variants (Dallago et al., 2021).
Second, we consider small molecule classification task involving detecting drug cardiac toxic- ity (Han et al., 2021). This dataset contains descriptions of the graph-structured molecules as SMILES strings along with assays for toxicity. To model this graph structured data, we derive kernels for deep graph neural networks, and fit these to the training set of 6,523 examples.
Developing specific kernels for such data is critical, since traditional covariance functions do not capture task-relevant similarities in digensinnnsioal inputs with complex structure. The connec- tion between neural networks and neural kernels the infnitte-wdth limit provides one promising direction to exploit the progress in neural network architecture design. For example, convolutional
"Specified in http://gghhaa..ommmmmmsyyyssehhh kemelss codehhiab/master/ennniggg
8

---

Under review as a conference paper at ICLR 2023
Method
Mut-Des
Des-Mut
1-vs-rest
2-vs-rest
7-vs-rest
low-vs-high
Best reported
0.79
0.75
0.48
0.74
0.74
0.39
Ours
0.81
0.84
0.64
0.71
0.71
0.34
Table 2: Spearman correlation for kernel regression compared to the best result reported for several methods in (Dallago et al., 2021)
Model
baseline
GP
SNGP
SNGP Ensemble
Ours
TEST-IID
0.919
0.937
0.932
0.942
0.943
TEST-OODI
0.786
0.823
0.836
0.850
0.800
TEST-OOD2
0.831
0.873
0.885
0.896
0.750
Table 3: AUROC for kernel regression compared to other methods reported in (Han et al., 2021).
operations common in neural networks used protein sequence data. Similarly, for graph data like molecules, graph convolution operations have grown in popularity. Any such operation has a corresponding kernel that can be used in approach to kernel regression.
In this vein, for the AAV dataset we considered deep. ID convolutional architectures and for car- dotoxicity considered deep, graph convolutional architectures. The remaining hyperparamters are depth, weight and bias initialization variance, activation function, diagonal regularization, the type of pooling function, filter size, and whether use the NNGP kernel or the NTK. To tune hy- perparameters for the cardiotox dataset, we used Ik trials on the Google Vizier service (Golovin al., 2017). Each trial used all examples from the training set with AUROC on the validation set as the objective. For the AAV dataset, tuning hyperparameters using the whole training set is not feasible due its size, since the kernel must be recomputed each time. Instead we used only 2k examples from the training set with MSE on 1 validation set examples as the objective. Again Ik trials of the Vizier service were used. For additional details on hyperparameters and their tuning, see Appendix C.
After selecting hyperparamters, we ran inference on all splits of the cardiotox and AAV datasets (see Tables. 2 and 3). In all cases, we find method is either better competitive with prior methods across all splits.
CONCLUSIONS
In this work we probe the performance kernel regression across two dimensions of scale, in the number of examples and the complexity of the kernel. We distributed the computation of large neural kernels with up to 5 million examples and then used distributed conjugate gradients solvers to solve the corresponding linear systems. In doing so, we not only set a new SotA for kernel methods on CIFAR-10 10 but also outperform highly tuned deep learning methods on a protein design problem and graph-structured molecule toxicity screening problem. Computational considerations aside, we see that these neural kernels scale remarkably well in terms of predictive performance. Extrapolating from our scaling curves, it appears that significant improvements could still be made by adding more data. These infinitely- advantages over the parametric, SGD trained alternatives. They offer stable solutions with relatively little optimization and hyperparameter tuning. In the context of Gaussian processes, they represent a distribution over deep networks, with corresponding uncertainty that can used across decision making tasks such experimental design. Indeed, our SotA results on protein function prediction for targeted gene therapy and molecule toxicity prediction suggest that these kernels may be very well suited to exciting structured design problems. While computing quantities such the marginal likelihood or test variances involves additional challenges, one can readily draw samples from the corresponding GP to be used, for example, with Thompson sampling approaches.
instantiations of deep neural networks offer some exciting properties and
The computational cost of computing the neural kernels and solving the corresponding linear sys- tems remains a limitation to scaling further. While there significant literature on approximating the kernels with regard to their scaling in the number of examples, approximations of the pairwise kernel computation itself similar Zandieh et (2021); Han al. (2022) seems necessary.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Ben Adlam, Jachoon Lee, Lechao Xiao, Jeffrey Pennington, and Jasper Snoek. Exploring the un- certainty properties of neural networks' implicit priors in the infinite-width limit. arXiv preprint arXiv:2010.07:55, 2020.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances Neural Information Processing Systems, pp. 8141-8150. Curran Associates, Inc., 2019.
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXi::2102.06700 2021.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
// //github. com/ google/jar
Drew Bryant, Ali Bashir, Sam Sinai, Nina Jain, Pierce Ogden, Patrick Riley, George Church, Lucy Colwell, and Eric Kelsic, Deep diversification of an aav capsid protein by machine learning. Nature Biotechnology, 39, 06 2021.
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 113-123, 2019.
Ekin Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops, pp. 702-703, 2020.
Christian Dallago, Jody Mou, Kadina E Johnston, Bruce Wittmann, Nick Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K Yang. FLIP: Benchmark tasks in fitness landscape inference for proteins. In Neural Information Processing Systems Datasets and Benchmarks Track, 2021.
Simon Du, Kangcheng Hou, Russ R Salakhutdinov, Barnabas Poczos, Ruosong Wang, and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. Advances in neural information processing systems, 32, 2019.
Adrià Garriga-Alonso, Laurence Aitchison, and Carl Edward Rasmussen. Deep convolutional net- works shallow gaussian processes. In International Conference on Learning Representations, 2019.
Daniel Golovin, Benjamin Solnik, Subhodeep Moitra, Greg Kochanski, John Karro, and David Scul- ley. Google vizier: A service for black-box optimization. In Proceedings of the 23rd ACM SIGKDD international conference on knowledge discovery and data mining. pp. 1487-1495, 2017.
Insu Han, Amir Zandich, Jachoon Lee, Roman Novak, Lechao Xiao, and Amin Karbasi. Fast neural keree embeddings for general activations. arXiv preprint ArXiv:2209.04133 2022.
Kehang Han, Balaji Lakshminarayanann and Jeremiah Liu. Reliable graph neural networks for drug discovery under distributional shift. November 2021.
James Hensman, Nicolò Fusi, and Neil D. Lawrence. Gaussian processes for big data. In Uncertainty in Artificial Intelligence, 2013.
Joel Hestness, Sharan Narang, Newsha Ardalani. Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Patwary, Mostofa Ali, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically. arXiv preprint 17Xiv:1712.004499 2017.
Jiri Hron. Yasaman Bahri, Jascha Sohl-Dickstein and Roman Novak. Infinite attention: NNGP and
NTK for deep attention networks. In International Conference Machine Learning, 2020.
10

---

Under review as a conference paper at ICLR 2023
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kemel: Convergence and gen- eralization in neural networks. In Advances in Neural Information Processing Systems, 2018.
Pranav Jeevan al. Convolutional xformers for vision. arXiv preprint arXiv:2201.102771 2022.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08366 2020.
Ya Le and Xuan Yang, Tiny imagenet visual recognition challenge. 2015.
Jachoon Lee, Yasaman Bahri, Roman Novak, Sam Schoenholz, Jeffrey Pemnington, and Jascha Sohl- dickstein. Deep neural networks gaussian processes. In International Conference Learning Representations, 2018.
Jachoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Advances in Neural Information Processing Systems, 33:15156 13115156-1517,, 2020.
Zhiyuan Li, Ruosong Wang, Dingli Yu, Simon S Du, Wei Hu, Ruslan Salakhutdinov, and Sanjeev Arora. Enhanced convolutional neural tangent kernels. arXiv preprint arXiv:1911.00809 2019.
Siyuan Ma and Mikhail Belkin. Diving into the shallows: a computational perspective on large-scale shallow learning. Advances in neural information processing systems, 30, 2017.
Siyuan Ma and Mikhail Belkin. Kernel machines that adapt to gpus for effective large batch training. Proceedings of Machine Learning and Systems, 1:360-373, 2019.
Wesley J. Maddox, Andres Potapcynski, and Andrew Gordon Wilson. Low-precision arithmetic for fast gaussian processes. In Proceedings of the "hirtt-Fighth Conference on Uncertainty in Artificial Intelligence, 2022.
Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahra- mani. Gaussian process behaviour in wide deep neural networks. In International Conference Learning Representations, 2018.
Giacomo Meanti, Luigi Carratino, Lorenzo Rosasco, and Alessandro Rudi. Kernel methods through the roof: handling billions of points efficiently. Advances in Neural Information Processing Systems, 33:14410-14422 2020.
Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. In International Conference on Learning Represen- tations. OpenReview.net, 2021.
Radford M. Neal. Priors for infinite networks (tech, rep. crg-tr-94-1). University of Toronto, 1994.
Timothy Nguyen, Zhourong Chen, and Jachoon Lee. Dataset meta-learning from kernel ridge- regression. In International Conference on Learning Representations, 2021a.
Timothy Nguyen, Roman Novak, Lechao Xiao, and Jachoon Lee. Dataset distillation with infinitely wide convolutional networks. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021b.
Partha Niyogi, Federico Girosi, and Tomaso Poggio. Incorporating prior information in machine learning by creating virtual examples. Proceedings of the IEEE, 86(1 11:21)):220011 1998.
Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang. Jiri Hron, Daniel A. Abo- lafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are gaussian processes. In International Conference on Learning Representations, 2019.
11

---

Under review as a conference paper at ICLR 2023
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 Learning Representations, 2020. URL /titt//gghhhbb// com/ googl e e/ neural Stauural-tnggntt...
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Advances in Neural Information Processing Systems, 2007.
Jonathan Rosenfeld. Scaling laws for deep learning. arXiv preprint arXiv:2108.07686, 2021.
Jonathan Rosenfeld, Amir Rosenfeld, Yonatan Belinkov, and Nir Shavit. A constructive prediction of the generalization error across scales. arXiv preprint arXiv: /26..,2667... 2019.
Alessandro Rudi, Luigi Carratino, and Lorenzo Rosasco. Falkon: An optimal large scale kernel
method. In Advances in Neural Information Processing Systems, 2017.
Bernhard Schölkopf, Chris Burges, and Vladimir Vapnik. Incorporating invariances in support vector learning machines. In International conference on artificial neural networks, pp. 47-52. Springer, 1996.
Vaishaal Shankar, Alex Fang, Wenshuo Guo, Sara Fridovich- Keil, Jonathan Selaa-KKllly,, Ludwig Schmidt, and Benjamin Recht. Neural kernels without tangents. International Conference on Machine Learning. 2020.
Jonathan Richard Shewchuk et al. An introduction the conjugate gradient method without the agonizing pain, 1994.
Samuel Stanton, Wesley Maddox, lan Delbridge, and Andrew Gordon Wilson. Kernel interpolation for scalable online gaussian processes. In International Conference on Artificial Intelligence and Statistics, 2021.
Mingxing Tan and Quoe Le. Efficientnet: Rethinking model scaling for convolutional neural net- works. In International conference on machine learning, pp. 6105-6114. PMLR, 2019.
Mingxing Tan and Quoc Le. Efficientnetv2: Smaller models and faster training. In International Conference on Machine Learning. pp. 10096-10106. PMLR, 2021.
Tresp. A Bayesian committee machine. Neural Computation, 12(11)2711,2741, 2000.
Ke Wang, Geoff Pleiss, Jacob Gardner, Stephen Tyree, Kilian Q Weinberger, and Andrew Gordon Wilson. Exact gaussian processes million data points. In Advances in Neural Information Processing Systems, 2019.
Christopher Williams. Computing with infinite networks. In Advances in Neural Information Pro- cessing Systems, 1996.
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning. volume 2. MIT press Cambridge, MA, 2006.
Andrew Gordon Wilson and Hannes Nickisch. Kernel interpolation for scalable structured gaus- sian processes (kiss-gp). In International Conference on International Conference on Machine Learning, 2015.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760, 2019.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006. 14044488..4448, 2020.
Amir Zandieh, Insu Han, Haim Avron, Neta Shoham, Chaewon Kim, and Jinwoo Shin. Scaling neural tangent kernels via sketching and random features. Advances Neural Information Pro- cessing Systems, 34:1062-1P 1073, 2021.
12

---

Under review as a conference paper at ICLR 2023
A NEURAL KERNEL ARCHITECTURE
Here, we provide details of neural network architectures use in the experiments.
A.I
ARCHITECTURE DEFINITION
For neural network architecture for our kernel used in the experiments, we provide layer specification code using Neural Tangent (Novak et al., 2020) library.
Fully-Connected Network (FC): We use vanilla 3-hidden fully-connected layers with ReLU ac- tivation function. We denote this neural architecture FC3.
from neural taaraltt import stax
def FC3 ():
depth
W_std = 2.0xx0.5
b_std 0.1
num classes 10
dense functools partial [)]
@_std=W_std, b_std-b_stad std)
layers =
for in range (depth)
layers + [dense (512), stax. Relu() ()
layers += [dense (num_classes))
return stax.seriae (+ (*layers)
Convolutional Neural Network with Vectorization aggregation (CNN-VEC): We us vanilla 8-hidden convolutional layers with ReL.U activation function. After the convolutional layers, aggre- gation is done by flattening (or vectorization). We denote this neural architecture as CV8.
from poural_tangente import stax
def CV8 ():
depth = 8
W_std 2.0*+0.5
b_std = 0.1
num cim_clasess
conv = functools. sunctools.partial (stax. Conv, M_std-W_stdh b_std-b_std,
padding='SAMS
layers
for in range (depth)
layers [conv (512, (3, 3)), Relx.
layers += [stax. [stax.Flatten() (),
stax. Dense (num_classes, R_std b_std-b_stdl ]
return stax. serial (* (+layers)
Myrtle Convolutional Neural Network (Myrtle): The neural kereel for Myrtle CNN was first introduced by Shankar et al. (2020). Using intermediate average pooling layers, the convolutional architecture is more efficient than global average pooling while achieving high performance with 10-layers.
from @oural_tangent import stax
def Myrtle() ():
W. std 2.0++0.5
b_std = 0.1
classes 10
13

---

Under review as a conference paper at ICLR 2023
activation fn stax.Relus
conv = functools, partial (stax. Conv, W_std=W_std, b_std=b_sta, Badding='SAMB1 layers = [conv (512, (3, 3)), activation_fn]
layers [
conv (512, (3, 3)), activation in] 2 layers [stax.AvgPook AvgPool ((2, 2), strides (2, 2))
layers +=
conv (512, (3, 3)), activation_fn,
layers + Estax.AvgPer ((2, 2), strides= (2, 2))
layers +=
conv (512, (3, 3)), activation_fi layers [stax. Gtoxx.Globalaggool.. ())
layers +-
stax. sta.. Flatten (),
Scax.Denser
num_classes, M_std, b_std)
]
return stax.serial (*layers)
Tuned Myrtle Convolutional Neural Network (Tuned Myrtle)
from neural Laaral_tangent import stax
def TunedMyrtle
W_std 2.0
b_std = 0.01
num_classes 10
activation fn stax. stax. ExpNormallzed
norm_layen stax. (ayxitlormmxxxiss 2, 3))
conv functools partial (stax. W_std-W_std, b_std=b_std,
padding='SAME')
layers [conv (512, (2, 2)), norm_layer, activation_fnj
layers [conv (512, (3, 3)), nom__ayyr,actiiatioo..
2
layers [stax. Avcaol ( (2,, 2), strides= 2))
layers = [conv (3, 3)), norm_layer, activation_fn)
layers
AtExxAvgFool1 ( (2, 2), strides= 2))
layers = [conv (512, (3, 3)), norm_layer, activation_fn)
layers += Estax.GlobalåvgPod (), stax. )))
for in range
layers += 1stax.Denss (512, W_std, b b_std),
stax. Ltye.Nay) activation fn]
layers += I Dense num_classes, w_std, b_std)
stax. stax.Dense(
return stax.serian (* layers)
Simple Graph Neural Network Our simple graph neural network is based off Graph NTK (Du et al,, 2019) where the aggregation node combines node features according to the graph connectivity.
from neural ttargett import stax
def (GNnp
14

---

Under review as a conference paper at ICLR 2023
depth 25
W_std = 0.6
b std 0.1
#_std_out = 0.1
b_std_eus 0.03
nun_classes 2
activation_fn stax.Erf()
dense = functools.partion
stax.Dense, #_std=H_std, b_std=b_std)
aggregate stax Aggregate (aggregate_axis-1) batch_axis=0, channel_axis-2,
implementatioon (DENSE')
layers []
for in range (depth)
layers +- [dense (512), activation_fi fn, aggregate]
layers layers.append (stax. www.allvgooooll.ooooooo.
layers.
append aayers.append
Iesta..Dense (num clamselasses,
W_tt b
return stax. stax.serial (ataaalalaaalla (+layers)
ID Convolution Neural Network
from neural poural_tangente import stax
def (neiimC
depth = 5
W_std 1.7
b_std = 0.5
W_std out 1.0
b_std_ out 0.
dimension_numbera None
filter_size 15
filter_tuple (Filter_size,)
use_gap True
padding = 'AMME' parameterization 'ntk'
activation_fn 'erf'
collect_layers []
conv_or_len functools, .unctools.partial
Stav, W_std-W_std, _dd padding=padding,
if depth 0:
collect_layers += [
conv len (512, filter_tuple dimension htw.:/ouuuuueerddiiieeiesoooeeee activation_fn
] for in range (1, depth) collect_layers += [conv_or_len (width, filter_tuple), activation_fni
collect_layers Gltax. ()]
collect_layers
stax.Flatten () stax.Densee (1, o_st b_std_out, parameterizal wwt.:rraaaaaaaaamteerr.o..ooo
return stax.serial (+collect_layers
15

---

Under review as a conference paper at ICLR 2023
A.2 REGULARIZED ZCA PREPROCESSING
For convolutional kernels on image data, (e.g- CV8, Myrtle, Tuned Myrtle), we preprocess inputs with regularized 7CA whitening. Following the convention of Lee et al. (2020), for flat- tened d-dimensional, n training inputs X E Rkdxn, the input data covariance is Σx 1 Xxx. Considering the SVD of Σx UDUT, the regularized ZCA whitening transform is given
(V "r( 1
Wica WZCA=U -U
(D+Edd
UT,
where is the regularization parameter.
has been demonstrated that regularzied ZCA significantly helps performance of convolutional
neural kernels on image classification tasks (Shankar et al., 2020; Lee et al., 2020; Nguyen et al., 2021b). For CV8 we use = 3.0, and for Myrtle le and Tuned Myrtle we use = 0.1.
B SEQUENCE DATA PREPROCESSING
The AAV data consists sequences amino acids of different lengths (around 700). The sequences were one-hot encoded and padded with special character so they were all the same length as the longest sequence. Finally the features and labels were normalized to have mean and standard deviation 1.
C ARCHITECTURE HYPERPARAMETER TUNING
C.I TUNING GRAPH CONVOLUTION ARCHITECTURE CARDIOTOX
The activation function tuned among ReLU, Erf. Identity, GeLU, Normalized Exponential, and Sigmoid. The kernel type was tuned among the NNGP kernel and the NTK. The weight (except the final layer) initialization standard deviation initialization standard deviation was tuned among {0., 0.001, 0.003, 0.01, 0.03, 0.1,0.3) The last- layer weight initialization standard deviation was tuned in [0.1, 10.0) on logarithmic scale. The depth was tuned among 21, 5, 7,9, 10, 12, 15, 20, 25, 30, 50}. Whether to include self-loops was also tuned. The diagonal regularization term was tuned among
tuned in [0.1, 2.5] on a linear scale. The bias
import numpy as np concatenate ( [[[.], np. 1ogspace (-10, 4, 49)
C.2 TUNING ID CONVOLUTION ARCHITECTURE ON AAV
The activation function was tuned among ReLU, Erf, Identity, GeL.u, and Normalized Exponential. The kernel type was tuned among the NNGP kernel and the NTK. The weight (except the final layer) initialization standard deviation was tuned in (0.1,2.5) on linear scale. The bias initialization standard deviation was tuned among (0., 0.01, 0.05, 0.1, 0.5). The last-layer weight initialization standard deviation was tuned in 10.0) on logarithmic scale. The depth was tuned among (1,2,3,5, 7, 9}. The diagonal regularization term was tuned among
import numpy as np np. concatenate [[]O], np. log.pacepcceeccc (-10, 49)))
As mentioned in the main text, only small subset of the training data was used for tuning. To investigate the sensitivity of hyperparameters to dataset size, we ran inference for several small dataset sizes and check the correlation of the objective (see Fig. 5). We saw good correlation (0.918) and rank agreement (Spearman's 's rank 0.700) for test MSE when 100 vs 2k training examples were used. This suggests the choosing hyperparamters based on a smaller subset of the training set reasonable.
C.3 TUNING MYRTLE ARCHITECTURE ON CIFAR-10
The activation function was tuned among ReLU, Erf, Identity, GeLU, Normalized Exponential and its variants. Whether to include layer normalization after affine layer and whether use
16

---

Under review as a conference paper at ICLR 2023
Test MSE
0.17
0.16
2k train set
0.15
0.14
0.255 0.260 0.265 0.270
0.275 0.280 0.285 0.290 0.295 0.300
100 train
Figure 5: A scatter plot of test MSE. Each point is a different assignment of the hyperparameters whose coordinates are the test MSE using 100 and 2k examples train respectively.
global average pooling after the convolution layer was tuned. The weight initialization standard linear scale. The bias initialization standard deviation was 0.01,0.03, 0.033, 0.3). The first convolution's filter size was tuned
deviation was tuned in [0.1, 2.5]
tuned among {0., 0.001, 0.003, 0.01, among 21,2, 3, 8, 10} and the remaining convolutional layers' filter size tuned among {1, 2, 4, 5, 8). The number of convolutional layers before each (2, 2) average pooling layer was tuned among {0, 1, 4,5.. After convolutional layers, option add {0, 2} dense layers was tuned. ZCA regularization strength was tuned in [0.003, 10.0] on a logarithmic scale. As described in the tuning of the AAV dataset, small subset of 1,280 training examples were used for tuning.
Among these, the most robust and high- seghoperformin setting was selected for the Tuned Myrtle archi- tecture used for CIFAR-10 data augmentation and Tiny ImageNet experiments.
17


Reviews:
[{"review_id": "cmoxR1dmAor", "reviewer": "Reviewer_c8Ru", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper explores further scaling up preconditioned conjugate gradient (PC-CG) based GP regression by parallelizing the computation of the Gram matrix as a one-off and storing it to disk in \"sub-blocks\" to be loaded in the RAM in a multi-threaded fashion as/when needed. \n\nThis contrasts with standard PC-CG where the Gram matrix is never computed entirely but rather partially recomputed on the fly as needed. \n\nThe standard PC-CG approach works reasonably well when the kernel is easy to compute (e.g. the square-exponential kernel), but breaks down when the kernel is more complex. \n\nThe authors demonstrate the effectiveness of the proposed approach using neural kernels, especially compared to alternatives consisting of approximating the linear system (as opposed to solving it within an error CG-style), and propose additional empirical studies such as how performance of GP regression under neural kernels scale with data sizes.\n", "summary_of_the_review": "Overall the main contribution is too simple to be worthy of a paper and, in any case, empirical evaluation should be strengthened.", "strength_and_weaknesses": "- **Strength**: Scalability is one of the major hurdles in bridging the gap between  GP methods and deep learning. This paper's attempt to doing so is laudable, and experimental results are encouraging.\n\n- **Weaknesses**: Experimental results seem to be missing error bars/intervals, hardware specs, and running durations. Overall the contribution seems too light to be worthy of publication.\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "_mrgjG5X9Z8", "reviewer": "Reviewer_wb1p", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper empirically investigates the performance of million-scale kernel regression with infinite-width neural networks using the preconditioned conjugate gradient method from the Gaussian process literature (Wang, 2019). The method is applied to several domains including testing scaling laws on CIFAR-5m, enabling data augmentation on CIFAR-10 by a factor of 20, and sequence/molecule classification tasks using convolutional/graph convolutional architectures.", "summary_of_the_review": "Given the lack of central message and limited amount of contribution compared to prior literature, I would recommend rejection. ", "strength_and_weaknesses": "## Strengths\n\n* The paper makes a nice attempt to scaling infinite-width neural-network kernels. This is an important problem faced by the area and a solution would be of interest to a wide audience. \n\n* The empirical study is comprehensive and covers several data modalities including images, sequences and graphs, all showing competitive performance against strong baselines.\n\n## Weaknesses\n\n* The results are good, but not surprising. Given the existing work of Wang et al. (2019) that already applies GPs to 1 million data points, the scaling in this work was not that far from reach. In fact, the paper is using the same preconditioning as Wang et al. (2019) and from the text I did not see any additional innovations needed to achieve the 5 million scale. Achieving 91% on CIFAR 10 with data augmentation that increases the training set by a factor of 20 does not seem surprising, either. \n\n* One big question i had in mind while reading this paper is---what is the central message? The focus in each section seems completely orthogonal to each other and the whole piece of work feels more like an experiment report. Simply showing that infinite-neural-network kernels achieves predictable good (but not state-of-the-art) results does not qualify a scientific contribution. \n\n* Some details are unclear about the 5 million experiment. What are the tolerance level needed to solve this linear system? One big concern I had about Wang et al. (2019) is that they used a very large tolerance (1 if I remember it correctly) during training and did not bother to get the GP hyperparameters right (partly the reason why the training is super fast/scalable). During test time, to ensure good performance, they had to switch to a much smaller tolerance level. Do you have to do the same? What is total training + test time when compared to a standard sparse method like Nystrom? (Nystrom-like ideas  is not totally infeasible for these kernels, see Deng, Z., Shi, J., & Zhu, J. (2022). NeuralEF: Deconstructing Kernels by Deep Neural Networks. arXiv preprint arXiv:2205.00165 \nwhere they show a nice recovery of NTK using sparse eigenfunctions. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "P-vNKuwkeCa", "reviewer": "Reviewer_nR21", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors consider the large-scale implementation of 'neural kernels', kernel\nmethods that are derived from the architecture of a specific neural network and\nthe correspondence between randomly-initialized neural networks and gaussian\nprocesses. In contrast to existing large-scale kernel methods for use with\nkernels like the RBF kernel, the fact that these neural kernels are\ncomputationally-expensive to compute necessitates additional parallelization.\nThe authors review methods for large-scale kernel regression, implement their\nown method in the context of neural kernels, and show experiments on various\nlarge datasets (CIFAR-5m, Tiny ImageNet) where these kernels achieve\nperformance beyond past kernel methods. They also demonstrate their approach on\nvarious basic science tasks, where it outperforms existing kernel methods.\n\n", "summary_of_the_review": "\nThe engineering contribution and its consequences are significant, and the\nauthors verify these well through various experiments. One hopes the code will\nbe released.\n", "strength_and_weaknesses": "## Strengths\n\n- The engineering contribution directly enables new evaluation of neural\n  kernels at novel scales, enabling their use on new datasets where they set\n  new record for performance by kernel methods.\n- The paper is well-written and does a good job of organizing various\n  approaches to large-scale kernel regression in section 1.2 and section 2. One\n  comes away with an excellent understanding of the landscape of approaches to\n  this problem, which should be useful for future work.\n- The authors do a good job of showing how the engineering advancements open\n  new opportunities for kernel methods. They demonstrate that neural kernels\n  achieve excellent performance on molecular prediction/classification tasks\n  compared to existing methods, and show that data augmentation can give good\n  improvements in performance for kernel methods, given the ability to work\n  with larger kernels.\n\n## Weaknesses\n\n- The precise methodological contributions seemed somewhat unclear to me -- is\n  the main technical contribution to identify the most relevant existing\n  large-scale kernel regression approaches for working with neural kernels (as\n  outlined in section 2), and then to develop and implement a new parallelization\n  scheme on top of these methods to facilitate computing neural kernels? The\n  latter feels like a somewhat limited (but nonetheless significant) technical\n  contribution, given that the paper's description of the parallelization scheme\n  seems to be mostly restricted to showing Figure 1.\n- It might be better to have a demonstration of the method on low-resolution\n  datasets other than CIFAR-5m, since this dataset is essentially synthetic. I\n  think the authors' experiments on Tiny ImageNet are more than appropriate in\n  this connection, I just point this out because the description of CIFAR-5m at\n  the bottom of page 5 as \"...ensuring that the additional data are sampled\n  i.i.d.\" is not correct and somewhat misleading (since the dataset is\n  generated from a trained generative model).\n- I think the section about neural scaling laws could be improved\n  by a more precise discussion of the issues -- once one is working with linear\n  models (whether or not the kernel is \"neural\"), one can make much more\n  precise statements about the asymptotic scaling of the test error as a\n  function of the number of samples given various distributional assumptions,\n  and it does not seem to be appropriate to me to treat this as a purely\n  empirical endeavor.\n  These types of questions have been studied in the literature on nonparametric\n  regression in RKHSes; under simple distributional models for the data (e.g.,\n  the target function is sufficiently smooth and the data come from a\n  distribution on a $d$-dimensional manifold), the RKHS associated to a given\n  kernel can be characterized in terms of differentiability properties of the\n  kernel (which are, in turn, inherited from the architectural choices and the\n  random initialization scheme), and the corresponding scaling behavior (i.e.,\n  the exponent) of the test error becomes a function of this RKHS, the\n  dimension $d$, and the target function's smoothness.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "8AgYwM3yYO", "reviewer": "Reviewer_7cjF", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work enables kernel regression with the so-called neural tangent kernel to work with significantly larger datasets compared to before (order of million datapoints) by building on top of the framework of [1] for Gaussian processes. The key difference to [1] stems from the difference of computational costs to obtain the kernel matrix; while for GPs, simpler kernels such as the RBF are often used, calculating the Gram matrix of the NTK associated with a convolutional neural network marks the bottleneck, instead of the linear system that needs to be solved. The authors distribute the computation of the kernel over many devices and store it on disk. To perform the matrix-vector product needed in conjugate gradient descent, the authors use a smart way to read chunks of the kernel efficiently, avoiding thus the need to recompute the kernel at every step. The more powerful method enables experiments that were not possible before, such as large scaling law experiments, experiments on TinyImageNet and heavy data augmentation leading to SOTA kernel performance on CIFAR10. Finally, the authors also show that these kernels can reach very strong performance on smaller datasets including molecular tasks.", "summary_of_the_review": "In summary, I would like to see this work accepted at ICLR. The framework developed in this work enables important experiments with NTK that were infeasible before. Scaling laws, usage of data augmentation and the possibility to scale to larger datasets such as TinyImageNet facilitate comparing finite and infinite width networks, which still remains an open problem.", "strength_and_weaknesses": "**Strengths**:\n1. The paper is very well-written and very easy to follow. I really enjoyed reading this paper.\n2. The resulting methodology in my opinion enables a lot of interesting experiments (provided enough compute) that were simply infeasible before. The obtained results provide a better and more fair comparison between finite and infinite networks, showing for instance that when standard training components such as data augmentation are included for kernels, very competitive accuracies can be achieved with the kernel method as well. It also marks a first step towards evaluating such kernels on ImageNet, which further facilitates comparisons.\n3. The experiments regarding scaling laws are very interesting and important as they provide a more fine-grained look into performance. To my knowledge, such experiments are missing in the literature. It would have been even more interesting to compare the kernel scaling laws with the corresponding finite network.\n\n**Weaknesses**\n1. The methodology largely relies on previous work (i.e. [1]) with the only difference that the kernel is stored on disk instead of computed on the fly. While the smart reading strategy of the chunks is certainly not trivial, the technical novelty of the approach is somewhat limited.\n2. While the protein/molecule experiments are very interesting in their own right, I don’t see how they fit into this paper. Most datasets are only moderately large and the involved methodology developed in this paper is thus not needed to perform accurate inference. \n\n[1] Ke Wang, Geoff Pleiss, Jacob Gardner, Stephen Tyree, Kilian Q Weinberger, and Andrew Gordon Wilson. Exact gaussian processes on a million data points. In Advances in Neural Information Processing Systems, 2019. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject