(An) infinite-width analysis on the Jacobian-regularised training of a neural network

Abstract

In this paper, we extend the Neural Tangent Kernel analysis to incorporate the Jacobian Regularisation, which includes additional term in the training objective. This additional regulariser forces the neural network to be flat around its training data, which makes the network to be robust against the adversarial attack. We first show that the Jacobian of neural network converges to the Gaussian Process jointly with the output as the widths of neural network tends to infinity, and characterise the kernel of this GP. We then show that the training under Jacobian regularisation is described by the inner products between the gradients of outputs and Jacobians with respect to the parameters, which we named this inner product as Jacobian Neural Tangent Kernel due to its similarity of NTK in literatures. We show that the properties of NTK also hold in JNTK, that at initialisation the JNTK converges to the deterministic asymptotic counterpart, and it stays constant during training if the width of neural network is sufficiently large. We finally prove that the Jacobian regularised training dynamics coincide with kernel regression with JNTK. Using this correspondence to kernel regression, we analyse how Jacobian regularised training enhances the robustness, and when is it possible