Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Training NTK to Generalize with KARE

Johannes Schwab, Bryan Kelly, Semyon Malamud, Teng Andrea Xu

TL;DR

The paper tackles generalization in overparameterized neural networks by proposing NTK-KARE, an approach that explicitly trains the Neural Tangent Kernel using Kernel Alignment Risk Estimator (KARE) rather than minimizing empirical risk directly. It shows that after-training NTK can be viewed as gradient boosting and introduces a concrete procedure to optimize the NTK via KARE, yielding the NTK-KARE estimator that relies on kernel learning alone. Empirically, NTK-KARE outperforms the standard DNN and the after-NTK across simulated data, MNIST, Higgs, and large-scale UCI benchmarks, achieving competitive or superior results including state-of-the-art comparisons on the UCI challenge. The work suggests explicit, high-capacity kernel learning as a viable alternative to end-to-end DNN optimization, while noting scalability challenges and outlining directions such as Falkon and EigenPro for large-scale kernel training and theoretical guarantees for KARE-based methods.

Abstract

The performance of the data-dependent neural tangent kernel (NTK; Jacot et al. (2018)) associated with a trained deep neural network (DNN) often matches or exceeds that of the full network. This implies that DNN training via gradient descent implicitly performs kernel learning by optimizing the NTK. In this paper, we propose instead to optimize the NTK explicitly. Rather than minimizing empirical risk, we train the NTK to minimize its generalization error using the recently developed Kernel Alignment Risk Estimator (KARE; Jacot et al. (2020)). Our simulations and real data experiments show that NTKs trained with KARE consistently match or significantly outperform the original DNN and the DNN- induced NTK (the after-kernel). These results suggest that explicitly trained kernels can outperform traditional end-to-end DNN optimization in certain settings, challenging the conventional dominance of DNNs. We argue that explicit training of NTK is a form of over-parametrized feature learning.

Training NTK to Generalize with KARE

TL;DR

The paper tackles generalization in overparameterized neural networks by proposing NTK-KARE, an approach that explicitly trains the Neural Tangent Kernel using Kernel Alignment Risk Estimator (KARE) rather than minimizing empirical risk directly. It shows that after-training NTK can be viewed as gradient boosting and introduces a concrete procedure to optimize the NTK via KARE, yielding the NTK-KARE estimator that relies on kernel learning alone. Empirically, NTK-KARE outperforms the standard DNN and the after-NTK across simulated data, MNIST, Higgs, and large-scale UCI benchmarks, achieving competitive or superior results including state-of-the-art comparisons on the UCI challenge. The work suggests explicit, high-capacity kernel learning as a viable alternative to end-to-end DNN optimization, while noting scalability challenges and outlining directions such as Falkon and EigenPro for large-scale kernel training and theoretical guarantees for KARE-based methods.

Abstract

The performance of the data-dependent neural tangent kernel (NTK; Jacot et al. (2018)) associated with a trained deep neural network (DNN) often matches or exceeds that of the full network. This implies that DNN training via gradient descent implicitly performs kernel learning by optimizing the NTK. In this paper, we propose instead to optimize the NTK explicitly. Rather than minimizing empirical risk, we train the NTK to minimize its generalization error using the recently developed Kernel Alignment Risk Estimator (KARE; Jacot et al. (2020)). Our simulations and real data experiments show that NTKs trained with KARE consistently match or significantly outperform the original DNN and the DNN- induced NTK (the after-kernel). These results suggest that explicitly trained kernels can outperform traditional end-to-end DNN optimization in certain settings, challenging the conventional dominance of DNNs. We argue that explicit training of NTK is a form of over-parametrized feature learning.
Paper Structure (14 sections, 1 theorem, 17 equations, 2 figures, 4 tables)

This paper contains 14 sections, 1 theorem, 17 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Main results
  4. Preliminaries
  5. Data-Dependent NTK: Gradient Boosting
  6. Training NTK Explicitly
  7. Experiments
  8. Simulated Data
  9. MNIST and Higgs Datasets
  10. UCI Classification Datasets
  11. Discussion
  12. Details for Numerical Experiments
  13. Hyperparameters
  14. Details

Key Result

Proposition 1

Let $\hat{y} = f(X; \theta)$ and $\ell_{\hat{y}} \coloneq \nabla_{\hat{y}} \ell$. Suppose that $f$ is differentiable, $\ell_{\hat{y}}$ is continuous, and that $\ell$ is such that, for any $R>0,$ the set $\{v\in {\mathbb R}:\ell(y,v)<R\}$ is bounded and $\ell\ge -A$ for some $A>0$. Suppose also that for some vector ${\mathcal{U}}_t$ that depends on the training data.

Figures (2)

  • Figure 1: Out-of-sample MSE on the synthetic dataset generated according to \ref{['eq:DGP-sim']} for NTK-KARE, after-NTK, and the DNN. We consider various parameterizations of the DGP by choosing $\gamma \in {1.0, 1.25, 1.5}$ (columns) and $\sigma \in {0.0, 0.1}$ (rows). NTK-KARE was fit using $z_{\rm KARE} = 0.1$. All other hyperparameters are provided in Table \ref{['tab:hyperparam-exp']}.
  • Figure 2: Out-of-sample mean-squared-error on the Mnist and Higgs dataset for NTK-KARE, after-NTK, and the DNN model using $n=1'000$ observations to train the models. We consider various DNN architecture by varying network $\text{depth} \in \{2, 4 \}$ (rows) and $\text{width} \in \{32, 64, 128 \}$ (columns). We used $z_{\rm KARE} = 0.1$ to fit NTK-KARE. All the other hyperparameters can be found in Table \ref{['tab:hyperparam-exp']}

Theorems & Definitions (1)

  • Proposition 1