Neural Tangent Kernel Perspective on Parameter Space Symmetries

Parameter-space symmetries are transformations that modify a model's parameters without changing its outputs. These transformations can be leveraged to accelerate optimization and enhance generalization. Remarkably, applying a single transformation, either before or during training, often suffices to realize these benefits. While the effectiveness of this approach is promising, its underlying mechanisms remain poorly understood. In this paper, we provide an explanation within the Neural Tangent Kernel (NTK) framework, by analyzing how such transformations affect the properties of the kernel. We show that maximizing the loss gradient norm via a symmetry is equivalent to maximizing the alignment between the NTK and the network. Since kernel alignment is known to correlate with the optimization rate in the NTK limit, this result elucidates how loss gradient optimization facilitates faster training. To establish the validity of this approach, we prove that parameter-space symmetries preserve the NTK limit.