An Infinite-Width Analysis on the Jacobian-Regularised Training of a Neural Network
This work provides theoretical insights into Jacobian regularization for neural networks, which is incremental to existing infinite-width research.
The authors extended infinite-width analysis to the Jacobian of deep neural networks, showing that both an MLP and its Jacobian converge to a Gaussian process at initialization and that robust training in this limit is described by a linear ODE. They experimentally validated these theoretical claims on wide finite networks and analyzed kernel regression to gain insights into Jacobian regularization.
The recent theoretical analysis of deep neural networks in their infinite-width limits has deepened our understanding of initialisation, feature learning, and training of those networks, and brought new practical techniques for finding appropriate hyperparameters, learning network weights, and performing inference. In this paper, we broaden this line of research by showing that this infinite-width analysis can be extended to the Jacobian of a deep neural network. We show that a multilayer perceptron (MLP) and its Jacobian at initialisation jointly converge to a Gaussian process (GP) as the widths of the MLP's hidden layers go to infinity and characterise this GP. We also prove that in the infinite-width limit, the evolution of the MLP under the so-called robust training (i.e., training with a regulariser on the Jacobian) is described by a linear first-order ordinary differential equation that is determined by a variant of the Neural Tangent Kernel. We experimentally show the relevance of our theoretical claims to wide finite networks, and empirically analyse the properties of kernel regression solution to obtain an insight into Jacobian regularisation.