
==== Front
ArXiv
ArXiv
arxiv
ArXiv
2331-8422
Cornell University

arXiv:2409.03749v1
2409.03749
1
preprint
Article
Dynamics of Supervised and Reinforcement Learning in the Non-Linear Perceptron
Schmid Christian Institute of Neuroscience, University of Oregon

Murray James M. Institute of Neuroscience, University of Oregon

cschmid9@uoregon.edu
5 9 2024
arXiv:2409.03749v1https://creativecommons.org/licenses/by/4.0/ This work is licensed under a Creative Commons Attribution 4.0 International License, which allows reusers to distribute, remix, adapt, and build upon the material in any medium or format, so long as attribution is given to the creator. The license allows for commercial use.
nihpp-2409.03749v1.pdf
The ability of a brain or a neural network to efficiently learn depends crucially on both the task structure and the learning rule. Previous works have analyzed the dynamical equations describing learning in the relatively simplified context of the perceptron under assumptions of a student-teacher framework or a linearized output. While these assumptions have facilitated theoretical understanding, they have precluded a detailed understanding of the roles of the nonlinearity and input-data distribution in determining the learning dynamics, limiting the applicability of the theories to real biological or artificial neural networks. Here, we use a stochastic-process approach to derive flow equations describing learning, applying this framework to the case of a nonlinear perceptron performing binary classification. We characterize the effects of the learning rule (supervised or reinforcement learning, SL/RL) and input-data distribution on the perceptron’s learning curve and the forgetting curve as subsequent tasks are learned. In particular, we find that the input-data noise differently affects the learning speed under SL vs. RL, as well as determines how quickly learning of a task is overwritten by subsequent learning. Additionally, we verify our approach with real data using the MNIST dataset. This approach points a way toward analyzing learning dynamics for more-complex circuit architectures.
==== Body
pmcIntroduction

Learning, which is typically implemented in both biological and artificial neural networks with iterative update rules that are noisy due to the noisiness of input data and possibly of the update rule itself, is characterized by stochastic dynamics. Understanding these dynamics and how they are affected by task structure, learning rule, and neural-circuit architecture is an important goal for designing efficient artificial neural networks (ANNs), as well as for gaining insight into the means by which the brain’s neural circuits implement learning.

As a step toward developing a full mathematical characterization of the dynamics of learning for multilayer ANNs solving complex tasks, recent work has made progress by making simplifying assumptions about the task structure and/or neural-circuit architecture. One fruitful approach has been to study learning dynamics in what is perhaps the simplest non-trivial ANN architecture: the individual perceptron. Even with this simplification, however, fully characterizing the mathematics of learning has been challenging for complex tasks, and further simplifications have been required.

One approach has been to analyze learning in the student-teacher framework [Gardner and Derrida, 1989, Seung et al., 1992], in which a student perceptron learns to produce an input-to-output mapping that approximates that of a teacher perceptron. This has led to insights about the differences in learning dynamics between different types of learning rules (e.g., supervised and reinforcement learning) [Werfel et al., 2003, Züge et al., 2023, Patel et al., 2023]. Such insights are highly relevant for neuroscience, where a longstanding goal has been to infer the learning mechanisms that are used in the brain [Lim et al., 2015, Nayebi et al., 2020, Portes et al., 2022, Humphreys et al., 2022, Mehta et al., 2023, Payeur et al., 2023]. However, by construction, the student-teacher setup in the perceptron only applies to input-output mappings that are linearly separable, which is seldom the case in practice. Another approach has been to study learning dynamics in the linearized perceptron [Werfel et al., 2003, Mignacco et al., 2020, Bordelon and Pehlevan, 2022], which enables exact solutions even for structured input data distributions that are not linearly separable. However, the dynamics of learning in nonlinear neural networks—even very simple ones—performing classification tasks are not fully understood. Further, whether and how the dynamics of learning might differ under different learning rules in such settings has not been investigated.

Here, we take a stochastic-process approach (similar to Yaida [2018] and Murray and Escola [2020]) to derive flow equations describing learning in the finite-dimensional nonlinear perceptron trained in a binary classification task (Fig. 1). These results are compared for two different online learning rules: supervised learning (SL, which corresponds to logistic regression) and reinforcement learning (RL). We characterize the effects of the input-data distribution on the learning curve, finding that, for SL but not for RL, noise along the coding direction slows down learning, while noise orthogonal to the coding direction speeds up learning. In addition, we verify our approach by training a nonlinear perceptron on the MNIST dataset. Finally, applying the approach to continual learning, we quantify how the input noise and learning rule affect the rate at which old classifications are forgotten as new ones are learned. Together, these results establish the validity of the approach in a simplified context and provide a path toward analyzing learning dynamics for more-complex tasks and architectures.

Stochastic-process approach for describing weight evolution

We consider a general iterative update rule of the form (1) wit+δt-wit=ηfiwt,

where wt∈Rn for arbitrary n>0, and η is the learning rate. The stochastic update term fi on the right-hand side is drawn from a probability distribution—it depends on the weights themselves, as well as the input to the network and, potentially, output noise. Starting from this update equation, our goal is to derive an expression characterizing the evolution of the probability distribution of the weights, p(w,t) (cf. Fig. 1C). We assume that fi(w) does not explicitly depend on η, and that all the moments ⟨fik⟩L,k=1,2,…, where ⟨⋅⟩L denotes an average over the noise in the update equation (1) (including the input distribution as well as, potentially, output noise), exist as smooth functions of w.

