Dynamics of Supervised and Reinforcement Learning in the Non-Linear Perceptron

Authors: christian  schmid
james m murray
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 21:28:25

Abstract: 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.

Main paper: 
 limitations: 
 Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We have included a paragraph in the Discussion section that describes the most-significant limitations of our approach. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? 
open access to data and: 
 Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: All the source code files can be found in the supplemental materials, which should allow a reader to straightforwardly recreate all experimental results.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: All parameters used for producing the results are stated in the figure captions or appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: The empirical curves in all relevant figures include shading to represent standard deviations.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: The modest compute resources required for our experiments are described in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: The research conducted in this paper conforms with the Code of Ethics in all respects.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This work is foundational research, and we do not foresee any direct societal impacts that could potentially arise from it.Guidelines:• The answer NA means that is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our work poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The only non-standard external code we used is [Haghighat et al., 2015] for the extraction of Gabor filters. As mentioned in the citation, the code is released under a BSD clause-2 license.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• Depending on the country in research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
introduction: 
 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.Figure 1 : Learning dynamics in the nonlinear perceptron. A: The perceptron, parametrized by weights w, maps an input x to the output ŷ. 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.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.We consider a general iterative update rule of the formEQUATIONwhere w t ∈ R n for arbitrary n > 0, and η is the learning rate. The stochastic update term f i 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 f i (w) does not explicitly depend on η, and that all the moments ⟨f k i ⟩ 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 p(w, t + δt) = dw ′ p(w, t + δt|w ′ , t)p(w ′ , t).(2)Denoting the weight update as δw := ww ′ , the integrand in this equation can be written as p(w, t + δt|w ′ , t)p(w ′ , t) = p(w + δw -δw, t + δt|w -δw, t)p(w -δw, t).(3)Changing the integration variable to δw and performing a Taylor expansion in δw, the right-hand side of (2) yieldsEQUATIONwhereEQUATIONandEQUATIONHere, 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 obtainEQUATIONFor 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 formEQUATIONwhere ⟨•⟩ L denotes an average over the noise in the update equation (1) (including the input distribution as well as, potentially, output noise). Thus, we findEQUATIONFinding the p(w, t) that solves this equation cannot in general be done exactly when f i 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 O((w -⟨w⟩) 3 ):EQUATIONEQUATIONTogether, these equations characterize the flow of p(w, t) for a general iterative learning algorithm in a general ANN architecture. 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 ∈ R N , multiplies it with a weight vector w ∈ R N , and produces an output ŷ. 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) = 1 2 . Each p(x|y) is given by a multivariate normal distribution with x ∼ N (µ y , Σ y ) (Fig. 1B ). We analyze both the case of SL with deterministic output, for which ŷ = ϕ(w • x), as well as RL, for which the stochastic output is given by π(ŷ = ±1) = ϕ(±w • x), where ϕ is the logistic sigmoid function. The goal of the model is to output a label ŷ that closely matches the ground truth y when given an input x.The supervised learning rule we consider is regularized stochastic gradient descent for a binary cross-entropy loss, which results in the weight update ruleEQUATIONwhere ỹ = 1 2 (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] :EQUATION)Here δ = y ŷ -⟨y ŷ⟩ is the reward prediction error, and ŷ is the stochastic output of the perceptron with probability π(ŷ = ±1) = ϕ(±w • x). To facilitate mathematical feasibility, we replace the perceptron activation function by a shifted error function ϕ(z) = 1 2 1 + Erf √ π 4 z . We first derive the learning dynamics for stochastic gradient descent. We assume that the initial condition is uniquely specified, with p(w, 0) = δ(w -w 0 ). In this case, the weight covariance will be zero, and the flow equations (10) simply reduce toEQUATIONTo make the formulas more concise, we set λ = 0. It can be reintroduced by simply adding the term -λw. We then getEQUATIONHere,Φ(z) = 1 2 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 quantitiesEQUATIONEQUATION)To aid interpretation of these results, we assume that µ ± = ±µ and Σ = σ 2 I. Then (15) impliesd dt ⟨µ • w⟩ = |µ| 2 1 -Φ µ • w 1 + σ 2 | w| 2 - 1 √ 2π σ 2 µ • w 1 + σ 2 | w| 2 e - (µ• w) 2 2(1+σ 2 | w| 2 ) w=⟨w⟩ (18)as well asd dt |⟨w⟩| 2 = 2w • µ 1 -Φ µ• w √ 1+σ 2 | w| 2 -1 2 σ 2 |w| 2 √ 1+σ 2 | w| 2 e - (µ• w) 2 2(1+σ 2 | w| 2 ) w=⟨w⟩ (19)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 calculateEQUATIONHere, T (•, •) is Owen's T function:EQUATIONAs for supervised learning, we can simplify this expression for isotropic distributions with means ±µ and get d dt ⟨µ • w⟩ =|µ| 2 4T µ • w 1 + σ 2 | w| 2 , 1 1 + 2σ 2 | w| 2 - 1 √ 2π 2σ 2 µ • w 1 + σ 2 | w| 2 e - (µ• w) 2 2(1+σ 2 | w| 2 ) Erf µ • w 1 + σ 2 | w| 2 2 + 4σ 2 | w| 2 w=⟨w⟩ (22) B C AEQUATION)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.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 Σ = σ 2 I 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 formEQUATIONwhere g iso and h iso 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.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 findEQUATIONwhere g an and h an 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.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 -w 0 ). 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 becomesEQUATIONand ( 11) implies thatEQUATION)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 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.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.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 Σ = σ 2 I, 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.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 continuallearning 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.In this section, we derive equations ( 10) and ( 11) from ( 9). We start withEQUATIONmultiply both sides by w i and integrate over w. The left-hand side simply becomesEQUATIONFor the right-hand side, we can use integration by parts to getEQUATION)To evaluate this expectation value, we introduce the mean-zero weight ŵ = w-⟨w⟩, which describes the fluctuations of w around its mean. If we expand ⟨f i ⟩ L (w) to second order in ŵ, (30) becomesEQUATIONThe derivation of (11) follows analogously.In order to analyze (10) and ( 11), we must evaluate the following expectation values:EQUATIONwith ϕ(x) = 1 2 1 + Erf √ π 4 x . Without loss of generality, we will calculate these integrals in a coordinate system where w = w 1 e 1 . We can then factorize p(x 1 , x 2 , . . .) = p m (x 1 )p c (x 2 , . . . |x 1 ). The marginal distribution is Gaussian with µ m = µ 1 and Σ m = Σ 11 , and the conditional distribution is also normal with(µ c ) i = µ i + 1 Σ11 Σ 1i (x 1 -µ 1 ), and (Σ c ) ij = Σ ij -1 Σ11 Σ 1i Σ 1j .Furthermore, to simplify notations, we introduce w = w • π/8, as well as the quantitiesEQUATIONEQUATION)whereEQUATION2) 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)x ⟩ x∼N (µ,Σ) , we first do the calculation for i = 1:EQUATION)For i ̸ = 1, we get⟨ϕ(w • x)x i ⟩ x∼N (µ,Σ) = ⟨ϕ(w 1 x 1 )x i ⟩ x∼N (µ,Σ) = 1 √ 2πΣ 11 R dx 1 e -1 2 (x1-µ1) 2 /Σ11 ϕ(w 1 x 1 ) µ i + Σ 1i Σ 11 (x 1 -µ 1 ) = µ i Φ a √ 1 + b 2 + 1 √ 2π Σ i1 w1 √ 1 + b 2 e -a 22(1+b 2 ) .(37)Thus, for general i and w we can writeEQUATION)We next calculateEQUATION)where T stands for Owen's T function.Analogously, we can calculate⟨ϕ 2 (w • x)x i ⟩ x∼N (µ,Σ) = µ i ⟨ϕ 2 (w • x)⟩ + 2(Σ w) i √ 2π √ 1 + b 2 Φ a √ 1 + b 2 √ 1 + 2b 2 e -a 22(1+b 2 ) .(40)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 formEQUATIONfor some functions f 1 and f 2 > 0. Also, it's easy to see that ⟨w * ⟩ • µ > 0. Thus, a fixed point ⟨w * ⟩ satisfiesEQUATION)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 thatEQUATION)where Erfc(z) = 1 -Erf(z) is the complementary error function.We can factor out the common exponential asymptotics of both terms to getEQUATION)Of the two terms in the parentheses, the first term starts at 1 2 for |⟨w * ⟩| = 0 and is monotonically decreasing to the value 1 2 e 1/2σ 2 Erfc( 1√ 2σ ) ≈ 1 √ 2π (σ -σ 3 + O(σ 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 d dt ⟨w⟩ = -∇ w ⟨L⟩ L , where L is the cross-entropy loss and ⟨•⟩ L is the average over the input distribution. Thus, ⟨L⟩ L is a Lyapunov function for the dynamical system. For RL, the same argument holds by the policy-gradient theorem, where ⟨L⟩ L additionally includes an average over the output noise.Note that, for λ = σ = 0, although |⟨w⟩| diverges, a calculation of d dt log (⟨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] . 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 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.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 Σ = σ 2 I, 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.1. Claims Question: the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: As we claim in the abstract, we study the learning dynamics of a nonlinear perceptron analytically, and apply the results to better understand the effects of the input distibution, learning rules and an actual perceptron training on MNIST. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
 
 Appendix:  


