Microscopic and collective signatures of feature learning in neural networks

TL;DR

This paper uses a statistical-mechanics approach to analyze how neural networks learn features, revealing both collective and microscopic signatures of feature learning in a large-width, large-data regime.

Contribution

It provides analytical insights into feature learning mechanisms in neural networks within a Bayesian framework, connecting microscopic weight changes and manifold geometry.

Findings

Distance between class manifolds varies nonmonotonically with temperature

Microscopic parameters undergo data-dependent displacements and develop correlations

Feature learning signatures are present even when the posterior resembles Gaussian process regression

Abstract

Feature extraction - the ability to identify relevant properties of data - is a key factor underlying the success of deep learning. Yet, it has proved difficult to elucidate its nature within existing predictive theories, to the extent that there is no consensus on the very definition of feature learning. A promising hint in this direction comes from previous phenomenological observations of quasi-universal aspects in the training dynamics of neural networks, displayed by simple properties of feature geometry. We address this problem within a statistical-mechanics framework for Bayesian learning in one hidden layer neural networks with standard parameterization. Analytical computations in the proportional limit (when both the network width and the size of the training set are large) can quantify fingerprints of feature learning, both collective ones (related to manifold geometry) and microscopic ones (related to the weights). In particular, (i) the distance between different class manifolds in feature space is a nonmonotonic function of the temperature, which we interpret as the equilibrium counterpart of a phenomenon observed under gradient descent (GD) dynamics, and (ii) the microscopic learnable parameters in the network undergo a finite data-dependent displacement with respect to the infinite-width limit, and develop correlations. These results indicate that nontrivial feature learning is at play in a regime where the posterior predictive distribution is that of Gaussian process regression with a trivially rescaled prior.

Equations

$$⟨\sigma (h_i^{\mu })\sigma (h_i^{\nu })⟩$$

$$+\underset{{\Delta }_2^{\mu \nu }}{\underbrace{\frac{\bar{Q}}{{\lambda }_1{N}_1}\sum _{\lambda ,\delta =1}^P{K}_{\mu \lambda }{K}_{\nu \delta }({\tilde{K}}_{(R)}^{-1}y{)}_{\lambda }({\tilde{K}}_{(R)}^{-1}y{)}_{\delta }}},$$

$$D^2={\Vert \frac{2}{P}\sum _{\mu =1}^{P/2}\sigma ({\mathbf{h}}_{+}^{\mu })-\frac{2}{P}\sum _{\mu =1}^{P/2}\sigma ({\mathbf{h}}_{-}^{\mu })\Vert }^2,$$

$$⟨D^2⟩$$

$$=\frac{4}{\alpha P}{y}^{T}Ky-\frac{4}{\alpha P}{y}^T{\Delta }_{1}y+\frac{4}{\alpha P}{y}^T{\Delta }_{2}y,$$

$$⟨D^2⟩\simeq \frac{4}{\alpha }\frac{1}{P}\sum _{\mu ,\nu =1}^P{y}_{\mu }{K}_{\mu \nu }{y}_{\nu }+\frac{{\lambda }_1}{\bar{Q}}({\bar{f}}_{+}-{\bar{f}}_{-}{)}^2,$$

$${\bar{f}}_{\pm }=\frac{2}{P}\sum _{\mu =1}^{P/2}⟨f_{\pm }^{\mu }⟩,⟨f_{\pm }^{\mu }⟩=\sum _{\lambda ,\sigma =1}^P(K_{(R)}{)}_{\mu \lambda }({\tilde{K}}_{(R)}^{-1}{)}_{\lambda \sigma }{y}_{\sigma }.$$

$$⟨w_{1h}{w}_{1k}⟩=\frac{{\delta }_{hk}}{{\lambda }_0}+\frac{({\Lambda }_1{)}_{hk}}{{\lambda }_0}+\frac{({\Lambda }_2{)}_{hk}}{{\lambda }_0},$$

$$({\Lambda }_1{)}_{hk}=\frac{\alpha \bar{Q}}{2{\lambda }_1}\frac{1}{P}\sum _{\mu ,\nu =1}^P({\tilde{K}}_{(R)}^{-1}{)}_{\mu \nu }{\Delta }_{hk}^{\mu \nu },$$

$$({\Lambda }_2{)}_{hk}=\frac{\alpha \bar{Q}}{2{\lambda }_1}\frac{1}{P}\sum _{\mu ,\nu =1}^P({\tilde{K}}_{(R)}^{-1}y{)}_{\mu }({\Delta }_{hk}^{\mu \nu })({\tilde{K}}_{(R)}^{-1}y{)}_{\nu }.$$

$$D^2=\frac{1}{N_1^{\eta }}{\Vert \frac{1}{N_{+}^{\delta }}\sum _{\mu =1}^{N_{+}}\sigma ({\mathbf{h}}_{+}^{\mu })-\frac{1}{N_{-}^{\delta }}\sum _{\mu =1}^{N_{-}}\sigma ({\mathbf{h}}_{-}^{\mu })\Vert }^2,$$