Given the stochastic process defined by (1), the probability distribution at time t+δt given the distribution at time t is (2) pw,t+δt=∫dw′pw,t+δt∣w′,tpw′,t.

Denoting the weight update as δw≔w-w′, the integrand in this equation can be written as (3) pw,t+δt∣w′,tpw′,t=pw+δw-δw,t+δtw-δw,tpw-δw,t.

Changing the integration variable to δw and performing a Taylor expansion in δw, the right-hand side of (2) yields (4) ∫dw′pw,t+δt∣w′,tpw′,t=pw,t-∑i∂∂wiαiwpw,t+12∑ij∂2∂wi∂wjβijwpw,t+Oδw3,

where (5) αi(w)=∫dδwδwip(w+δw,t+δt∣w,t)

and (6) βijw=∫dδwδwiδwjpw+δw,t+δtw,t.

Here, we assumed that the probability distribution describing the weight updates f has bounded derivatives with respect to w.

Although (2) is only defined for discrete time steps, we assume a continuous probability density p(w,t) interpolates between the updates and exists as a smooth function for all values of t. We can then expand the left-hand side of (2) to obtain (7) pw,t+δt=pw,t+δt∂∂tpw,t+Oδt2.

For the iterative update rules that we will consider, we have δw∝η, where η is a learning rate. In order to take a continuous-time limit, we let η≔δt and take the limit δt→0. For the general learning rule (1), the coefficients in (4) have the form (8) αiw=fiL,βijw=fifjL,

where ⟨⋅⟩L denotes an average over the noise in the update equation (1) (including the input distribution as well as, potentially, output noise). Thus, we find (9) η∂p∂tw,t=-η∑i∂∂wipw,tfiL+𝒪η2.

Finding the the p(w,t) that solves this equation cannot in general be done exactly when fi is nonlinear. However, by multiplying (9) with powers of w and integrating, as well as expanding in w-⟨w⟩, where ⟨⋅⟩ denotes the average with respect to p(w,t), we can derive a system of equations for the moments of p(w,t) [Risken, 1996]. As we derive in the appendix, this gives the following expressions for the first two moments up to 𝒪((w-⟨w⟩)3): (10) ddtwi=1+12∑k,lCovwk,wl∂k∂lfiLw,

(11) ddtCovwi,wj=∑k[Covwi,wk∂kfjL(⟨w⟩)+Covwj,wk∂kfiL(⟨w⟩)].

Together, these equations characterize the flow of pw,t for a general iterative learning algorithm in a general ANN architecture.

Learning dynamics in the nonlinear perceptron

While the above approach is general and could be applied to any iterative learning algorithm for any ANN architecture, for the remainder of this work we will focus on its application to the nonlinear perceptron (Fig. 1A), a one-layer neural network that receives an input x∈RN, multiplies it with a weight vector w∈RN, and produces an output yˆ. The task we study is a binary Gaussian classification task, in which the model is presented with samples x drawn from two distributions p(x∣y) with labels y=±1, where p(y=±1)=12. Each p(x∣y) is given by a multivariate normal distribution with x~𝒩μy,Σy (Fig. 1B). We analyze both the case of SL with deterministic output, for which yˆ=ϕ(w⋅x), as well as RL, for which the stochastic output is given by π(yˆ=±1)=ϕ(±w⋅x), where ϕ is the logistic sigmoid function. The goal of the model is to output a label yˆ that closely matches the ground truth y when given an input x.

Derivation of the flow equations

The supervised learning rule we consider is regularized stochastic gradient descent for a binary cross-entropy loss, which results in the weight update rule (12) f(w)i=(y˜-yˆ)xi-λwi,

where y˜=12(y+1)∈0,1 is the shifted input label, and λ is the regularization hyperparameter. This learning rule describes online logistic regression.

For reinforcement learning, we use the REINFORCE policy-gradient rule with reward baseline [Williams, 1992, Sutton and Barto, 2018]: (13) f(w)i=yˆδϕ(-yˆw⋅x)xi-λwi.

Here δ=yyˆ-⟨yyˆ⟩ is the reward prediction error, and yˆ is the stochastic output of the perceptron with probability π(yˆ=±1)=ϕ(±w⋅x). To facilitate mathematical feasibility, we replace the perceptron activation function by a shifted error function ϕ(z)=121+Erfπ4z.

We first derive the learning dynamics for stochastic gradient descent. We assume that the initial condition is uniquely specified, with p(w,0)=δw-w0. In this case, the weight covariance will be zero, and the flow equations (10) simply reduce to (14) ddtwi=fiLw=⟨w⟩.

To make the formulas more concise, we set λ=0. It can be reintroduced by simply adding the term -λw. We then get (15) fiL(w)=(y˜-ϕ(w⋅x))xix,y=12(1-ϕ(w⋅x))xix~𝒩μ+,Σ+-12ϕ(w⋅x)xix~𝒩μ-,Σ-=12μi+1-Φa+1+b+2-1212πΣ+⋅w˜i1+b+2e-a+22(1+b+2)-12μi-Φa-1+b-2-1212πΣ-⋅w˜i1+b-2e-a-22(1+b-2).