Decision: The paper provides a theoretical study of the behavior of a nonlinear perception on a simple problem. 

Reviewers recommended Accept, Accept and Borderline reject. While the settings are restricted, the contribution seems valuable. Reviewers appreciated the authors' modifications in response to reviews.

Overall I am willing to accept the paper.

Reviews: review: paper_summary: paper_summary: This paper analyzes the weight dynamics of single layer neural networks with nonlinear output functions.
strength_weakness: strength_weakness: The theory is tested only in binary classification task.
questions: questions: Why is the probabilistic output is considered as reinforcement learning? It should be driven by reward signal.
limitations: limitations: It is not clear hoe this extends to multi-layer cases.
paper_summary: paper_summary: This paper provides what seems to be the first derivation of equations for the dynamics of SGD-style learning in a nonlinear perceptron outside of the standard student-teacher paradigm. Specifically, a Fokker-Planck-like PDE is derived for the temporal dynamics of the weight distribution and from this ODEs are derived for the first and second moments of the weights. The equations can be applied generally to different learning rules; in the paper supervised learning of logistic regression and REINFORCE are studied. Through numerical solving and analysis of the dynamical equations several interesting results are found. In particular, it was observed that increasing noise has opposite effects on supervised learning compared with REINFORCE, and that noise orthogonal to the classification boundary speeds learning in both REINFORCE and supervised learning.
strength_weakness: strength_weakness: In my view the paper’s primary weaknesses is its clarity. There are many points in the paper where equations are derived and the sub-steps are not thoroughly explained. This might be less of a problem in a physics or mathematics publication, but for the general audience encountered at NeurIPS I believe that there should be more “holding of the reader’s hand” w.r.t. derivations. There are also several sections where it could be useful to explain assumptions a little more thoroughly (e.g. around the continuum limit). If it wasn't for this point regarding assumptions I would have given the paper a soundness rating of 4. If it wasn't for the issues of clarity I would have given a higher presentation rating.