$$D^2=\frac{1}{N_1^{\eta }{N}_{+}^{2\delta }}\left[\begin{array}{c}\sigma ({\mathbf{h}}_{+})\sigma ({\mathbf{h}}_{-})\end{array}\right]\cdot \left[\begin{array}{cc}1& -\gamma \\ -\gamma & {\gamma }^2\end{array}\right]\cdot \left[\begin{array}{c}\sigma ({\mathbf{h}}_{+})\\ \sigma ({\mathbf{h}}_{-})\end{array}\right],$$

$$z_{\mu }=\{\begin{array}{l}1\text{for}\mu =1,...,N_{+}\\ -\gamma \text{for}\mu =N_{+}+1,...,P\end{array},$$

$$D^2=\frac{1}{N_1^{\eta }{N}_{+}^{2\delta }}\sum _{\mu \nu }^P{z}_{\mu }{z}_{\nu }\sum _i^{N_1}\sigma (h_i^{\mu })\sigma (h_i^{\nu }).$$

$$⟨D^2⟩=\frac{1}{N_1^{\eta }{N}_{+}^{2\delta }}\sum _{\mu \nu }^P{z}_{\mu }{z}_{\nu }\sum _i^{N_1}⟨\sigma (h_i^{\mu })\sigma (h_i^{\nu })⟩.$$

$$⟨\sigma (h_i^{\mu })\sigma (h_i^{\nu })⟩=K_{\mu \nu }-\frac{\bar{Q}}{{\lambda }_1{N}_1}\sum _{\lambda \delta }{K}_{\mu \lambda }{K}_{\nu \delta }\left[({\tilde{K}}_{(R)}^{-1}{)}_{\lambda \delta }-({\tilde{K}}_{(R)}^{-1}y{)}_{\lambda }({\tilde{K}}_{(R)}^{-1}y{)}_{\delta }\right].$$

$$⟨D^2⟩=\frac{1}{N_1^{\eta -1}{N}_{+}^{2\delta }}\sum _{\mu \nu }^P{z}_{\mu }{z}_{\nu }⟨\sigma (h_i^{\mu })\sigma (h_i^{\nu })⟩,$$

$$⟨D^2⟩$$

$$={\tilde{D}}_1^2-{\tilde{D}}_2^2+{\tilde{D}}_3^2.$$

$${\tilde{D}}_1^2=\frac{1}{N_1^{\eta -1}{N}_{+}^{2\delta }}\sum _{\mu \nu }^P{z}_{\mu }{K}_{\mu \nu }{z}_{\nu },$$

$${\tilde{D}}_2^2=\frac{\bar{Q}}{\lambda }\frac{1}{N_1^{\eta }{N}_{+}^{2\delta }}\sum _{\mu \nu }^P{z}_{\mu }{z}_{\nu }\sum _{\lambda \delta }{K}_{\mu \lambda }{K}_{\nu \delta }({\tilde{K}}_{(R)}^{-1}{)}_{\lambda \delta },$$

$${\tilde{D}}_3^2$$

$$=\frac{\bar{Q}}{{\lambda }_1}\frac{1}{N_1^{\eta }{N}_{+}^{2\delta }}\sum _{\mu \nu }^P{z}_{\mu }{z}_{\nu }\left(\sum _{\lambda \sigma }{K}_{\mu \lambda }({\tilde{K}}_{(R)}^{-1}{)}_{\lambda \sigma }{y}_{\sigma }\right)\left(\sum _{\delta \rho }{K}_{\nu \delta }({\tilde{K}}_{(R)}^{-1}{)}_{\delta \rho }{y}_{\rho }\right)=$$

$$=\frac{{\lambda }_1}{\bar{Q}}\frac{1}{N_1^{\eta }{N}_{+}^{2\delta }}{\left(\sum _{\mu }^P{z}_{\mu }⟨f{⟩}_{\mu }\right)}^2,$$

$$\sum _{\mu }^P{z}_{\mu }⟨f_{\mu }⟩$$

$$=N_{+}\left({\bar{f}}_{+}-{\gamma }^{\frac{\delta -1}{\delta }}{\bar{f}}_{-}\right).$$

$$P=\left(1+\frac{1}{\sqrt[\delta ]{\gamma }}\right)N_{+}=\tilde{\gamma }{N}_{+},$$

$${\tilde{D}}_1^2=\frac{{\alpha }^{\eta -1}{\tilde{\gamma }}^{2\delta }}{P^{\eta +2\delta -1}}\sum _{\mu \nu }^P{z}_{\mu }{K}_{\mu \nu }{z}_{\nu },$$

$${\tilde{D}}_2^2=\frac{\bar{Q}}{{\lambda }_1}\frac{{\alpha }^{\eta }{\tilde{\gamma }}^{2\delta }}{P^{\eta +2\delta }}\sum _{\mu \nu }^P{z}_{\mu }{z}_{\nu }\sum _{\lambda \delta }{K}_{\mu \lambda }{K}_{\nu \delta }({\tilde{K}}_{(R)}^{-1}{)}_{\lambda \delta },$$

$${\tilde{D}}_3^2=\frac{{\lambda }_1{\alpha }^{\eta }{\tilde{\gamma }}^{2\delta -2}}{\bar{Q}}\frac{1}{P^{\eta +2\delta -2}}{\left({\bar{f}}_{+}-{\gamma }^{\frac{\delta -1}{\delta }}{\bar{f}}_{-}\right)}^2.$$

$$⟨D^2⟩=\frac{1}{\alpha }{\left(\frac{\gamma +1}{\gamma }\right)}^2\frac{1}{P}\sum _{\mu \nu }^P{z}_{\mu }{K}_{\mu \nu }{z}_{\nu }+\frac{{\lambda }_1}{\bar{Q}}({\bar{f}}_{+}-{\bar{f}}_{-}{)}^2.$$

$$⟨D^2⟩=\frac{{\gamma }^2+1}{{\gamma }^2}\frac{1}{P}\sum _{\mu \nu }{z}_{\mu }{K}_{\mu \nu }{z}_{\nu }+\alpha \frac{{\lambda }_1}{\bar{Q}}\frac{{\gamma }^2}{{\gamma }^2+1}{\left({\bar{f}}_{+}-\frac{1}{\gamma }{\bar{f}}_{-}\right)}^2$$

$$\mathcal{Z}(\{J\})=\int dvdW{e}^{-\frac{\beta }{2}|f-y{|}^2}{e}^{-\frac{{\lambda }_0}{2}||W|{|}^2}{e}^{-\frac{{\lambda }_1}{2}||v|{|}^2}{e}^{-\frac{{\lambda }_0}{2}{\left({\sum }_{k=1}^{N_1}{J}_k{W}_{1k}\right)}^2}.$$

$$\langle W_{1k}W_{1h}\rangle=-\frac{1}{\lambda_{0}}\frac{1}{\mathcal{Z}(0)}\frac{\partial^{2}\mathcal{Z}(\{J\})}{\partial J_{k}\partial J_{h}}\Bigg{|}_{J=0}.$$

$$\int \prod _{i\ne 1}^{N_1}d{W}_{ij}\delta \left(h_i^{\mu }-\frac{1}{N_0}\sum _j{W}_{ij}{x}_j^{\mu }\right)e^{-\frac{{\lambda }_0}{2}||W|{|}^2}\int d{W}_{1j}\delta \left(h_1^{\mu }-\frac{1}{N_0}\sum _j{W}_{1j}{x}_j^{\mu }\right)e^{-\frac{{\lambda }_0}{2}||W|{|}^2}{e}^{-\frac{{\lambda }_0}{2}{\left({\sum }_{k=1}^{N_1}{J}_k{W}_{1k}\right)}^2}.$$

Full text

*Introduction—*The empirical success of deep learning is fundamentally linked to the ability of neural networks (NNs) to extract meaningful features from raw data [1, 2, 3]. Nevertheless, the mechanistic definition of such process, referred to as feature learning, remains debated. Even for simple models, such as standard-scaled fully-connected (FC) NNs, there is no agreed-upon set of observables (properties) of their trainable parameters that definitively describe feature learning [4, 5, 6, 7, 8].

Recent progress on the theoretical analysis of NNs has established an equivalence between standard-scaled Bayesian FC NNs and kernel regressions, particularly in two limits. (i) In the infinite-width framework, where the width of the hidden layers $N_{1}$ is much larger than the number of training examples $P$, standard-scaled NNs are known to be mathematically equivalent to fixed Gaussian Processes (GPs), which only depend on the prior statistics of the weights [9, 10, 11, 12, 13, 14, 15, 16, 17]. In this case, training produces an infinitesimal displacement of both the network’s weights and the hidden-layer features, which is enough to fit the dataset examples. (ii) In the less overparameterized proportional regime, where $P/N_{1}=\alpha>0$, one lacks such formal equivalence, yet GPs still come to hand when analysing predictive performance. Specifically, the average generalisation error of the NN is found to be the same as that of a GP regression with a set of free hyperparameters that are fine-tuned to the task at hand. Differently from the infinite-width limit, this GP is not fixed at initialisation, and depends on the posterior statistics of the weights [18, 19, 20, 21, 22, 23, 24, 25, 26, 27]. (This regime allows for the analysis of continual learning and transfer learning as well [28, 29].)

These descriptions in terms of GPs suggest that looking at the predictive distribution may not be enough to characterise feature learning in FC models. In fact, consistent experimental evidence points to the fact that finite, yet overparametrized, FC networks are eventually outperformed by a suitable GP [30] (at least in most computer vision tasks), prompting the identification of a different set of observables to describe feature learning in this context.

A key insight emerges from recent empirical studies of feature dynamics in FC NNs. Experiments show a robust behaviour that occurs during training: collective observables of the features are nonmonotonic in training time, with a quasi-universal inversion point that is consequential for good generalization properties [31, 32]. More precisely, for binary classification problems, the class manifolds (the two sets containing the representations in feature space of data with the same label [33, 34, 35, 36, 37]) initially become well separated, and then approach each other when the network is learning the most “challenging” data samples. These results, in addition to pre-existing numerical evidence [38, 39], indicate that the geometry of class manifolds in feature space is intimately linked to feature extraction.

In this manuscript, we investigate analytically both collective and microscopic equilibrium observables linked to hidden-layer features in Bayesian one-hidden-layer (1HL) FC NNs. In the proportional regime, we compute the posterior average squared distance $\langle D^{2}\rangle$ between class manifolds in a binary classification problem, and analyse the second order statistics of the hidden layer weights. Our results can be summarised in three points. (i) The posterior distance remains unaltered in the infinite-width limit, while it departs from its prior value at proportional width (see sketch in panel (a) of Fig. 1). (ii) The observable $\langle D^{2}\rangle$ is a non-monotonic function of the Gibbs temperature $T$. The nonmonotonicity disappears in the infinite-width limit, or when information about the features is removed from the data by randomising the labels. (iii) The hidden-layer weights become correlated in the proportional regime, undergoing a finite displacement, differently from the infinite-width limit, where the parameter statistics is known to be unaffected by training.

Setting of the learning problem— We consider a training set $\mathcal{D}=\{x^{\mu},y^{\mu}\}_{\mu=1}^{P}$, with $x^{\mu}\in\mathbb{R}^{N_{0}}$ and $y^{\mu}\in\mathbb{R}$, and a 1HL FC network $f(x)=v\cdot\sigma[h(x)]/\sqrt{N_{1}}$, where the pre-activations read $h(x)=wx/\sqrt{N_{0}}$, and $\sigma$ is a pointwise nonlinear activation function. The matrix $w\in\mathbb{R}^{N_{1}\times N_{0}}$ and vector $v\in\mathbb{R}^{N_{1}}$ contain the microscopic degrees of freedom of the system. We refer to them collectively as $\theta=\{v,w\}$.

Multiplying the network’s output by $1/\sqrt{N_{1}}$ corresponds to the standard scaling [4]. For the sake of simplicity, we restrict our analysis to zero-mean activation functions (e.g., erf, tanh). The statistics of $v$ and $w$ in the Bayesian setting are determined by their priors, which we take to be rescaled normals, and by a likelihood function, for which we use the mean-squared error loss. The properties of the Bayesian network are then determined by the partition function

$\mathcal{Z}(X,y)=\int\text{d}\mu(\theta)\,\text{exp}(-\beta||f_{\theta}(X)-y||^{2}/2)$. The shorthand notation $f_{\theta}(X)=(f_{\theta}(x^{\mu}),x^{\mu}\in\mathcal{D})$ indicates the collection of the network’s outputs to the training samples, while the measure d$\mu(\theta)$ indicates integration over the prior on the weights $w\sim\mathcal{N}(\bm{0},\mathds{1}/\lambda_{0})$, $v\sim\mathcal{N}(\bm{0},\mathds{1}/\lambda_{1})$. As shown in [18, 19, 25] for linear activation functions and in [20] for general activations,

the partition function can be evaluated analytically in the proportional-width limit, where $N_{1},P\rightarrow\infty$ with $P/N_{1}=\alpha$ fixed. The solution is obtained via a saddle-point evaluation of an integral, which corresponds to optimizing a scalar-dependent effective action $S(Q)$ through the condition

$S^{\prime}(\bar{Q})=0$.

Importantly, the optimization depends on the training dataset (besides the temperature $T=1/\beta$), so that $\bar{Q}$ contains information on both the input data and their corresponding labels. In the following, we will compute average observables over the posterior Gibbs distribution $P(\theta|X,y)$ associated to $\mathcal{Z}$, denoted with

$\langle\star\rangle=\langle\star\rangle_{P(\theta|X,y)}$.

Collective displacement of the features— As shown in [24], within this setting it is possible to compute the posterior statistics of the features, which are collective variables of the first layer weights: $\sigma(h_{j}^{\mu})=\sigma(\sum_{i_{0}}w_{ji_{0}}x_{i_{0}}^{\mu}/\sqrt{N_{0}})$. This is captured by the average similarity matrix, which reads [24]:

[TABLE]

The NNGP kernel $K_{\mu\nu}$ is defined as the averaged similarity matrix of the features over the prior distribution, $K_{\mu\nu}=\langle\sigma(h^{\mu})\sigma(h^{\nu})\rangle_{P(h|X)}$, and we have defined $\tilde{K}_{(R)}=\mathds{1}/\beta+K_{(R)}$, where $K_{(R)}=\bar{Q}K/\lambda_{1}$ is the renormalized kernel. Both in the proportional and in the infinite-width limit, the posterior corrections $\Delta_{1}^{\mu\nu}$, $\Delta_{2}^{\mu\nu}$ are of order $\mathcal{O}(1/N_{1})$, thus irrelevant whenever $N_{1}\to\infty$. While this fact is evident in the infinite-width limit, where the sums involve a finite number of terms $P$, it requires further explanation when $P\rightarrow\infty$. Indeed, if $\beta\to\infty$, then $\tilde{K}_{(R)}^{-1}$ is equal to $\lambda_{1}K^{-1}/\bar{Q}$. As a result, the first sum gives $\Delta_{1}^{\mu\nu}=K_{\mu\nu}/N_{1}$, and the second gives $\Delta_{2}^{\mu\nu}=(\lambda_{1}/\bar{Q})y_{\mu}y_{\nu}/N_{1}$. On the other hand, if $\beta\to 0$, we have $\tilde{K}_{(R)}^{-1}\sim\beta\mathds{1}\to 0$, which makes both the terms in the sum vanish. For all intermediate values of $\beta$, the system continuously interpolates between these two regimes and the two terms will remain of order $\mathcal{O}(1/N_{1})$.

These considerations point to the fact that the second-order statistics of the features in the proportional-width limit is the same as in the infinite-width limit, where no feature learning happens [9, 40]. We now show that it is still possible to define observables that turn the subleading terms $\Delta_{1}$ and $\Delta_{2}$ into finite corrections in the proportional limit. While the previous results hold for a regression task with general training labels, in what follows we consider binary labels $y^{\mu}=\pm 1$, and training datasets where the data points are equally split into the two different classes of cardinality $P/2$ (the generalization to unbalanced datasets [41, 42] and other pairs of label values is shown in the Supplementary Material (SM)).

The squared distance $D^{2}$, defined as the separation between the mean post-activations of the two different classes, has been used in [31] as a measure of class manifold separation between features:

[TABLE]

where $\bm{h}_{\pm}^{\mu}$ is the pre-activation corresponding to the $\mu$-th sample of the first and second class, respectively. Averaging over the posterior, we obtain:

[TABLE]

with $\Delta_{1}$ and $\Delta_{2}$ as in Eq. (Microscopic and collective signatures of feature learning in neural networks). It is important to emphasize that the $y$’s appearing explicitly in this formula do not originate from the posterior distribution, but rather from the relative signs in the definition (2). The labels stemming from the posterior average are instead embedded within $\Delta_{1}$ and $\Delta_{2}$.

In fact, averaging over the prior yields the first term alone, $\langle D^{2}\rangle_{P(h|X)}=4y^{T}Kt/(\alpha P)$.

When $P\rightarrow\infty$, the scaling of the three terms in Eq. (3) is nontrivial, but it can be discussed in the zero-temperature limit $\beta\rightarrow\infty$, where the expressions simplify. In this case, the second term reduces to $4y^{T}Ky/(\alpha PN_{1})$, while the third simplifies to $\frac{4\lambda_{1}}{\bar{Q}}$. Thus, assuming that $y^{T}Ky/P=\mathcal{O}(1)$, the second term vanishes, and the last term provides a finite, feature-dependent correction to the distance averaged over the prior. In general, the assumption is numerically satisfied in realistic regimes for computer vision tasks where the dataset is noisy enough (e.g., CIFAR10, see SM). In such cases, the large but finite-size regime, where the theory is applicable [26], creates a scenario in which the overall squared distance is primarily determined by the first and third terms. This numerical observation is an exact statement for iid random data: in this case, the components of $K$ are independent random variables.

According to the central limit theorem, $\sum_{\nu}^{P}K_{\mu\nu}y_{\nu}\sim\sqrt{P}$ and the sum over two indices is $\sim P$. In the SM, we provide a numerical study of this scaling and more details about the typical situations where it does not hold. At finite temperature, disregarding the subleading (or effectively small) second term, the squared distance reads:

[TABLE]

where $\langle f_{\pm}^{\mu}\rangle$ are the posterior expected outputs of the NN in the proportional limit [20, 21]. The symbols $\pm$ indicate that the index $\mu$ runs over the kernel evaluated on samples from the first and second class, respectively.

Note that the squared distance diverges in the infinite-width limit $\alpha\rightarrow 0$, since the post activations $\sigma(\bm{h})$ are $N_{1}$-dimensional vectors, and their squared norms are $\mathcal{O}(N_{1})$. We plot $\left<D^{2}\right>$ as a function of the temperature in Fig. 1(b) (and $\alpha\left<D^{2}\right>$ in Fig. 1(c), which is rescaled so that it converges in the infinite-width limit). The dependence of the squared distance on the temperature is nonmonotonic for $\alpha>0$, with a characteristic peak at positive temperature. (An alternative way to normalize the distance is presented in the SM, with similar results as here.) The emergence of the term $(\bar{f}_{+}-\bar{f}_{-})^{2}\lambda_{1}/\bar{Q}$ in the proportional limit signals the network’s ability to learn the features by exploiting contributions from the posterior, which influences their separation. Thus, the strongly nonmonotonic behavior of $\left<D^{2}\right>$ is understood as a consequence of feature learning. This is confirmed by the fact that assigning random labels to the data nearly removes the peak (dashed line in Fig. 1(b)). Furthermore, reducing $\alpha$ towards the infinite-width limit reduces the relative contribution of the posterior with respect to the prior, while simultaneously eliminating the peak (Fig. 1(c)). This is in accordance with the fact that the posterior statistics becomes indistinguishable from that of the prior in the infinite-width limit.

Microscopic pairwise correlations of the weights— In the previous section, we isolated a macroscopic observable that carries information about feature learning at proportional width. Here, we investigate microscopic features, showing that the statistics of the first-layer weights $w$ depends on the data (which does not happen in the infinite-width limit). We compute the two-point function $\langle w_{1k}w_{1h}\rangle$ with respect to the posterior distribution. In the SM, we show that it is possible to use the effective action formalism to express the similarity matrix $\langle w_{1h}w_{1k}\rangle$ in terms of the order parameter $\bar{Q}$:

[TABLE]

$\Delta_{hk}^{\mu\nu}$ is a data-dependent kernel defined in the SM.

In the case of a linear activation function ($\sigma=\mathrm{Id}$), these results hold exactly for finite $N_{1}$ and $P$. This fact is in line with the use of the effective action formalism carried out in [25] for finite-width deep linear networks.

In the infinite-width limit, the terms $\Lambda_{1}$ and $\Lambda_{2}$ vanish,

the components of $w$ become uncorrelated, and averages over the prior and over the posterior coincide. In the proportional case, since the kernel $\Delta_{hk}^{\mu\nu}$ cannot be expressed in simple terms by using $K_{\mu\nu}$, the zero-temperature limit is not useful to investigate the magnitude of the proportional correction terms. In Fig. 2, we show results of numerical simulations performed with a linear activation function. In this case, the kernel is easily expressed through the data as $\Delta_{hk}^{\mu\nu}=-({x_{h}^{\mu}x_{k}^{\nu}+x_{h}^{\nu}x_{k}^{\mu}})/(\lambda_{0}N_{0})$. We note that, while the first correction $\Lambda_{1}$ is negligible, the second $\Lambda_{2}$ contributes to the statistics of the weights with a finite term that depends almost linearly on $\alpha$. The linearity is consistent with a scaling argument that can be invoked also in the context of GD dynamics: in this case, the time derivative of the weights is proportional to $\sqrt{\alpha}$.

Discussion and conclusions—

The squared distance $D^{2}$ is a collective observable of the features that signals non-trivial feature learning in Bayesian FC shallow networks.

Other measures of feature segregation have been found to display interesting behavior empirically [43, 44, 7]; the Bayesian proportional-width framework employed here should be able to capture those as well.

Our computations show that the posterior average $\langle D^{2}\rangle$ is a nonmonotonic function of the temperature $T$. The peak at the inversion point is more pronounced for larger values of $\alpha$, and eventually disappears in the infinite-width limit $\alpha=0$. In the Bayesian setting, the temperature plays the role of a regularizer, acting as early stopping in the optimization dynamics [45]. The behavior of the distance then aligns with the inversion dynamics observed in [31]. In that work, the nonmonotonic trend observed during training is interpreted as a trade-off between the segregation of the two class manifolds and the fine tuning required for classifying the last hard samples.

In this sense, feature learning manifests itself, during the optimization dynamics, through a well-defined transition between an easy and a hard phase of training. The results summarized in Fig. 1 can be seen as the equilibrium counterpart to this non-equilibrium phenomenon. A possible interpretation arises from the observation that increasing the temperature allows the posterior to sample regions of the loss landscape further away from the optima. The temperature at the inversion peak then corresponds to the typical loss values associated with the transition between easy and hard samples.

The (posterior) second-order statistics of the hidden-layer weights develop finite correlations in the proportional limit. This is remarkable. The Gaussian process of the output receives a trivial modification with respect to the infinite-width limit, because the scalar renormalization by $\bar{Q}$ can be reabsorbed in the prior parameter $\lambda_{1}$ [24]). In contrast, the second-order statistics of $w$ depends on the input patterns in a way that cannot be traced back to the infinite-width limit.