Here, Φ(z)=12(1+Erf(z/2))=ϕ(z⋅8/π) is the cumulative distribution function of the standard normal distribution. To simplify notation, we have introduced w˜=w⋅π/8, as well as the quantities (16) ay=μy⋅w˜,

(17) by=w˜TΣyw˜.

To aid interpretation of these results, we assume that μ±=±μ and Σ=σ2I. Then (15) implies (18) ddt⟨μ⋅w⟩=|μ|21-Φμ⋅w˜1+σ2|w˜|2-12πσ2μ⋅w˜1+σ2|w˜|2e-(μ⋅w˜)22(1+σ2|w˜|2)w=⟨w⟩

as well as (19) ddt|⟨w⟩|2=2w⋅μ1-Φμ⋅w˜1+σ2|w˜|2-12σ2|w|21+σ2|w˜|2e-(μ⋅w˜)22(1+σ2|w˜|2)w=⟨w⟩

An interpretation of (18) is that the first term pushes the weight vector in the decoding direction, while the second term acts as a regularization, whereby the cross-entropy loss penalizes misclassifications more as μ⋅w increases. An increase in the input noise leads to a higher overlap of the distributions, which means that even the Bayes-optimal classifier will make more mistakes.

For RL, we need to calculate (20) fiL(w)=yˆδϕ(-yˆw⋅x)xix,y,yˆ=ϕ(-w⋅x)ϕ(w⋅x)xix~𝒩μ+,Σ+-ϕ(-w⋅x)ϕ(w⋅x)xix~𝒩μ-,Σ-=Σ+⋅w˜i2π1+b+2e-a+22(1+b+2)1-2Φa+1+b+21+2b+2+2μi+Ta+1+b+2,11+2b+2-Σ-⋅w˜i2π1+b-2e-a-22(1+b-2)1-2Φa-1+b-21+2b-2-2μi-Ta-1+b-2,11+2b-2.

Here, T(⋅,⋅) is Owen’s T function: (21) Th,a=12π∫0ae-12h2(1+x2)1+x2dx.

As for supervised learning, we can simplify this expression for isotropic distributions with means ±μ and get (22) ddt⟨μ⋅w⟩=|μ|24Tμ⋅w˜1+σ2|w˜|2,11+2σ2|w˜|2-12π2σ2μ⋅w˜1+σ2|w˜|2e-(μ⋅w˜˜)22(1+σ2|w˜|2)Erfμ⋅w˜1+σ2|w˜|22+4σ2|w˜|2w=⟨w⟩

and (23) ddt|⟨w⟩|2=8w⋅μTμ⋅w˜1+σ2|w˜|2,11+2σ2|w˜|2-122σ2|w|21+σ2|w˜|2e-(μ⋅w˜)22(1+σ2|w˜|2)Erfμ⋅w˜1+σ2|w˜|22+4σ2|w˜|2w=w.

As we show in the appendix, and as demonstrated in Fig. 2, the flow equations for both SL and RL have a unique, globally stable fixed point whenever λ>0 or the input noise σ>0 (Fig. 2A,B). The solutions of (15) and (20) exhibit agreement with learning curves obtained by direct simulation of (1) (Fig. 2C,D), where the small remaining discrepancy arises from the fact that, for the simulation, we used a standard logistic sigmoid function instead of the error function sigmoid curve used for the analytical calculations. We also see that the asymptotic weight norm decreases approximately linearly with lnλ (Fig. 2E). Of particular note is the observation that, perhaps counter-intuitively, higher levels of noise appear to lead to faster learning for SL, though the effect is more ambiguous in the case of RL. This will be analyzed in more detail in the following section.

Impact of noise on learning time

We next investigate the effect of different types of input noise on the dynamics of learning and whether differences arise for the supervised and reinforcement algorithms. We begin with the case of isotropic input noise, with Σ=σ2I and means ±μ with |μ|=1. In this case, the optimal alignment μ⋅⟨w⟩|⟨w⟩| of 1 is always reached asymptotically, so we focus on how quickly this value is approached as a function of the input noise.

In the case of SL, analytically analyzing the logarithmic derivative of the alignment between μ and ⟨w⟩ yields a flow equation of the form (24) ddtlogμ⋅ww=giso(μ,w)+σ2hiso(μ,w)2+𝒪σ4,

where giso and hiso do not depend on σ. Thus, the higher the input noise, the faster the task is learned. The analogous relationship for RL is indeterminate, such that input noise may either speed up or slow down learning in this case, depending on the parameters. As is illustrated in Fig. 3A, numerical integration of the flow equations reveals qualitatively distinct trends for the dependence of learning speed on noise.

Anisotropic input distributions

To analyze the case of anisotropic input noise, we divide the total noise into two components: a component σ‖2=1+ε in the direction of μ and the noise σ⊥2=1-ε orthogonal to it, while keeping the total noise σ‖2+σ⊥2 fixed (Fig. 3B). For both SL and RL, we find that learning slows down when the noise is shifted to the decoding direction and speeds up when it is shifted to orthogonal directions (Fig. 3C). To confirm this analysis analytically, we calculate the logarithmic derivative of the alignment between μ and ⟨w⟩ and find (25) ddtlogμ⋅ww=gan(μ,w)+εhan(μ,w)2+𝒪ε2,

