Kernel Regression with Infinite-Width Neural Networks on Millions of Examples

TL;DR

This paper demonstrates large-scale kernel regression with neural kernels on millions of examples, achieving state-of-the-art results and exploring diverse data modalities through distributed computation and advanced algorithms.

Contribution

It introduces a massively parallelized computation method for neural kernels enabling large-scale kernel regression on millions of data points.

Findings

Achieved 91.2% accuracy on CIFAR-10 with kernel methods.

Scalable kernel regression on datasets with up to five million examples.

Competitive results on protein and small molecule prediction tasks.

Abstract

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. In this work, 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 five 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). Moreover, we explore neural kernels on other data modalities, obtaining results on protein and small molecule prediction tasks that are competitive with SotA methods.