Representations and generalization in artificial and brain neural networks

Abstract

Humans and animals excel at generalizing from limited data, a capability yet to be fully replicated in artificial intelligence. This perspective investigates generalization in biological and artificial deep neural networks (DNNs), in both in-distribution and out-of-distribution contexts. We introduce two hypotheses: First, the geometric properties of the neural manifolds associated with discrete cognitive entities, such as objects, words, and concepts, are powerful order parameters. They link the neural substrate to the generalization capabilities and provide a unified methodology bridging gaps between neuroscience, machine learning, and cognitive science. We overview recent progress in studying the geometry of neural manifolds, particularly in visual object recognition, and discuss theories connecting manifold dimension and radius to generalization capacity. Second, we suggest that the theory of learning in wide DNNs, especially in the thermodynamic limit, provides mechanistic insights into the learning processes generating desired neural representational geometries and generalization. This includes the role of weight norm regularization, network architecture, and hyper-parameters. We will explore recent advances in this theory and ongoing challenges. We also discuss the dynamics of learning and its relevance to the issue of representational drift in the brain.

Keywords: deep neural networks; few-shot learning; neural manifolds; representational drift; visual cortex.

Fig. 1.

(A) Illustration of three layers in a visual hierarchy where the population response of the first layer is mapped into intermediate layer by F1 and into the last layer by F2 (Top) (10). The transformation of per-stimuli responses is associated with changes in the geometry of the object manifold, the collection of responses to stimuli of the same object (colored blue for a “dog” manifold and pink for a “cat” manifold). Changes in geometry may result in transforming object manifolds that are not linearly separable (in the first and intermediate layers) into separable ones in the last layer (separating hyperplane, colored orange). (B) Changes in classification capacity ${\alpha }_C$, manifold radius $R_M$, manifold dimension $D_M$, and classification margin $\kappa$ across the layers of pre-trained DNNs (ResNets).

Fig. 2.

(A and B) Examples of novel objects, here “coatis” (blue) and “numbats” (green), are presented to the ventral visual pathway (Top), modeled by a trained DNN (Bottom), eliciting a pattern of activity across IT-like neurons in the feature layer. We model concept learning as learning a linear readout $\mathit{w}$ to classify these activity patterns. (C), Generalization accuracy is very high across pairs of novel objects from the ImageNet21k dataset when using a pre-trained DNN (orange), but poor when using a randomly initialized DNN (blue), or a linear classifier in the pixel space of input images (gray). (D) Few-shot learning improves along the ventral visual hierarchy from pixels to V1 to V4 to IT, due to orchestrated transformations of object manifold geometry. The layerwise behavior of a trained ResNet50 (blue), Alexnet (light blue), and an untrained ResNet50 (gray) is included for comparison. We align V1, V4, and IT to the most similar ResNet layer under the BrainScore metric (20) (see ref. for details).

Fig. 3.

(A) We compare the empirical generalization error in 1-, 2-, and 5-shot learning experiments to the prediction from our geometric theory (Eq. 3) on all pairs of objects from the ImageNet21k dataset, using object manifolds derived from a trained ResNet50. x-axis: SNR obtained by estimating neural manifold geometry. y-axis: Empirical generalization error measured in few-shot learning experiments. Theoretical prediction (dashed line) shows a good match with experiments. (B) We provide additional examples of 5-shot prototype learning experiments in a ResNet50 (colored points), along with the prediction from our geometric theory (dashed line), on four randomly selected novel visual objects from the ImageNet21k dataset. Each panel plots the generalization error of one novel visual object (e.g., “Virginia bluebell”) against all 999 other novel visual objects. Each point represents the average generalization error on one such pair of objects. x-axis: SNR (Eq. 3) obtained by estimating neural manifold geometry. y-axis: Empirical generalization error measured in few-shot learning experiments. Theoretical prediction (dashed line) shows a good match with experiments. (C) In a pre-trained ResNet50 (blue) dimensionality expands dramatically in the early layers and contracts in the later layers, while in the primate visual pathway (Black) dimensionality contracts from the V1-like layer to V4, then expands from V4 to IT. (D) Single-manifold eigenspectra in macaque V4 (black) and the corresponding layer of a pre-trained ResNet50 (blue).

Fig. 4.

(A) Two stages of learning of Langevin dynamics with small T, ${\sigma }_0$ controls the width of weight distribution at initialization, $\sigma$ controls the size of the solution space, and $T$ relates to the sampling speed. (B) Example trajectories of the predictor from three different initializations, the dynamics is initially deterministic and starts to fluctuate as ${\mathrm{\Theta }}_t$ drifts in the solution space after reaching zero training error.

Fig. 5.

(A) Schematics for the BPKR approach. A renormalization factor is introduced at each step during backward integration until all the network weights are averaged out. (B) Theory (black solid line) and simulation (blue points) of generalization error ${\epsilon }_g={⟨{ϵ}_g(\mathbf{x},\mathbf{y}(\mathbf{x}))⟩}_{\{\mathbf{x},\mathbf{y}(\mathbf{x})\}}$ on binary MNIST classification in fully connected ReLU networks, for small (Top) and large (Bottom) $\sigma$. The approximate theory for ReLU networks agrees remarkably well with the numerics.

Fig. 6.

(A and B) NNGP and mean layer-wise kernels in classifying eight MNIST digits SI Appendix, 3, (57). (C) Generalization error averaged across test examples for finite width random feature model (blue), infinitely wide network following the NNGP theory (yellow), and the learned network following the BPKR theory (red, overlaying NNGP theory). (D) The NNGP kernel, and the mean layer-wise kernel of hidden layer $l=2,4$. For a 4-hidden-layer ReLU network trained on four MNIST digits grouped into two higher-order categories of even vs. odd. The values of the kernel are small since we take relatively small $\sigma$ (SI Appendix, 3).

Fig. 7.

(A) Signal as a function of hidden layer depth $l$. For ReLU networks in the thermodynamic limit, signal increases with layer depth (blue). For linear networks (yellow, purple) and ReLU networks in the infinite width limit (red), signal remains unchanged across $l$. Error bars are across all distinct pairs of manifolds/digits. (B) Dimension as a function of $l$. In the infinite width limit, dimension remains constant with $l$ in linear networks (purple) and increases with $l$ in ReLU networks (red). In the thermodynamic limit, dimension decreases in linear networks (yellow), and is non-monotonic in ReLU networks (blue), similar to Fig. 3C. Error bars are across all manifolds/digits.

Fig. 8.

(A) Comparison of the generalization error dynamics between a network fully trained under Langevin dynamics (Eq. 14, shown in red), and a network with $\mathbf{a}$ frozen at time $t_0$ in the diffusive learning stage, and ${\mathcal{W}}_t$ randomly drifting afterward (shown in blue). $\tau$ denotes the difference between the current time $t$ and $t_0$. (B) Both the temporal correlation of the top right singular vector (${\rho }_{\mathbf{u}}(\tau )$) and the correlation between $\mathbf{u}(t)$ and $\mathbf{Y}$ (${\rho }_{\mathbf{u},\mathbf{Y}}$) remain close to 1, representing the constant alignment between the Top Right singular vector of the representation and the training labels. Temporal correlation of the Top Left singular vector (${\rho }_{\mathbf{v}}(\tau )$) gradually decreases with the time difference $\tau$, representing the random drift in the feature space.