Finally, we would like to point out two distinctions between our work and other lines of research in this field. First, our results hold for standard-scaled neural networks. Another possibility would be to consider the mean-field scaling, where microscopic quantifiers of feature learning have been found already in the infinite-width limit [4, 5, 46, 47, 48, 49]. Second, the proportional regime represents a genuinely overparameterized scenario. Therefore, our theory does not help to characterize feature learning closer to the interpolation threshold. This setting, which may also be relevant for modern deep learning applications, has been very recently investigated for Gaussian training inputs [50, 51, 52, 53], leading to more complex feature learning mechanisms than the ones under scrutiny here [54, 55, 56].

Acknowledgements.– P.R. is supported by $\#$NEXTGENERATIONEU (NGEU) and funded by the Ministry of University and Research (MUR), National Recovery and Resilience Plan (NRRP), project MNESYS (PE0000006) “A Multiscale integrated approach to the study of the nervous system in health and disease” (DN. 1553 11.10.2022). R.P. is funded by the MUR, project PRIN 2022HSKLK9.

Supplementary Material

Derivation of the average squared distance

In the main text, we mentioned that the squared distance can be defined to be finite also in the infinite-width regime and in the case of unbalanced datasets, where each class has a different number of data points. A more general definition can be considered:

[TABLE]

