Task-Driven Kernel Flows: Label Rank Compression and Laplacian Spectral Filtering

TL;DR

This work develops a dynamically driven theory of feature learning for wide networks with a $C$-dimensional linear head, framing learning as a two-time-scale, kernel-flow process. By deriving an explicit kernel ODE under fast readout and regularization, the authors show a label-driven rank compression with $ ext{rank}(K_ ext{infty})\le C$, and, for squared loss, a precise water-filling spectral truncation that enforces neural collapse. They further demonstrate intrinsic low-rank SGD noise confined to the task subspace, provide a population-analytic view of generalized risk, and unify supervised, self-supervised, and semi-supervised learning within a thermodynamics of feature learning. The framework highlights a trade-off between reachability and variance, with SSL exhibiting high-rank diffusion on smooth graph spectra, while supervision compresses representations to the label subspace. Collectively, the results offer a physics-inspired lens on representation learning and guidance for architecture and optimization design.

Abstract

We present a theory of feature learning in wide L2-regularized networks showing that supervised learning is inherently compressive. We derive a kernel ODE that predicts a "water-filling" spectral evolution and prove that for any stable steady state, the kernel rank is bounded by the number of classes ($C$). We further demonstrate that SGD noise is similarly low-rank ($O(C)$), confining dynamics to the task-relevant subspace. This framework unifies the deterministic and stochastic views of alignment and contrasts the low-rank nature of supervised learning with the high-rank, expansive representations of self-supervision.

Figures (3)

• Figure 1: Dynamics of Geometric Alignment. (Left) The evolution of the subspace overlap score during training (with $D=N=100$). The system exhibits a distinct phase transition, escaping a saddle point to achieve perfect alignment (Score $\approx 1.0$) with the target subspace. (Right) The heatmap of the final feature kernel $K_\infty$ entries. The checkerboard pattern confirms the commutativity structure $[K_\infty, M_Y] \approx 0$, validating that the learned features share the same eigenbasis as the labels.
• Figure 2: Empirical Verification of the Spectral Truncation Law. A rigorous comparison between theory and experiment. The red dashed line represents the analytical prediction from Theorem \ref{['thm:spectral_truncation']} (Eq. \ref{['eq:water_filling']}), while blue dots show the eigenvalues of the kernel trained via gradient descent. The grey vertical dotted line indicates the theoretical truncation threshold $\tau = \lambda \mu$. The experimental data perfectly matches the "Water-Filling" curve, confirming that eigenmodes with signal strength $\sigma_i \le \tau$ are strictly pruned ($k_i=0$).
• Figure 3: Rank Collapse Phase Transition. The effective rank of the learned representation versus weight decay $\mu$. The red dashed line denotes the number of classes ($C=10$). As predicted by the spectral truncation law, increasing $\mu$ raises the noise threshold $\tau$. When $\tau$ surpasses the signal strength of intra-class variations, the rank undergoes a sharp phase transition, collapsing from the ambient dimension ($N$) onto the label subspace ($C$), marking the onset of the Neural Collapse regime.

Theorems & Definitions (46)

• Theorem 1: Global Convergence
• proof
• Theorem 2: Label-Driven Rank Compression
• proof
• Lemma 3: Variational Alignment Principle
• proof
• Theorem 4: Spectral Truncation Law
• Theorem 5: Equivalence of Weight Decay and Nuclear Norm
• proof
• Remark 6: Over-parameterization Condition