where gan and han are independent of ε. From this expression, we see that, at least to leading order in ε, noise anisotropy orthogonal to the decoding direction tends to increase the speed of learning, while anisotropy along the decoding direction tends to decrease the speed of learning. This is in apparent contrast to a recent study in two-layer networks, where input variance along the task-relevant dimension was found to increase the speed of learning [Saxe et al., 2019]. The reason for these seemingly opposite results is because, in the the task studied in that work, variance along the coding direction is a signal that facilitates learning, while, in our case of binary classification, variance along the coding direction is noise that impairs learning.

Input noise covariance

So far, we have assumed that the initial weight distribution, which can be thought of as characterizing an ensemble of networks with different initializations, is specified deterministically, i.e. p(w,0)=δw-w0. In this case, according to (11), the covariance of w will remain zero at later times. If training is instead initiated with a distribution p(w,0) having nonzero covariance, then we can ask how this covariance evolves with training—in particular, whether the covariance of this distribution diverges, converges to 0, or approaches a finite value as t→∞.

This calculation can be easily performed in the limit σ→0 where the inputs are just x=±μ. Then (15) simply becomes (26) fiLw=μi1-ϕμ⋅w-λwi,

and (11) implies that (27) ddttr(Cov(w))=-e-π(μ⋅w)2/164μTCov(w)μ-2λtr(Cov(w)).

Since Cov(w) is positive semidefinite, both terms on the right-hand side of (27) are always nonpositive for λ>0 and lead to exponential decay of tr(Cov(w)), so the eigenvalues of Cov(w) approach zero. Thus, the covariance of the distribution p(w,t) vanishes as t→∞ (Fig. 4A).

The same calculation can be performed for the RL algorithm, again with the result that tr(Cov(w))→0 0 as t→∞ whenever λ>0 (Fig. 4B). As can be seen in Fig. 4, the total variance continues to decay to zero upon including input noise (in the η→0 limit we are working in), with the decay speeding up as the noise is increased.

Application to real tasks

In order to test whether the theoretical equations derived above apply to realistic input data, we next train a perceptron with stochastic gradient descent to perform binary classification with cross-entropy loss on the MNIST dataset (Fig. 5A). To obtain suitable input representations, the images corresponding to the digits 0 and 1 are first convolved with a set of 1440 Gabor filters [Haghighat et al., 2015]. (In the appendix, we perform the same analysis on the raw MNIST data without the Gabor convolution and obtain similar results.) We then model these two input classes as multivariate Gaussians with covariances Σ0,1 and means μ0,1 (or ±μ after a translation). The evolution of the weight vector during training is found by numerically integrating (15). To quantify the test accuracy during training, an approximation of the expected error at each time step is derived by integrating the Gaussian approximations to the two input distributions up to the hyperplane orthogonal to the weight vector. As can be seen in Fig. 5B, this theoretically derived learning curve closely matches the actual generalization performance of the trained classifier on the hold-out set.

To further illustrate that the flow equations capture non-trivial aspects of the learning dynamics, Fig. 5C shows the alignment of w with μ, which continues to evolve after the task has been learned. The close alignment of the experimental results with the analytical predictions shows that the flow equations can capture learning dynamics in a realistic task with input data distributions that are not necessarily Gaussian.

Continual learning

In addition to describing the dynamics of learning a single task, the flow equations derived above can also describe the learning and forgetting of multiple tasks. In continual learning, natural and artificial agents struggle with catastrophic forgetting, which causes older learning to be lost as it is overwritten with newer learning [Hadsell et al., 2020, Kudithipudi et al., 2022, Flesch et al., 2023]. Here, we ask how the number of tasks that can be remembered by the perceptron depends on the level of noise and the learning algorithm. The weights are first trained on Task 0, with input distribution defined by μ=μ0 and Σ=σ2I, until the alignment of w with μ0 has reached 80%. We then train on subsequent tasks μ=μ1,μ2,…. This yields a forgetting curve that decays exponentially with the number of tasks, as shown in the simulation results in Fig. 6. The decay constant does not significantly depend on the learning algorithm being used, but we observe that a higher input noise leads to faster forgetting. Together with the results in the preceding subsections, this hints toward a trade-off between the learning speed and forgetting of previously learned tasks as the amount of input noise is varied.

Discussion

In this work, we have used a stochastic-process framework to derive the dynamical equations describing learning in the nonlinear perceptron performing binary classification. We have quantified how the input noise and learning rule affect the speed of learning and forgetting, in particular finding that greater input noise leads to faster learning for SL but not for RL. Finally, we have verified that our approach captures learning dynamics in an MNIST task that has a more-complex input data distribution. Together, the results characterize ways in which task structure, learning rule, and neural-circuit architecture significantly impact learning dynamics and forgetting rates.

One limitation of our approach is the assumption that the input distributions are multivariate Gaussians, which may not be the case for real datasets. While the agreement between the theoretical and empirical results applied to the MNIST data in Fig. 5 is encouraging in this regard, there may be greater discrepancies in cases where the input distributions are more complex. Indeed, recent work on the nonlinear perceptron has shown that, while the first- and second-order cumulants of the input distribution are learned early in training, later stages of training involve learning beyond-second-order (i.e. non-Gaussian) statistical structure in the input data [Refinetti et al., 2023], suggesting that our theory’s ability to describe late-stage training in complex datasets may be somewhat limited. Another limitation is the choice to neglect higher-order terms in w-⟨w⟩ (Equations (10), (11)) and η (Equation (9)). This may limit the ability to characterize instabilities and noise effects induced by non-infinitesimal learning rates. Future work will be needed to assess these effects.

