Neural Networks as Kernel Learners: The Silent Alignment Effect

Abstract: Neural networks in the lazy training regime converge to kernel machines. Can neural networks in the rich feature learning regime learn a kernel machine with a data-dependent kernel? We demonstrate that this can indeed happen due to a phenomenon we term silent alignment, which requires that the tangent kernel of a network evolves in eigenstructure while small and before the loss appreciably decreases, and grows only in overall scale afterwards. We show that such an effect takes place in homogenous neural networks with small initialization and whitened data. We provide an analytical treatment of this effect in the linear network case. In general, we find that the kernel develops a low-rank contribution in the early phase of training, and then evolves in overall scale, yielding a function equivalent to a kernel regression solution with the final network's tangent kernel. The early spectral learning of the kernel depends on the depth. We also demonstrate that non-whitened data can weaken the silent alignment effect.

• The paper demonstrates that silent alignment enables the NTK to rapidly align its eigenvectors with task-relevant directions before scaling to achieve kernel regression equivalence.
• Empirical experiments with ReLU and Tanh activations show that small initial weights and whitened data are crucial for engaging the silent alignment effect.
• Analytical findings reveal that NTK evolution is depth-dependent, offering insights to enhance generalization and transfer learning in deep neural networks.

Neural Networks as Kernel Learners: The Silent Alignment Effect

The paper "Neural Networks as Kernel Learners: The Silent Alignment Effect" explores the behavior of neural networks (NNs) in two distinct training regimes: the lazy and rich regimes. In the lazy regime, NNs are known to converge to kernel machines utilizing the neural tangent kernel (NTK) that remains static throughout the training process. This static NTK results in dynamics that resemble kernel regression solutions. The rich regime, on the other hand, is characterized by evolving internal representations and NTK dynamics that are sensitive to the data's structure.

This study presents a perspective on the rich training regime by investigating the possibility of neural networks operating as kernel machines with dynamically evolving, data-dependent kernels. The authors introduce the concept of "silent alignment," wherein the NTK rapidly evolves in its eigenstructure early in the training before the network’s loss significantly decreases. Subsequently, the NTK grows in scale while maintaining its learned orientation, allowing the network to achieve a kernel regression solution with the final NTK.

Key Insights and Findings

1. Silent Alignment Phenomenon: Silent alignment is identified as a training phase where the NTK initially evolves its eigenvectors toward task-relevant directions while its scale remains small, followed by growth in overall scale. This phase allows for a kernel regression equivalent representation of the learned NN function using the final NTK.
2. Empirical Basis: The effect is demonstrated empirically through experiments on both linear and nonlinear networks, such as ReLU and Tanh activations. Silent alignment is notably evident in networks initialized with small weights and trained on whitened data.
3. Analytical Treatment: The paper provides an analytical account of silent alignment within fully connected linear networks, exploring how the kernel develops a low-rank contribution early in training. The authors show that the kernel’s evolution depends significantly on the network's depth.
4. Impact of Data Cochlearity: Non-whitened data can interfere with silent alignment, as demonstrated experimentally. This interference is attributed to the NTK needing to change its eigenvectors during times when the loss is significantly decreasing, which disrupts the silent alignment effect.

Practical and Theoretical Implications

• Kernel Regression Equivalence: The findings suggest a closer link between kernel methods and fully trained deep learning models. It posits that even in feature-rich regimes, NN functions can mirror kernel regression solutions with the plant NTK, offering a practical method of understanding and improving NN generalization and transfer learning properties.
• Transfer Learning Performance: The alignment of the NTK with task-relevant features suggests enhanced transfer learning capabilities when tasks are highly correlated with learned NNs. The final NTK, being aligned with training task directions, may enable NN to adapt efficiently to related tasks.

Future Directions in AI Developments

The results presented open up avenues for further exploration into the mechanics of NTK evolution and its implications for generalization, robustness, and transfer learning. An additional exploration could be in quantifying NTK changes in networks employing adaptive learning rates and how these affect overall learning trajectories.

The concept of silent alignment could guide the development of new training protocols that explicitly harness NTK dynamics to optimize representation learning. Moreover, better understanding kernel regression equivalence could refine transfer learning strategies, especially in environments where task covariance with the training objective is high. This approach could also be instrumental in devising unified frameworks for both kernel and NN methodologies in AI applications.