# Notes on my review
In part because of this issue of clarity–and in part because of lack of time on my part and a large reviewing load for this conference–I have been unable to verify many of the equations in the paper. Specifically, I have only verified the math up to and including Equation 9. Thus, it is difficult for me to comment on the soundness of many of the results, leading to my low confidence score. Critically, my score for the paper assumes that the math that I have not verified is sound.
questions: questions: I am arranging this section in bullet points for ease of reading. Properly addressing these questions/comments could lead to me improving my rating of the paper during the rebuttal. I have divided these questions/comments based on whether I perceive them to be highly important (Primary), or regularly important (Secondary). I have also included a section with basic writing related comments at the end (Syntax). When a question/comment corresponds with a line (or lines) of the paper I list that line at the start of the comment.

# Primary
- Previous work studying SGD dynamics has had trouble deriving a continuum limit that meaningfully takes into account noise (i.e. that doesn’t become gradient dynamics in the limit – see e.g. section 2.3.3 in Yaida 2019, a reference in your paper). It could be nice to discuss a little (beyond the comment on line 70) in the text how your work is able to derive a continuous time F-P equation and avoid this issue
- Related to the above: it would be nice to discuss the assumptions necessary for the continuum limit in the main derivation. If there are key assumptions, this could be listed in the limitations in the discussion section
- 78 (equation 9): this is of course basically a Fokker-Planck equation but with a particularly chosen drift term and no diffusion term. Could you elaborate on the similarities and differences with the Fokker-Planck Equation (FPE) please?
- 80: I am not sure why the definition $\delta w = \langle w \rangle - w$ can be made. Is this a different variable than $\delta\mathbf{w}$ (with $w$ boldface) from before?
- 80-81: Risken is referenced as the means of deriving the moment equations. Given that the PDE studied is a very particular version of the FPE, and that NeurIPS represents a fairly general audience, it would be nice to include derivations of the moment equations. These could be included in the appendix. It would also be nice to have page numbers for where the referenced derivations in Risken appear
- 125-126: it is claimed that there is a unique globally stable fixed point, but this is not immediately clear to me from the equations. Perhaps the analysis of the equations leading to this conclusion could be listed in the appendix?