While other work has approached SGD learning in neural networks within a stochastic-process framework, most of these works have not derived the noise statistics from the noisy update rule (as done here and in Yaida [2018] and Murray and Escola [2020]), but rather have added Gaussian noise to the mean update (e.g. [He et al., 2019, Li et al., 2019, 2021]). While the results for the flow of the weights’ mean ⟨w⟩(t) are the same under both approaches, the approach that we take enables us to additionally derive the flow of the weight covariance. Further, it allows for the possibility of describing effects arising from finite learning rate by including higher-order terms in η from the expansion of (4)—a topic that we will address in an upcoming publication.

In our results on continual learning, we found that only a few tasks could be remembered by the perceptron before being overwritten. This is perhaps somewhat surprising given recent work [Murray and Escola, 2020] showing that the binary perceptron can recall O(N) individual random patterns in a continual-learning setup. This difference may arise in part from the fact that that work used a more efficient, margin-based supervised learning rule [Crammer et al., 2006] rather than the stochastic gradient descent rule used here, as well as the fact that input noise and weight regularization were not included. This difference suggests that there is likely room for significant improvements in continual-learning performance with the setup studied here. This would be another interesting direction for future work, given that recent work has found that nonlinearity can drastically increase the amount of catastrophic forgetting in continual learning [Dominé et al., 2023].

Finally, we speculate that qualitative differences between learning rules such as that shown in Fig. 3 may provide a path for designing experiments to distinguish between learning rules implemented in the brain. More work will be needed, however, to formulate testable experimental predictions for more-realistic learning rules and network architectures. More generally, the approach developed here paves the way for analyzing numerous questions about learning dynamics in more-complex circuit architectures and diverse task structures.

Acknowledgements

We are grateful to Elliott Abe for early collaboration related to this project. Support for this work was provided by NIH-BRAIN award RF1-NS131993.

Derivation of the general evolution equations

In this section, we derive equations (10) and (11) from (9). We start with (28) ∂p∂t(w,t)=-∑j∂∂wj(p(w,t)fjL),

multiply both sides by wi and integrate over w. The left-hand side simply becomes (29) ddtwiw.

For the right-hand side, we can use integration by parts to get (30) -∫dwwi∑j∂∂wj(pw,tfjL)=∫dwpw,tfiL=fiLww.

To evaluate this expectation value, we introduce the mean-zero weight wˆ=w-⟨w⟩, which describes the fluctuations of w around its mean. If we expand fiL(w) to second order in wˆ, (30) becomes (31) fiL(w)w=fiL(⟨w⟩)+∑jwˆj∂jfiL(⟨w⟩)+12∑j,kwˆjwˆk∂j∂kfiL(⟨w⟩)+𝒪wˆ3w=fiL(⟨w⟩)+∑jwˆjw∂jfiL(⟨w⟩)+12∑j,kwˆjwˆkw∂j∂kfiL(⟨w⟩)+𝒪wˆ3=fiLw+12∑j,kCovwk,wj∂j∂kfiLw+𝒪wˆ3,

The derivation of (11) follows analogously.

Derivation of explicit SL and RL flow equations

In order to analyze (10) and (11), we must evaluate the following expectation values: (32) ϕ(w⋅x)xix~𝒩(μ,Σ)andϕ2(w⋅x)xix~𝒩(μ,Σ)

with ϕ(x)=121+Erfπ4x. Without loss of generality, we will calculate these integrals in a coordinate system where w=w1e1. We can then factorize px1,x2,…=pmx1pcx2,…∣x1. The marginal distribution is Gaussian with μm=μ1 and Σm=Σ11, and the conditional distribution is also normal with μci=μi+1Σ11Σ1ix1-μ1, and Σcij=Σij-1Σ11Σ1iΣ1j.

Furthermore, to simplify notations, we introduce w˜=w⋅π/8, as well as the quantities (33) a=μ⋅w˜,

(34) b=w˜TΣw˜.

Let’s first calculate ⟨ϕ(w⋅x)⟩x~𝒩(μ,Σ): (35) ⟨ϕ(w⋅x)⟩x~𝒩(μ,Σ)=ϕw1x1x~𝒩(μ,Σ)=12πΣ11∫Rdx1e-12x1-μ12/Σ11ϕw1x1=12π∫Rdue-u2/2ϕw1(μ1+uΣ11)=Φa1+b2,

where Φ(x)=12+12Erf(x/2) is the cumulative distribution function of the standard normal distribution. To evaluate the last line of this and the following integrals, we used the reference [Owen. 1980].

For the integral ϕ(w⋅x)xix~𝒩(μ,Σ), we first do the calculation for i=1: (36) ϕ(w⋅x)x1x~𝒩(μ,Σ)=ϕw1x1x1x~𝒩(μ,Σ)=12πΣ11∫Rdx1e-12x1-μ12/Σ11ϕw1x1x1=12π∫Rdue-u2/2ϕw1(μ1+uΣ11)(μ1+uΣ11)=μ1Φa1+b2+12πΣ11w˜11+b2e-a22(1+b2).

