The Neural Tangent Kernel for Classification

2026-05-17

AI Summary

This work addresses the challenge of extending Neural Tangent Kernel (NTK) theory—originally developed for regression—to classification settings, where cross-entropy loss typically drives logits to diverge, thereby violating the linearization assumption underpinning NTK. By introducing either parameter-space regularization or non-degenerate target conditions, the paper establishes, for the first time, sufficient conditions under which sufficiently wide neural networks maintain a "lazy training" regime in classification tasks, ensuring the NTK remains approximately constant throughout training. This advancement enables a rigorous extension of NTK theory to classification, allowing precise characterization of both training dynamics and generalization behavior. Moreover, it reveals a theoretical connection between the predictive distribution induced by random initialization and Bayesian inference.

Abstract

In wide neural networks, the Neural Tangent Kernel (NTK) remains approximately constant during training, providing a powerful theoretical tool for studying training dynamics, generalization, and connections to kernel methods. However, this theory is largely restricted to regression losses. It was previously thought that training on a classification loss, or more generally losses involving nonlinear output transformations, breaks this property, leading to divergent logits and a breakdown of the linearization. In this paper, we extend NTK theory to classification by identifying conditions under which wide neural networks remain in the lazy training regime. We show that parameter-space regularization ensures a constant NTK during training for cross-entropy loss, while in the absence of regularization the regime is recovered when targets are non-degenerate, i.e. when all classes have strictly positive probability. Under these conditions, training is well-approximated by the linearized model, yielding an explicit characterization of the solution in terms of the NTK. We further analyze the distribution of trained predictors induced by random initialization and relate this notion of model uncertainty to Bayesian methods.

Problem

• Neural Tangent Kernel
• classification
• lazy training regime
• cross-entropy loss
• wide neural networks

Innovation

• Neural Tangent Kernel
• classification
• lazy training
• cross-entropy loss
• model uncertainty