where $\eta\geq 0$, $\delta>0$, and $N_{+}$ and $N_{-}$ (with $N_{+}+N_{-}=P$) denote the number of training data points belonging to the first and second class, respectively. Assuming that $N_{+}$ and $N_{-}$ represent non-negligible fractions of the total number of points $P$ — which is large in the proportional limit — we can define $\gamma=\left(\frac{N_{+}}{N_{-}}\right)^{\delta}$ and express the distance as

[TABLE]

where the matrix and the vectors are to be understood in block form. By defining the vector $z$, with components

[TABLE]

it is possible to write the block matrix as a tensor product whose components are $z_{\mu}z_{\nu}$, which allows writing the squared distance as

[TABLE]

Note that in the case of a balanced dataset with $N_{+}=N_{-}$, we have $\gamma=1$ and $z_{\mu}=\pm 1$, which are the values assumed by the labels $y_{\mu}$, as in the case of the main text. Averaging over the posterior distribution, we have:

[TABLE]

The averaged similarity matrix has been computed in [24], which provides an explicit expression for the averaged similarity matrix:

[TABLE]

Since the averaged similarity matrix does not depend on the index $i$, we can trivially perform the internal sum to get:

[TABLE]

which can be expressed as

[TABLE]

As mentioned in the main text, the terms that come from the posterior distribution (which involve $y$) are clearly distinguished from the ones which come from the relative minus sign of the norm (represented by $z$). We defined

