Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes

An Infinite-Width Analysis on the Jacobian-Regularised Training of a Neural Network

ICML 2024 · Probabilistic and scientific modeling › Uncertainty and Gaussian processes · approach: Infinite-width networks as Gaussian processes · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows that an MLP and its input-output Jacobian converge jointly to a Gaussian process at infinite width, and that robust training with a Jacobian regulariser follows a linear ODE governed by a Jacobian NTK.

Problem

Whether infinite-width theory (GP limit at initialisation, NTK training dynamics) extends to the input-output Jacobian of an MLP, and how Jacobian-regularised ("robust") training behaves in the infinite-width limit.

Approach

Extends NNGP and NTK theory to jointly cover the network output and its input Jacobian, defining a Jacobian NNGP kernel (via tensor programs) and a Jacobian NTK (JNTK) that stays constant during robust training, so the trained infinite-width MLP equals a kernel regression solution.

Key results

• Finite-width Jacobian NNGP and JNTK estimates converge toward their infinite-width limits as width grows (2^6 to 2^13), validating the theory empirically.
• The minimum eigenvalue of the JNTK Gram matrix is positive only for deeper networks (e.g., depth 11 with GeLU, depth 6 with erf), versus always-positive NTK.
• With Jacobian regularisation, eigenfeatures of the kernel regression solution show alignment between test accuracy and robustness, unlike standard training.

Measured effects reported in the paper

Surprising findings

• The minimum eigenvalue of the JNTK Gram matrix is positive only for sufficiently deep networks, so the full-rank assumption can fail for shallow MLPs.
• In the kernel-regression analysis, standard-training eigenfeatures show little correlation between accuracy and robustness, while Jacobian-regularised ones show alignment.

Limitations

• Width requirement F is implicit (proof relies on tensor programs with no explicit rate), and kappa must be small.
• Full-rank JNTK assumption is not known theoretically and only holds empirically for sufficiently deep networks.
• Experiments are on small MLPs and simple UCI datasets; the Figure-level validation does not directly verify Gaussianity.

Experimental setup

Datasets: synthetic Fibonacci-lattice hypersphere dataset (256 points in R^4) Algerian forest fire dataset (UCI) banknote authentication (UCI) connectionist bench (UCI)

Benchmarks: max-norm distance of finite vs limiting NNGP/JNTK kernels UCI forest fire test accuracy under perturbation

Metrics: max-norm distance test accuracy test accuracy after input perturbation minimum eigenvalue of Gram matrix

Baselines: standard NTK training (without Jacobian regulariser)

Builds on: neural tangent kernel tensor programs NNGP

Closest papers by meaning

• Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs (NeurIPS 2025 · similarity 0.89)
• Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $\mu$ Parametrization (ICML 2025 · similarity 0.88)
• On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling (NeurIPS 2025 · similarity 0.88)
• Quantitative convergence of trained neural networks to Gaussian processes (NeurIPS 2025 · similarity 0.88)
• Unveiling The Matthew Effect Across Channels: Assessing Layer Width Sufficiency via Weight Norm Variance (NeurIPS 2024 · similarity 0.87)
• Adaptive kernel predictors from feature-learning infinite limits of neural networks (ICML 2025 · similarity 0.87)