# Secondary
- It could be nice to discuss how your work compares with the related work on SGD dynamics: *Homogenization of SGD in high-dimensions: Exact dynamics and generalization properties* – Paquette et al. 2022.
- 59: might be nice to remind the reader here that $w \in \mathbb{R}^n$ for any $n$ (unless I am mistaken). This is an important aspect of the paper and deserves to be highlighted.
- 67: would suggest changing “Performing a Taylor expansion” to a sentence that describes that you are changing the variable of integration from $w’$ to $\delta w$ and then Taylor expanding w.r.t. $\delta w$. This doesn’t take much extra space and is helpful for the reader.
- 67 (equation 4): it would be good to include the smoothness assumptions on $p$ that are required to interchange the order of integration and differentiation in this equation
- 131: “particularly for SL” => “only for SL”.
- 175: is the convolution of the input data with Gabor filters necessary? If so, a little elaboration on why would be useful
- Figure 6: not immediately clear which of the two plots are derived from the flow equations and which are from numerical simulations. Could this be explained?
- 223-224: it is suggested that this approach would allow for effects due to finite learning rate, but the learning rate is taken to zero for the continuum limit, no? Does this future direction involve simply working in discrete time?
- 321: why was $\lambda$ set to $10$ for the forgetting task? This seems abnormally high

# Syntax
- Fig 1.B: might suggest replacing multinormal => multivariate normal as the latter seems more conventional
- 34-39: long sentence. Suggest ending at “... (e.g., supervised and reinforcement 36 learning)” <period>. And starting a second sentence. For the second sentence, avoid saying both “important goal” and “longstanding goal”, for more concise writing
- 67: “becomes” => “yields” or similar
- 110 (equation 15 line 1): parentheses is on right side of $x_i$ when it should be on left
- 77: would suggest defining the average w.r.t. $p(w,t)$ later (e.g. line 83) as it doesn’t show up in the equations yet.
limitations: limitations: In this section I will primarily discuss issues related to the NeurIPS checklist. I will do so in point form in the order of appearance of the topics in the checklist. Below, I only include bullet points where I have concerns about what the authors list. Addressing these issues could also lead me to improve my rating of the paper.

