The Neural Tangent Kernel for Classification
Pith reviewed 2026-05-20 14:21 UTC · model grok-4.3
The pith
Wide neural networks with cross-entropy loss remain in the lazy training regime when parameter regularization is used or when all class probabilities are strictly positive.
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
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.
What carries the argument
The infinite-width Neural Tangent Kernel, which remains constant under the stated conditions and thereby lets the entire training trajectory be replaced by a linear model whose solution is expressed directly through the kernel.
If this is right
- Training dynamics on classification are well approximated by a linear model in parameter space.
- The final predictor admits an explicit closed-form characterization written in terms of the NTK.
- The distribution of predictors arising from different random initializations connects directly to Bayesian posterior uncertainty.
- The same constant-NTK regime that was previously known only for regression losses now applies to cross-entropy classification.
Where Pith is reading between the lines
- The explicit NTK solution could be used to derive generalization bounds for classification without running full gradient descent.
- Adding modest parameter regularization may be a practical way to keep wide networks inside the theoretically tractable lazy regime.
- The same non-degeneracy condition on targets might allow NTK analysis for other nonlinear output transformations beyond cross-entropy.
Load-bearing premise
The networks are wide enough for the infinite-width NTK limit to hold and for random initialization to place the network in the linearization regime that persists for the whole training run.
What would settle it
Train a sufficiently wide network on a multi-class dataset with cross-entropy loss, measure the NTK at initialization and after many steps, and check whether it stays approximately constant when weight decay is present or when every class has positive probability in the targets.
Figures
read the original 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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The manuscript extends Neural Tangent Kernel (NTK) theory from regression to classification with cross-entropy loss. It shows that parameter-space regularization keeps the NTK constant during training of wide networks. Without regularization, the lazy regime holds when targets are non-degenerate (all classes have strictly positive probability). Under these conditions, dynamics are approximated by the linearized model, yielding an explicit NTK-based solution characterization. The work also analyzes the distribution of predictors induced by random initialization and relates it to Bayesian methods.
Significance. If the derivations and conditions are rigorously established, the result would be a meaningful advance by bringing NTK analysis to classification, a dominant practical setting. The explicit solution form and initialization-induced uncertainty analysis supply concrete, testable predictions that connect to kernel methods and Bayesian views. The non-degenerate-target restriction is a genuine limitation on scope but does not invalidate the conditional claims.
major comments (2)
- [§4] §4, Definition 4.1 and surrounding text: the non-degenerate targets condition (strictly positive probability for every class) excludes standard one-hot labels. The abstract and title frame the contribution as addressing classification, yet the unregularized case therefore applies only to softened targets; the manuscript should either prove an approximate constancy result in the limit of vanishing probabilities or supply empirical evidence that the NTK remains nearly constant for near-degenerate targets.
- [§3.2] §3.2, Eq. (12): the constancy proof for the regularized cross-entropy case is derived under continuous-time gradient flow. It is unclear whether the same invariance holds for the discrete SGD steps used in all reported experiments; a short perturbation argument or numerical check would strengthen the claim that the NTK remains constant throughout practical training.
minor comments (2)
- [Introduction] Introduction, paragraph 3: the phrase 'non-degenerate, i.e., when all classes have strictly positive probability' could be accompanied by a one-sentence remark on how this relates to label smoothing or temperature scaling commonly used in practice.
- [Figure 2] Figure 2 caption: the experimental details (network width, depth, dataset, and regularization strength) are missing; adding them would make the NTK-evolution plots easier to interpret.
Simulated Author's Rebuttal
We thank the referee for the constructive comments on our manuscript. We address each major point below and describe the revisions we will incorporate.
read point-by-point responses
-
Referee: [§4] §4, Definition 4.1 and surrounding text: the non-degenerate targets condition (strictly positive probability for every class) excludes standard one-hot labels. The abstract and title frame the contribution as addressing classification, yet the unregularized case therefore applies only to softened targets; the manuscript should either prove an approximate constancy result in the limit of vanishing probabilities or supply empirical evidence that the NTK remains nearly constant for near-degenerate targets.
Authors: The non-degenerate target condition is required for the exact invariance result in the unregularized case, as the proof shows that degenerate targets (such as one-hot labels) can cause logits to diverge and break constancy. This is a genuine scope limitation for the unregularized setting, which we already state explicitly. The regularized case holds without this restriction. To strengthen the manuscript, we will add a new subsection with numerical experiments demonstrating that the NTK remains nearly constant (within 1-2% relative change) for targets with small positive probabilities down to 10^{-6}, approximating one-hot labels. We will also clarify in the introduction that practical classification often employs label smoothing, which satisfies the non-degenerate condition. revision: partial
-
Referee: [§3.2] §3.2, Eq. (12): the constancy proof for the regularized cross-entropy case is derived under continuous-time gradient flow. It is unclear whether the same invariance holds for the discrete SGD steps used in all reported experiments; a short perturbation argument or numerical check would strengthen the claim that the NTK remains constant throughout practical training.
Authors: We agree that the invariance is proven under continuous-time gradient flow. For discrete SGD, we will add a brief perturbation analysis in §3.2 showing that the per-step change in the NTK is O(η) where η is the learning rate; thus for the small step sizes used in practice the NTK remains approximately constant, with the error accumulating linearly in the number of steps but remaining small over typical training horizons. We will also include a numerical check in the experiments confirming that the NTK variation under SGD matches the continuous prediction within measurement noise. revision: yes
Circularity Check
No significant circularity; derivation self-contained under stated assumptions
full rationale
The paper identifies explicit conditions (parameter-space regularization or non-degenerate targets with strictly positive class probabilities) under which the NTK remains constant for cross-entropy loss, then uses the linearized model to characterize the solution. These conditions are derived from the gradient-flow dynamics and Jacobian evolution equations rather than being fitted to data or defined in terms of the target result. No load-bearing self-citations, self-definitional steps, or renamings of known results are present that would make the central claim equivalent to its inputs by construction. The work extends prior NTK regression results with domain-specific analysis that remains independently falsifiable via the infinite-width limit assumptions.
Axiom & Free-Parameter Ledger
axioms (1)
- domain assumption Infinite-width limit in which the NTK remains constant during training
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.