[TABLE]

[TABLE]

[TABLE]

where $\langle f_{\mu}\rangle$ is defined as in the main text $\langle f_{\mu}\rangle=\sum_{\lambda,\sigma=1}^{P}(K_{(R)})_{\mu\lambda}(\tilde{K}^{-1}_{(R)})_{\lambda\sigma}y_{\sigma}$ and represents the expected output of the Gaussian Process associated with the Neural Network in the proportional limit. It can be observed that the expected output, although defined through a sum involving an increasing number of terms in the proportional limit, has finite size. In fact, if $\beta\rightarrow\infty$, then $\tilde{K}{(R)}^{-1}$ becomes proportional to $K^{-1}$. The result is thus a delta, and the output satisfies $\langle f_{\mu}\rangle=y_{\mu}$. On the other hand, if $\beta\rightarrow 0$, then $\tilde{K}{(R)}^{-1}\sim\beta\mathds{1}\rightarrow 0$, which drives the outputs to zero. One can numerically check that all intermediate values of $\beta$ smoothly interpolate between these two regimes, while keeping $\langle f_{\mu}\rangle$ of finite size. With this in place, we can proceed further by writing

[TABLE]

We defined $\bar{f}_{\pm}^{\mu}=1/N_{\pm}\sum_{\mu}\langle f_{\pm}^{\mu}\rangle$. Since in the proportional limit all the quantities $N_{1},P,N_{+},N_{-}\rightarrow\infty$ at the same rate, it is useful to express everything as a function of $P$. From the relations $N_{1}=P/\alpha$ and

