On the Effect of Stochastic Noise on Kernel-Target Alignment: An NTK Analysis

Neural Tangent Kernel (NTK) has emerged as an important framework to understand the behavior of neural networks. An important notion of the kernel is kernel–target alignment, which measures how much the top eigenfunctions of NTK are aligned to the target function. In a rich feature learning regime, it is observed that NTK evolves to increase the kernel-target alignment. Recent work has shown that this particular evolution is prominent during the edge of stability (EoS) from a large learning rate of gradient descent (GD). In this work, we examine the evolution of NTK under stochastic gradient descent (SGD). Across networks of various architectures, we observe that stochastic noise (e.g. from a small batch size of SGD) can induce a larger kernel-target alignment. Theoretically, in a toy example of a diagonal neural network, we show that stochastic noise can selectively suppress weakly target-aligned modes and thereby increase the kernel-target alignment. We also connect such mode suppression to low-rank bias of SGD. Empirically, we observe that low-rank bias is one of the elements that are correlated with the kernel-target alignment. However, we also find that rank reduction alone does not determine kernel-target alignment; rather, the effect depends critically on whether the remaining spectral directions align to the target.