For i≠1, we get (37) ϕ(w⋅x)xix~𝒩(μ,Σ)=ϕw1x1xix~𝒩(μ,Σ)=12πΣ11∫Rdx1e-12x1-μ12/Σ11ϕw1x1μi+Σ1iΣ11x1-μ1=μiΦa1+b2+12πΣi1w˜11+b2e-a22(1+b2).

Thus, for general i and w we can write (38) ϕ(w⋅x)xix~𝒩(μ,Σ)=μiΦa1+b2+12π(Σw˜)i1+b2e-a22(1+b2).

We next calculate (39) ϕ2(w⋅x)x~𝒩(μ,Σ)=ϕ2w1x1x~𝒩(μ,Σ)=12πΣ11∫Rdx1e-12x1-μ12/Σ11ϕ2w1x1=Φa1+b2-2Ta1+b2,11+2b2,

where T stands for Owen’s T function.

Analogously, we can calculate (40) ϕ2w⋅xxix~𝒩μ,Σ=μiϕ2w⋅x+2(Σw˜)i2π1+b2Φa1+b21+2b2e-a22(1+b2).

Fixed point analysis

In this section, we analyze the fixed points of the systems of equations (18) & (19) for SL and (22) & (23) for RL. We will first show the intuitive result that any fixed point w* is maximally aligned with μ, i.e. w*⋅μ=|μ|⋅w*, as long as σ>0. For simplicity, we set the regularization parameter λ=0. Note that for both SL and RL, the flow equations take the form (41) ddt⟨w⟩⋅μ=|μ|2f1μ,w,σ-μ⋅wf2μ,w,σ,12ddt|⟨w⟩|2=μ⋅⟨w⟩f1(μ,w,σ)-|⟨w⟩|2f2(μ,w,σ)

for some functions f1 and f2>0. Also, it’s easy to see that w*⋅μ>0. Thus, a fixed point w* satisfies (42) f1f2=μ⋅w*|μ|2,f1f2=w*2μ⋅w*.

Thus, setting these equal to one another, we find that w*⋅μ=|μ|⋅w* and the two equations reduce to a single equation for w*.

Assume without loss of generality that |μ|=1. For supervised learning, (19) then implies that (43) 0=w*2Erfcw*π16+2πσ2w*2-14σ2w*21+σ2w*2π/8e-|⟨w*⟩|2π16+2πσ2|⟨w*⟩|2,

where Erfc(z)=1-Erf(z) is the complementary error function.

We can factor out the common exponential asymptotics of both terms to get (44) 0=e-|⟨w*⟩|2π16+2πσ2|⟨w*⟩|2e+|⟨w*⟩|2π16+2πσ2|⟨w*⟩|212Erfcw*π16+2πσ2w*2-14σ2w*1+σ2w*2π/8.

Of the two terms in the parentheses, the first term starts at 12 for w*=0 and is monotonically decreasing to the value 12e1/2σ2Erfc12σ≈12πσ-σ3+𝒪σ5, while the second term starts at 0 and increases to the larger value σ2π. Thus, there is a unique fixed point. For RL, uniqueness of the fixed point can be shown using an analogous argument.

This fixed point w* is stable, because for SL, (14) implies that ddt⟨w⟩=-∇w⟨ℒ⟩L, where ℒ is the cross-entropy loss and ⟨⋅⟩L is the average over the input distribution. Thus, ⟨ℒ⟩L is a Lyapunov function for the dynamical system. For RL, the same argument holds by the policy-gradient theorem, where ⟨ℒ⟩L additionally includes an average over the output noise.

Note that, for λ=σ=0, although |⟨w⟩| diverges, a calculation of ddtlog(⟨w⟩⋅μ)2|μ|2|⟨w⟩|2 shows that the alignment between μ and ⟨w⟩ still converges to 1. This result is consistent with those of Ji and Telgarsky [2020].

MNIST details

Figure 7: Comparison of the theory with training on raw MNIST. A: Comparison of the empirical test classification accuracy with the theoretical prediction. B: Just like for the Gabor-filtered inputs, the theory accurately captures non-trivial ongoing learning dynamics.

When applying our theory to the MNIST dataset, we compare SGD applied to actual data (orange curves) with calculated SGD curves from our theory (blue curves). In the main text, we preprocess the data by convolving raw pixel values with a bank of Gabor filters to approximate a more realistic scenario where the binary classifier appears at the end of a convolutional neural network. For the plots in Figure 7, we use raw pixel values to demonstrate that our evolution equations (15) work in general settings without relying on Gabor filter representation. In both cases, we globally translate all input vectors such that the dataset’s mean is zero. We directly evaluate test set accuracy (Figures 5B and 7A) and the correlation of w at each SGD step with the mean μ of digit ‘1’ (Figures 5C and 7B)). To calculate the theoretical accuracy curve, we first numerically solve the differential equations (15) for the mean μ and covariances Σ0,1 obtained from the empirical dataset to derive w(t). We then integrate two multivariate normal distributions with these μ and Σ values in the half-spaces bounded by w(t) (expressible as an error function) and plot the result as the theoretical accuracy curve in Figures 5B and 7A. The tSNE embedding in Figure 5A is included for illustrative purposes only and is not used in calculations.

