Efficient kernel surrogates for neural network-based regression

TL;DR

The Conjugate Kernel (CK), an efficient approximation to the NTK that has been observed to yield fairly similar results, is studied, and the regularity of the kernel is identified as the key determinant of performance.

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.

Figures (13)

• Figure 1: The results of function regression for 100 different NNs, trained for 2400 epochs, and the corresponding approximations using the NTK, CK, and CKJ extracted from the NN at the end of the training.
• Figure 2: The test errors for NNs over the course of being trained to approximate the given functions, and for the corresponding NTK, CK, and CKJ approximations. The errors are averaged over ten iterations to reduce the effects of random initialization.
• Figure 3: Function regression results for 100 NNs with the widths of the hidden layers set to 256, and the corresponding NTK, CK, and CKJ approximations.
• Figure 4: The test errors for function regression for 100 trained NNs using the ReLU activation function, and the corresponding NTK and CK approximations.
• Figure 5: Averaged test errors for function regression over the course of training ten NNs using ReLU activations, and the corresponding NTK and CK approximations.

Theorems & Definitions (11)

• Remark 3.1
• Lemma 4.1
• Lemma 4.2
• Lemma 4.3
• Lemma 4.4
• Theorem 4.1
• Theorem 4.2
• Lemma 5.1
• Lemma 5.2
• Theorem 5.1