[TABLE]

we can explicitly write the three contributions to the distance:

[TABLE]

[TABLE]

[TABLE]

The zero temperature limit can be used to justify the different scalings in $P$ of the three quantities. When $T\rightarrow 0$, the internal sum of the second contribution returns $K_{\mu\nu}$, telling us that $\tilde{D}_{2}^{2}$ can be neglected. Under the assumption that $1/P\sum_{\mu\nu}z_{\mu}K_{\mu\nu}z_{\nu}\sim\mathcal{O}(1)$, the term $\tilde{D}_{1}^{2}$ is of the same order of $\tilde{D}_{3}^{2}$, making it a feature dependent correction that reflects the fact that the posterior is learning through the labels. As mentioned in the main text, the assumption is true in the case of Gaussian data, where the kernel has independent entries, which make the sum $\sum_{\nu}^{P}K_{\mu\nu}z_{\nu}\sim\sqrt{P}$ and $1/P\sum_{\mu\nu}z_{\mu}K_{\mu\nu}z_{\nu}=\mathcal{O}(1)$ in virtue of the central limit theorem. In the case of fairly noisy data, it is reasonable to expect the term to remain finite also for large values of $P$. As shown in the graphs below ($\eta=0$, $\delta=1$), this is the case for CIFAR10 dataset.

