Random Features for the Neural Tangent Kernel

TL;DR

This paper introduces an efficient method for approximating the Neural Tangent Kernel using random features and sketching, enabling large-scale neural network training with comparable accuracy and significantly reduced computation.

Contribution

It proposes a novel feature map construction for NTK that is scalable to large datasets, combining random features with sketching and leverage score sampling.

Findings

Achieves comparable error bounds with smaller feature dimensions.

Runs tens of times faster than exact kernel methods on large datasets.

Demonstrates superior performance on various machine learning tasks.

Abstract

The Neural Tangent Kernel (NTK) has discovered connections between deep neural networks and kernel methods with insights of optimization and generalization. Motivated by this, recent works report that NTK can achieve better performances compared to training neural networks on small-scale datasets. However, results under large-scale settings are hardly studied due to the computational limitation of kernel methods. In this work, we propose an efficient feature map construction of the NTK of fully-connected ReLU network which enables us to apply it to large-scale datasets. We combine random features of the arc-cosine kernels with a sketching-based algorithm which can run in linear with respect to both the number of data points and input dimension. We show that dimension of the resulting features is much smaller than other baseline feature map constructions to achieve comparable error bounds both in theory and practice. We additionally utilize the leverage score based sampling for improved bounds of arc-cosine random features and prove a spectral approximation guarantee of the proposed feature map to the NTK matrix of two-layer neural network. We benchmark a variety of machine learning tasks to demonstrate the superiority of the proposed scheme. In particular, our algorithm can run tens of magnitude faster than the exact kernel methods for large-scale settings without performance loss.