Experiments

The numerical code implementing the model and performing the analyses was mostly written in JAX [Bradbury et al., 2018], as well as Wolfram Mathematica and SciPy [Virtanen et al., 2020]. For Fig. 2, the flow fields were plotted for the limit of zero input noise and a regularization parameter of λ=0.1. The learning curves are plotted for λ=0 and Σ=σ2I, with σ=0.1 and σ=1. As in all other figures besides Fig. 5, we set |μ|=1. Each experiment was repeated for 10 runs with different random seeds, with the standard deviation indicated as a (barely visible) shaded region. The curve of w* is plotted for the limit of zero input noise. For Fig. 3, we numerically integrated (10) with λ=0. For the total variance plotted in Fig. 4, we set λ=0.1 and integrate the differential equations (11) numerically. For Fig. 5, the Gabor-filtered inputs were created using [Haghighat et al., 2015] with default parameters. The perceptron was trained using SGD with a logistic sigmoid output, and the orange curve in panel B shows test accuracy. The training was repeated 10 times with shuffled data. We set λ=1. For the forgetting curves in Fig. 6, we set λ=10 and the learning rate to η=10-2. The curves shown are averages over 50 different random initializations each. The input dimension is set to N=500. For all other simulations, the learning rate was set to η=10-3 The computations were performed on an NVIDIA Titan Xp GPU, with runtimes of at most a few minutes.

Figure 1: Learning dynamics in the nonlinear perceptron. A: The perceptron, parametrized by weights w, maps an input x to the output yˆ. B: The inputs are drawn from two multivariate normal distributions with labels y = ±1. The weight vector w is orthogonal to the classification boundary. C: Due to the stochasticity inherent in the update equations, the weights are described by the flow of a probability distribution in weight space.

Figure 2: Learning dynamics in a perceptron classification task. A, B: Flow fields determining the weight dynamics with trajectories for different initial conditions for SL (A) and RL (B). C, D: Learning dynamics from simulations closely follow the analytical results for SL (C) and RL (D). E: Dependence of the asymptotic weight norm on the regularization parameter λ.

Figure 3: Relationship between input noise and time to learn the task. A: The time required for the alignment μ⋅⟨w⟩/|⟨w⟩| to reach 80% depends on the noise σ of the isotropic input distributions. B: To characterize anisotropic input noise, the total input variance is split into a noise component σ‖2 parallel to and a component σ⊥2 orthogonal to the decoding direction. C: Shifting the input noise into the decoding direction slows down learning.

Figure 4: Dynamics of the total variance of w for isotropic input noise. Higher noise leads to a faster decay in tr(Cov(w)) for supervised learning (A) and for reinforcement learning (B).

Figure 5: Comparison of the theory with training on MNIST. A: A nonlinear perceptron is trained to classify the digits 0 and 1 in the MNIST dataset. B: Comparison of the empirical test classification accuracy with the theoretical prediction. C: Even after the task has been learned, the theory accurately captures non-trivial ongoing learning dynamics.

Figure 6: Forgetting curves. A: Learning curves for multi-task learning, where w are trained on Task 1 μ=μ1 after training to 80% on Task 0 μ=μ0. B. The alignment of ⟨w⟩ with μ0 after training on additional tasks 1, ..., 9.
==== Refs
References