- Theory assumptions and proofs: 
  - I disagree with the justification for this section. Certain assumptions are not included (e.g. assumptions necessary to interchange differentiation and integration mentioned above) and extra insight could be provided on derivations (see “clarity” mentioned above).
- Broader impact: 
  - I would argue that, even with theoretical work, one should consider downstream interests. E.g., is it ethical to study general algorithms that can easily be employed for harm (e.g. in weapons industry) without taking steps to support policy that restrict the negative application of such innovations?
  - Safegaurds: this should be “No” instead of “N/A”, for the above reasons.
paper_summary: paper_summary: In this work, the authors derive a set of equations describing gradient flow in a non-linear finite-dimensional perceptron under the assumptions that the data is multinormally distributed, there is a small learning rate, and the task is binary classification. They develop the equations both for the case of a supervised learning rule and for a reinforcement learning rule. From these equations and simple simulations, they study the impact of noise on learning time, the influence of anisotropy in input distributions, input noise covariance and continual learning. They furthermore use a preprocessed MNIST dataset to asses their derived equations beyond toy datasets.
strength_weakness: strength_weakness: The work is seemingly presented as a general theory of learning in non-linear perceptrons; however, the assumptions are very strong and far from practical settings. The assumptions might even render the learning dynamics equivalent to those of a linear model (see questions). This should be more clear in abstract and/or introduction.  

It is hard to assess the novelty of the results, as the work seems to lack references to closely related works that derive equations in similar settings and/or draw similar conclusions. Notably a comparison to the work by Refinetti [5] would be beneficial (see questions). 

Far more relevant related work exists on the dynamics of learning in complex and/or non-linear networks. Currently, the introduction presents the related work as merely falling into two approaches: the student-teacher setup and the linearized perceptron. However, several works exists for far more complex and even non-linear setups exist. References below are just a few of the possible additions. 

The MNIST dataset is heavily preprocessed and it is mentioned that the two classes used are modfied in the following way: (L177) ‘We then model these two input classes as multivariate Gaussians with covariances Σ0, and means μ0,1 178 (or °æμ after a translation).’ It is unclear how this dataset is still a deviation from the multinormal assumption.
questions: questions: 1)	Are the insights following the derivations of the gradient flow equations (i.e., sections impact noise on training time, Anisotropic input distributions, etc), also based on the assumption that the covariance of the weights is zero (L107)? If so, please elaborate on to what degree this assumption limits the applicability of your results in practical settings. 

2)	How does your work differ from the theory developed in [5]? 

3)	Provided I understand [5] and your work correctly, it seems that the dynamics of learning in non-linear perceptrons when the input data is Guassian is equivalent to the dynamics of a linear model. As such, your theory would not really capture the effects of the non-linearity. Please elaborate; if this is correct, please indicate how you would rephrase the paper to reflect this fact. 

4)	Please compare your insights to the analyis of Saxe [4]. Specifically, they mention an analysis of the following: we find that perceptual correlations can either speed up or slow down learning, depending on whether they are aligned or misaligned with the input dimensions governing task-relevant semantic distinctions.

5)	Please compare your insights on continual learning with the work of [3]

6)	Please explain better how you used/modeled the MNIST dataset, cfr L177. 

7)	How would you expand the related works to include more relevant papers? 

In general: if you agree something would be a useful contribution to your work, please explicitly state how you would update the current manuscript.