This type of argument holds for every choice of the scaling that meets the condition $\eta+2\delta-2=0$. One possible choice, which is the one presented in the main text, consists in selecting $\eta=0$, $\delta=1$, which returns (disregarding the second contribution):

[TABLE]

As argued in the main text, the divergence in the infinite-width limit $\alpha\rightarrow\infty$ is completely understood in terms of the original definition of the distance, which diverges as the norm of a vector with a growing number of components. As an alternative, the definition with $\eta=1$ and $\delta=1/2$ returns a well-defined infinite-width limit where the distance reads:

[TABLE]

Note that in both cases the relative magnitude between the prior term and the posterior correction is proportional to $\alpha$. Because of this, in the infinite-width limit the prior term is dominant and the posterior correction can be disregarded, returning a distance that does not depend on the labels.

Derivation of the second order statistics of the weights

To compute the second order statistics of the internal weights, it is possible to introduce a partition function with a source term:

[TABLE]

With this definition, we have

[TABLE]

Note that, being interested in mean values, it is possible not to take care of the normalizations (unless they are $J$-dependent) while computing the partition function. Considering the inner integral, we separate the contribution of the terms with $i\neq 1$ from that with $i=1$ (omitting the products over the indices that are being integrated in an obvious way).

[TABLE]

After the integration, the previous expression assumes the form:

