neural-feature-learning

name

Neural Feature Learning

description

Neural Feature Learning methods open the black box of *how* a trained neural network discovers useful representations from data, by deriving the directions along which each layer's effective input geometry is reshaped during training. The unifying object is the Neural Feature Matrix (NFM) — for a fully-connected layer, the NFM is the weight matrix multiplied by its transpose, and acts as a metric on that layer's input space. Empirically, the NFM aligns at convergence with the Expected Gradient Outer Product (EGOP), defined as the average outer product of the network output's gradient with respect to the same layer's input. The NFM's top eigenvectors point in the directions the network has learned to use — turning feature learning into an analyzable, kernel-shaped object.

Advantage

The framework supplies a kernel-shaped interpretation of feature learning: any fully-connected or convolutional architecture can be analyzed through its NFM / EGOP, the same diagnostic transfers across layers and widths, and the resulting features are far more interpretable than raw weight matrices. Because the EGOP only needs per-sample input gradients (cheap to obtain via standard reverse-mode autograd), it also powers kernel-style algorithms — Recursive Feature Machines (RFM) iteratively refit a kernel using the EGOP and recover neural-network-style features without backpropagation through model weights — and yields clean theoretical predictions in the infinite-width limits, where the choice of parameterization decides whether the network truly learns features (muP) or freezes into a fixed kernel (NTK / NNGP).

Limitation

Predictions hinge on width, parameterization, and training-dynamics assumptions: the NFM-EGOP alignment is an empirical ansatz verified on standard architectures but is not universal, and the infinite-width limits (NTK, NNGP, muP) describe finite-width training only approximately — and only when initialization scales and learning rates are tuned consistently with the chosen parameterization. Results are also dataset- and layer-specific — features extracted from one site rarely transfer verbatim — and EGOP estimation is expensive on high-resolution inputs because every per-sample input gradient must be accumulated into a dense matrix whose size grows quadratically with the input dimension.

Submethods

The category comprises 4 main forms, progressing from the diagnostic (DNFA), to the algorithmic instantiation (RFM), to the theoretical limits (TP4), and finally to the convolutional extension (ConvRFM):

• Eigenvector Feature Direction (Deep Neural Feature Ansatz): States that for each fully-connected layer, the principal eigenvectors of the NFM converge during training to the principal eigenvectors of the EGOP measured at that layer's input — so the directions a network ends up using can be read directly from its weights at convergence. This Deep Neural Feature Ansatz (DNFA) supplies the spectral lens for feature learning: visualize the top eigenvectors per layer to see which input directions the network has learned to rely on. You can find a demo for this method in ./eigenvector-feature-direction. This demo shows deep-neural-feature-ansatz: Use this skill when working with the Deep Neural Feature Ansatz (DNFA) — verifying feature learning in neural networks, training fully connected networks on image/tabular datasets, computing Neural Tangent Kernels (NTK) and Neural Network Gaussian Processes (NNGP), or reproducing experiments from the paper "The Deep Neural Feature Ansatz" (arXiv:2212.13881).

• Gradient Outer Product (Recursive Feature Machines): Lifts the EGOP from a diagnostic to an algorithm. Recursive Feature Machines (RFM) wrap a Mahalanobis kernel — parameterized by a feature matrix initialized to the identity — around kernel ridge regression: fit the kernel, re-estimate the feature matrix as the empirical EGOP of the fitted predictor, and iterate. The fixed point reproduces neural-network-like feature learning without backpropagation through model weights, demonstrating that the EGOP captures the essential mechanism of feature acquisition. You can find a demo for this method in ./gradient-outer-product. This demo shows recursive-feature-machines: Use this skill when working with Recursive Feature Machines (RFM) for kernel-based machine learning, feature learning, or when implementing backpropagation-free models that learn features similarly to neural networks.

• Kernel / NTK Feature Regime (Tensor Programs IV): Characterizes the limits of feature learning. In the infinite-width limit, the parameterization determines which of three regimes a network falls into — Gaussian-process inference at initialization (NNGP), lazy / kernel-regression dynamics with frozen features (NTK), or genuine feature learning (muP, the maximal-update parameterization). Tensor Programs IV (TP4) makes these regimes precise and supplies the recipe for picking initialization scales and learning rates so that finite-width networks actually learn features rather than degenerate to a fixed kernel. You can find a demo for this method in ./kernel-ntk-feature-regime. This demo shows tp4-feature-learning: Use this skill when working with infinite-width neural networks for feature learning, replicating Word2Vec or MAML experiments from the Tensor Programs series (TP4), or implementing infinite-width limits (GP, NTK, muP) for meta-learning and word embedding tasks.