1)	Ji, Z., & Telgarsky, M. (2020). Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 33, 17176–17186.
2)	Saxe, A. M., Sodhani, S., & Lewallen, S. (n.d.). The Neural Race Reduction: Dynamics of Abstraction in Gated Networks.
3)	J Dominé, C. C., Braun, L., Fitzgerald, J. E., & Saxe, A. M. (2023). Exact learning dynamics of deep linear networks with prior knowledge *. Journal of Statistical Mechanics: Theory and Experiment, 2023(11), 114004. https://doi.org/10.1088/1742-5468/ad01b8
4)	Saxe, A. M., McClelland, J. L., & Ganguli, S. (2019). A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116(23), 11537–11546. https://doi.org/10.1073/pnas.1820226116
5)	Refinetti, M., Ingrosso, A., & Goldt, S. (n.d.). Neural networks trained with SGD learn distributions of increasing complexity.
6)	Pinson, H., Lenaerts, J., & Ginis, V. (2023). Linear CNNs Discover the Statistical Structure of the Dataset Using Only the Most Dominant Frequencies. International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, 202, 27876–27906. https://proceedings.mlr.press/v202/pinson23a.html
limitations: limitations: See questions above; no ethical considerations.

Comments: comment: We thank the reviewer for providing us with all the missing citations, and we've fixed the incorrect references to these two papers. We are very grateful for the careful consideration of our submission.
We thank the reviewer for the adjustment, and we hope that the updated version of our paper will help elucidate the missing details of the mathematical derivations.
Thank you to the authors for their detailed responses; I have increased my score accordingly. While I still have some uncertainty regarding the paper on account of my own failing to check every detail of the math, from the aspects that I have understood this paper seems impactful and entirely worthy of publication at NeurIPS.
I have read the rebuttal, The experiments with the raw MNIST data are convincing, and the added discussion and comparisons with related works strengthen the paper considerably. I would like to thank the authors for their effort and clarifications.

I have one last remark: the work of (Saxe et al, 2019) and (Domine et al, 2023) is not based on the student-teacher setup. Cfr. in (Domine, 2023):  "A line of theoretical research has considered online learning dynamics in teacher-student settings [ 36 , 37, 38 ], deriving ordinary differential equations for the average learning dynamics even in nonlinear networks. However, solving these equations requires numerical integration. By contrast, our approach provides explicit analytical solutions for the more restricted case of deep linear networks." Please consider this for the camera-ready version. 

I am by now convinced this is a solid contribution to the NeurIPS conference with a moderate to high impact on the field of theoretical deep learning, and I will raise my score to a 7.
Thanks to the reviewer for their overall positive assessment of our work, including its importance, soundness, and exposition. Thanks also for the numerous constructive comments, which we have addressed as described below.

The main weaknesses identified by the reviewer involve assumptions that we failed to adequately explain–particularly the Gaussianity of the input distribution and preprocessing of the MNIST dataset–as well as connections of our work to recent related literature. As we describe below, these shortcomings have been fully addressed in our updated manuscript.

For a detailed discussion of the preprocessing of the MNIST dataset,please refer to the explanation in the main rebuttal.

The other major issue concerned whether our theory is somehow equivalent to a linear model given that we model the input data as Gaussian. It is not, as we describe in detail in the rebuttal below.

Finally, the reviewer pointed out connections of our work to other recently published work. We are grateful to the reviewer for these and discuss how these works bear on ours in detail below.
We appreciate the reviewer’s assessment of our work as impactful and novel. We also appreciate the reviewer’s constructive suggestions to improve our paper’s clarity, which was described as the primary weakness of the paper. In particular, the reviewer pointed out that our work’s soundness and presentation would be enhanced by filling in more steps of our derivations and by explaining certain assumptions more thoroughly. In response to the reviewer’s suggestion, we have written a new appendix in which we provide the details of our derivations of the moment equation (10) and (11), the flow equations, as well as another appendix in which we prove the existence of a unique, stable fixed point. We also explicitly added mathematical assumptions where they are being used. We describe additional changes we have made to address the reviewer’s comments in the rebuttal.