[TABLE]

where we defined $D(J)$ as:

[TABLE]

The inverse is computed by employing the Sherman-Morrison formula. Note that we defined $J^{2}=\sum_{k}J_{k}^{2}$ and that $\text{Det}(\delta_{ij}+J_{i}J_{j})=1+J^{2}$. That said, the partition function is simplified by the introduction of an integral representation of a Dirac delta function $\prod_{\mu}\delta\left(s^{\mu}-f(v,h^{\mu})\right)$. Integrating the read-out weights, we obtain:

[TABLE]

where the Kernel $K(J)$ is defined by

[TABLE]

A tedious but straightforward computation of the derivatives yields

[TABLE]

where the second derivative of the Kernel assumes the form of

[TABLE]

The Fourier variable $\bar{h}$ can be integrated out by means of a translation, leaving

[TABLE]

Note that the expression appears to be ill-defined since it depends on the single components of the inverse of the covariance matrix, which is non-existent as soon as $P>N_{0}$, which is typically the case of interest. It turns out that it is not the case and that the expression can be recast in such a way that the dependence on the single components of $C^{-1}$ is removed. Note also that the presence of $C^{-1}$ in the Gaussian measure associated with the variable $h$ does not represent an issue since the distribution is well defined through its Fourier transform. To show the independence of the previous expression on $C^{-1}$, we first note that the integrals define two different kernels: while the first is the usual NNGP kernel, which we denote $K_{\mu\nu}$, the second depends on more indices and is defined as

[TABLE]

To recast this expression, we express it in terms of the NNGP kernel. To do so, we introduce a source such that:

[TABLE]

After completing the square and performing a translation, the kernel is brought in the form of:

[TABLE]

The computation of the derivatives returns:

[TABLE]

where $\sigma^{\prime}_{\mu}$ and $\sigma^{\prime\prime}_{\mu}$ are the first and second derivatives of the activation function evaluated on $h^{\mu}$. Inserting this expression in Eq. (Derivation of the second order statistics of the weights), the sums over the Greek indices remove the dependence on $C^{-1}$:

[TABLE]

which is well defined also in case of $P>N_{0}$. Defining

[TABLE]

[TABLE]

we denote for convenience \Big{(}\partial_{h}\partial_{k}K(J)\Big{|}_{J=0}\Big{)}_{\mu\nu}=\Delta_{hk}^{\mu\nu}, with:

[TABLE]

It is now possible to continue with the computation of the second order statistics. From the definition and the previous results, we have:

[TABLE]

From Eq. (Derivation of the second order statistics of the weights), by plugging $J=0$ one can check that the integral in the first term is exactly $\mathcal{Z}(0)$, so that the first contribution to the expectation value is simply $\delta_{hk}/\lambda_{0}$. To make further progress with the second term, we use the integral definition of the Gamma function (the Gamma itself plays the role of a normalization constant and is not reported):

[TABLE]

The second line is obtained by rescaling the integration variable $\tilde{Q}=QN_{1}/2$. The integral in the second term reads:

[TABLE]

where the effective action $S(Q)$ is defined as

[TABLE]

with

[TABLE]

By introducing an integration variable $Q$ in the same way as it was done in the previous lines, one can easily check that the partition function $\mathcal{Z}(0)=\int_{Q>0}\text{exp}[-N_{1}S(Q)/2]$. Because of that, in the proportional limit where $N_{1}$ is large, the integral is dominated by the saddle-point contribution ($\bar{Q}$ such that $\partial_{Q}S(Q)|_{Q=\bar{Q}}=0$), returning:

[TABLE]

As mentioned in the main text, the first term is obtained as the expectation value over the prior distribution of the weights $W$: $\langle W_{1h}W_{1k}\rangle_{P(W)}=\delta_{hk}/\lambda_{0}$ or over the posterior in the infinite-width limit. The latter statement can be easily checked by noting that when $P$ is finite the term proportional to $1/N_{1}\sum^{P}\rightarrow 0$. In the proportional limit, as argued in the main text, this term can be nontrivial e bring finite corrections to the two-point function.

Consistency check for linear activation function

In the case of a linear activation function $\sigma=Id$, we can compute

[TABLE]

and the result must match the one presented in Eq. (13), which is obtained by an independent computation. First of all, the expression of $\Delta_{hk}^{\mu\nu}$ simplifies as long as $\sigma^{\prime}(h^{\mu})=1$ and $\sigma^{\prime\prime}(h^{\mu})=0$, returning

[TABLE]

Furthermore, $K_{\mu\nu}=C_{\mu\nu}$. Plugging the previous equation into the expectation value:

[TABLE]

where we used the notation $\tilde{C}_{(R)}=\mathds{1}/\beta+\bar{Q}C/\lambda_{1}$. The computation of $\langle h_{i}^{\mu}h_{i}^{\nu}\rangle$ involves two further sums over the indices $h,k$. The terms involved return:

[TABLE]

[TABLE]

With this in mind, plugging the previous expressions in Eq. (50), we obtain:

[TABLE]

which is the equivalent of Eq. (13) in the case of linear activation function.