Visualising Feature Learning in Deep Neural Networks by Diagonalizing the Forward Feature Map

2024-10-05

arXiv.org

AI Summary

The tight coupling between feature richness and model performance hinders independent evaluation of feature learning mechanisms. Method: We propose a decoupled analytical framework that decomposes deep networks into a forward feature mapping Φ and a linear classifier, and diagonalize the Jacobian operator of Φ under gradient descent dynamics—a novel operator diagonalization technique. Contribution/Results: This reveals two distinct learning paradigms—Minimum Features (MF) and Expanded Features (EF)—emerging early in training, unifying the explanation of neural collapse. Extensive validation across diverse architectures and datasets shows that optimal generalization typically occurs during the MF phase, and neural collapse generalizes beyond classification to regression tasks. Our framework yields an interpretable, performance-agnostic visualization tool for feature dynamics, establishing a new analytical paradigm for studying feature learning mechanisms.

Technology Category

• Machine Learning: Learning with Manifolds
• Computer Vision: Learning & Optimization for CV
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

Abstract

Deep neural networks (DNNs) exhibit a remarkable ability to automatically learn data representations, finding appropriate features without human input. Here we present a method for analysing feature learning by decomposing DNNs into 1) a forward feature-map $Phi$ that maps the input dataspace to the post-activations of the penultimate layer, and 2) a final linear layer that classifies the data. We diagonalize $Phi$ with respect to the gradient descent operator and track feature learning by measuring how the eigenfunctions and eigenvalues of $Phi$ change during training. Across many popular architectures and classification datasets, we find that DNNs converge, after just a few epochs, to a minimal feature (MF) regime dominated by a number of eigenfunctions equal to the number of classes. This behaviour resembles the neural collapse phenomenon studied at longer training times. For other DNN-data combinations, such as a fully connected network on CIFAR10, we find an extended feature (EF) regime where significantly more features are used. Optimal generalisation performance upon hyperparameter tuning typically coincides with the MF regime, but we also find examples of poor performance within the MF regime. Finally, we recast the phenomenon of neural collapse into a kernel picture which can be extended to broader tasks such as regression.

Problem

• Measure rich regime independently of performance metrics
• Develop interpretable feature metrics for visualization
• Analyze lazy-to-rich transitions without performance proxies

Innovation

• Independent measure for rich regime dynamics
• Interpretable feature metrics for visualization
• Low-rank bias generalizes neural collapse metrics