Efficient kernel surrogates for neural network-based regression

TL;DR

This paper investigates the use of efficient kernel surrogates, specifically the Conjugate Kernel (CK), as practical approximations to Neural Tangent Kernels (NTKs) for improving the performance and understanding of neural network-based regression and classification tasks.

Contribution

It provides a theoretical analysis of CK performance relative to NTK, demonstrating near-equivalent results and offering a practical framework for enhancing neural network accuracy inexpensively.

Findings

CK performs nearly as well as NTK in regression and classification.

The regularity of the kernel significantly influences performance.

The approach improves GPT-2 classification and physics-informed operator training.

Abstract

Despite their immense promise in performing a variety of learning tasks, a theoretical understanding of the limitations of Deep Neural Networks (DNNs) has so far eluded practitioners. This is partly due to the inability to determine the closed forms of the learned functions, making it harder to study their generalization properties on unseen datasets. Recent work has shown that randomly initialized DNNs in the infinite width limit converge to kernel machines relying on a Neural Tangent Kernel (NTK) with known closed form. These results suggest, and experimental evidence corroborates, that empirical kernel machines can also act as surrogates for finite width DNNs. The high computational cost of assembling the full NTK, however, makes this approach infeasible in practice, motivating the need for low-cost approximations. In the current work, we study the performance of the Conjugate Kernel (CK), an efficient approximation to the NTK that has been observed to yield fairly similar results. For the regression problem of smooth functions and logistic regression classification, we show that the CK performance is only marginally worse than that of the NTK and, in certain cases, is shown to be superior. In particular, we establish bounds for the relative test losses, verify them with numerical tests, and identify the regularity of the kernel as the key determinant of performance. In addition to providing a theoretical grounding for using CKs instead of NTKs, our framework suggests a recipe for improving DNN accuracy inexpensively. We present a demonstration of this on the foundation model GPT-2 by comparing its performance on a classification task using a conventional approach and our prescription. We also show how our approach can be used to improve physics-informed operator network training for regression tasks as well as convolutional neural network training for vision classification tasks.