Bordelon Blake and Pehlevan Cengiz . Learning curves for SGD on structured features. International Conference on Learning Representations. ArXiv preprint arXiv:2106.02713, 2022.
Bradbury James , Frostig Roy , Hawkins Peter , Johnson Matthew James , Leary Chris , Maclaurin Dougal , Necula George , Paszke Adam , VanderPlas Jake , Wanderman-Milne Skye , and Zhang Qiao . JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax.
Crammer Koby , Dekel Ofer , Keshet Joseph , Shalev-Shwartz Shai , Singer Yoram , and Warmuth Manfred K . Online passive-aggressive algorithms. Journal of Machine Learning Research, 7 (3 ), 2006.
Dominéz Clémentine CJ , Braun Lukas , Fitzgerald James E , and Saxe Andrew M . Exact learning dynamics of deep linear networks with prior knowledge. Journal of Statistical Mechanics: Theory and Experiment, 2023 (11 ):114004, 2023.38524253
Flesch Timo , Saxe Andrew , and Summerfield Christopher . Continual task learning in natural and artificial agents. Trends in Neurosciences, 46 (3 ):199–210, 2023.36682991
Gardner Elizabeth and Derrida Bernard . Three unfinished works on the optimal storage capacity of networks. Journal of Physics A: Mathematical and General, 22 (12 ):1983, 1989.
Hadsell Raia , Rao Dushyant , Rusu Andrei A , and Pascanu Razvan . Embracing change: Continual learning in deep neural networks. Trends in cognitive sciences, 24 (12 ):1028–1040, 2020.33158755
Haghighat Mohammad , Zonouz Saman , and Abdel-Mottaleb Mohamed . Cloudid: Trustworthy cloud-based and cross-enterprise biometric identification. Expert Systems with Applications, 42 (21 ):7905–7916, 2015. URL https://github.com/mhaghighat/gabor. BSD 2-clause license.
He Fengxiang , Liu Tongliang , and Tao Dacheng . Control batch size and learning rate to generalize well: Theoretical and empirical evidence. Advances in neural information processing systems, 32 , 2019.
Humphreys Peter C , Daie Kayvon , Svoboda Karel , Botvinick Matthew , and Lillicrap Timothy P . Bci learning phenomena can be explained by gradient-based optimization. bioRxiv, pages 2022–12, 2022.
Ji Ziwei and Telgarsky Matus . Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 33 :17176–17186, 2020.
Kudithipudi Dhireesha , Aguilar-Simon Mario , Babb Jonathan , Bazhenov Maxim , Blackiston Douglas , Bongard Josh , Brna Andrew P , Raja Suraj Chakravarthi , Cheney Nick , Clune Jeff , Biological underpinnings for lifelong learning machines. Nature Machine Intelligence, 4 (3 ):196–210, 2022.
Li Qianxiao , Tai Cheng , and Weinan E . Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research, 20 (40 ):1–47, 2019.
Li Zhiyuan , Malladi Sadhika , and Arora Sanjeev . On the validity of modeling sgd with stochastic differential equations (sdes). Advances in Neural Information Processing Systems, 34 :12712–12725, 2021.
Lim Sukbin , McKee Jillian L , Woloszyn Luke , Amit Yali , Freedman David J , Sheinberg David L , and Brunel Nicolas . Inferring learning rules from distributions of firing rates in cortical neurons. Nature neuroscience, 18 (12 ):1804–1810, 2015.26523643
Mehta Yash , Tyulmankov Danil , Rajagopalan Adithya E , Turner Glenn C , Fitzgerald James E , and Funke Jan . Model based inference of synaptic plasticity rules. bioRxiv, pages 2023–12, 2023.
Mignacco Francesca , Krzakala Florent , Urbani Pierfrancesco , and Zdeborová Lenka . Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification. Advances in Neural Information Processing Systems, 33 :9540–9550, 2020.
Murray James M and Escola G Sean . Remembrance of things practiced with fast and slow learning in cortical and subcortical pathways. Nature Communications, 11 (1 ):1–12, 2020.
Nayebi Aran , Srivastava Sanjana , Ganguli Surya , and Daniel L Yamins. Identifying learning rules from neural network observables. Advances in Neural Information Processing Systems, 33 :2639–2650, 2020.
Owen Donald Bruce . A table of normal integrals. Communications in Statistics-Simulation and Computation, 9 (4 ):389–419, 1980.
Patel Nishil , Lee Sebastian , Mannelli Stefano Sarao , Goldt Sebastian , and Saxe Andrew M . The rl perceptron: Dynamics of policy learning in high dimensions. In ICLR 2023 Workshop on Physics for Machine Learning, 2023.
Payeur Alexandre , Orsborn Amy L , and Lajoie Guillaume . Neural manifolds and learning regimes in neural-interface tasks. bioRxiv, pages 2023–03, 2023.
Portes Jacob , Schmid Christian , and Murray James M . Distinguishing learning rules with brain machine interfaces. Advances in neural information processing systems, 35 :25937–25950, 2022.37101843
Refinetti Maria , Ingrosso Alessandro , and Goldt Sebastian . Neural networks trained with sgd learn distributions of increasing complexity. In International Conference on Machine Learning, pages 28843–28863. PMLR, 2023.
Risken Hannes . The Fokker-Planck Equation. Springer, 1996.
Saxe Andrew M , McClelland James L , and Ganguli Surya . A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116 (23 ):11537–11546, 2019.
Seung Hyunjune Sebastian , Sompolinsky Haim , and Tishby Naftali . Statistical mechanics of learning from examples. Physical review A, 45 (8 ):6056, 1992.
Sutton Richard S and Barto Andrew G . Reinforcement learning: An introduction. MIT press, 2018.
Virtanen Pauli , Gommers Ralf , Oliphant Travis E. , Haberland Matt , Reddy Tyler , Cournapeau David , Burovski Evgeni , Peterson Pearu , Weckesser Warren , Bright Jonathan , van der Walt Stéfan J. , Brett Matthew , Wilson Joshua , Millman K. Jarrod , Mayorov Nikolay , Nelson Andrew R. J. , Jones Eric , Kern Robert , Larson Eric , Carey C J , Polat İlhan , Feng Yu , Moore Eric W. , VanderPlas Jake , Laxalde Denis , Perktold Josef , Cimrman Robert , Henriksen Ian , Quintero E. A. , Harris Charles R. , Archibald Anne M. , Ribeiro Antônio H. , Pedregosa Fabian , van Mulbregt Paul , and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17 :261–272, 2020. doi: 10.1038/s41592-019-0686-2.32015543
Werfel Justin , Xie Xiaohui , and Seung H . Learning curves for stochastic gradient descent in linear feedforward networks. Advances in neural information processing systems, 16 , 2003.
Williams Ronald J . Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8 :229–256, 1992.
Yaida Sho . Fluctuation-dissipation relations for stochastic gradient descent. In International Conference on Learning Representations, 2018.
Züge Paul , Klos Christian , and Memmesheimer Raoul-Martin . Weight versus node perturbation learning in temporally extended tasks: Weight perturbation often performs similarly or better. Physical Review X, 13 (2 ): 021006, 2023.
