Coding schemes in neural networks learning classification tasks

TL;DR

The paper addresses how neural networks learn interpretable, task-dependent representations by analyzing the Bayesian weight posterior in the non-lazy, mean-field regime and how the observed coding schemes depend on neuronal nonlinearity. It develops a theory that yields closed-form single-neuron posteriors and posterior-averaged kernels, revealing three distinct coding schemes—analog (linear), redundant (sigmoidal), and sparse (ReLU)—that emerge as the hidden nonlinearity changes. Across toy tasks and standard datasets like MNIST and CIFAR10, the work shows that non-lazy networks learn strong features but improve generalization mainly through reduced predictor variance rather than mean-predictor improvements, with symmetry breaking and neural collapse concepts offering a unifying view of the representations. The study highlights how weight-scaling and nonlinearity profoundly shape representations, offering insights into representation learning, generalization, and potential transfer learning implications in deep networks, while noting regime-specific limitations and the role of regularization in the observed phenomena.

Abstract

Neural networks posses the crucial ability to generate meaningful representations of task-dependent features. Indeed, with appropriate scaling, supervised learning in neural networks can result in strong, task-dependent feature learning. However, the nature of the emergent representations, which we call the `coding scheme', is still unclear. To understand the emergent coding scheme, we investigate fully-connected, wide neural networks learning classification tasks using the Bayesian framework where learning shapes the posterior distribution of the network weights. Consistent with previous findings, our analysis of the feature learning regime (also known as `non-lazy', `rich', or `mean-field' regime) shows that the networks acquire strong, data-dependent features. Surprisingly, the nature of the internal representations depends crucially on the neuronal nonlinearity. In linear networks, an analog coding scheme of the task emerges. Despite the strong representations, the mean predictor is identical to the lazy case. In nonlinear networks, spontaneous symmetry breaking leads to either redundant or sparse coding schemes. Our findings highlight how network properties such as scaling of weights and neuronal nonlinearity can profoundly influence the emergent representations.

Figures (26)

• Figure 1: (A) Sketch of the last layer feature space in the lazy regime, where the predominantly random features are almost orthogonal to the readout weight vector ($O(\sqrt{N})$ overlap), and in the non-lazy regime, where the learned features are aligned to the readout weight vector ($O(N)$ overlap). (B) Fully connected network with two hidden layers ($L=2$) and two outputs ($m=2$). Readouts are scaled by $1/N$ instead of $1/\sqrt{N}$ in non-lazy networks after training, i.e., under the posterior, which enforces the network to learn strong representations. (C) Posterior-averaged kernels of non-lazy networks with a single hidden layer for random binary classification of orthogonal data. Nonlinearity from left to right: linear, sigmoidal, and ReLU. (D) Activations for a given weight sample from the networks shown in (C); for ReLU only the $20$ most active neurons are shown. Parameters: $N=P=500$, $N_{0}=510$, classes assigned with probability $1/2$, targets $y_{+}=1$ and $y_{-}=-1$.
• Figure 2: Analog coding in two hidden layer linear networks on random classification task. (A) Readout weight posterior on first class; theoretical distribution (\ref{['eq:a_posterior_linear']}) as black dashed line. (B) Sampled readout weights for first class and two arbitrary neurons. (C,D) Activations of all neurons on all training inputs for a given weight sample in the second (C) and first (D) hidden layer. Neurons sorted by mean squared activity; inputs are sorted by class. (E,F) Kernel on training data from sampling in the first (E) and second (F) hidden layer. Parameters: $N=P=200$, $N_{0}=220$, classes assigned with fixed ratios $[1/2,1/4,1/4]$, targets $y_{+}=1$ and $y_{-}=0$.
• Figure 3: Coding scheme and generalization of two hidden layer linear networks on MNIST. (A) Samples of the readout weights of all three classes. (B,C) Activations of all neurons on all training (B) and $100$ test (C) inputs for a given weight sample. (D) Kernel on $100$ test inputs from sampling. (E) Mean predictor for class 0 from sampling (gray) and theory (eq. (\ref{['eq:predictor_linear']}), black dashed). (F) Generalization error for each class averaged over $P_{*}=1,000$ test inputs from sampling (gray bars), theory (Eq. (\ref{['eq:predictor_linear']}), black circles), and GP theory (back triangles). Parameters: $N=P=100$, $N_{0}=784$, classes 0, 1, 2 assigned randomly with probability $1/3$, targets $y_{+}=1$ and $y_{-}=0$.
• Figure 4: Redundant coding in single layer sigmoidal networks on random classification task. (A) Distribution of readout weights for first class across neurons and samples (blue); theoretical distribution (black dashed) includes finite $P$ corrections (see supplement). (B) Sampled readout weights for first class and two arbitrary neurons. (C) Samples of the readout weights of all three classes. (D) Activations of all neurons on all training inputs for a given weight sample. (E) Distribution of training activations of neurons with 1-1-0 code on inputs from class 0 across neurons for a given weight sample (blue) and according to theory (black dashed). (F) Kernel on training data from sampling. Parameters: $N=P=500$, $N_{0}=520$, classes assigned with fixed ratios $[1/2,1/4,1/4]$, targets $y_{+}=1$ and $y_{-}=1/2$.
• Figure 5: Coding scheme transitions in single layer sigmoidal networks on random classification task. (A-D) Activations of all neurons on all training inputs for a given weight sample with changing training target $y_{-}$. Parameters as in Fig. \ref{['fig:sigmoidal_code']} except that $y_{-}$ changes as indicated in the figure.