A theory of learning with constrained weight-distribution
Abstract
A central question in computational neuroscience is how structure determines function in neural networks.
Recent large-scale connectomic studies have started to provide a wealth of structural information such as the distribution of excitatory/inhibitory cell and synapse types as well as the distribution of synaptic weights in the brains of different species.
The emerging high-quality large structural datasets raise the question of what general functional principles can be gleaned from them.
Motivated by this question, we developed a statistical mechanical theory of learning in neural networks that incorporates structural information as constraints.
We derived an analytical solution for the memory capacity of the perceptron, a basic feedforward model of supervised learning, with constraint on the distribution of its weights.
Interestingly, the theory predicts that the reduction in capacity due to the constrained weight-distribution is related to the Wasserstein distance between the cumulative distribution function of the constrained weights and that of the standard normal distribution.
To test the theoretical predictions, we use optimal transport theory and information geometry to develop an SGDbased algorithm to find weights that simultaneously learn the input-output task and satisfy the distribution constraint.
We show that training in our algorithm can be interpreted as geodesic flows in the Wasserstein space of probability distributions.
We further developed a statistical mechanical theory for teacher-student perceptron rule learning and ask for the best way for the student to incorporate prior knowledge of the rule (i.e., the teacher).
Our theory shows that it is beneficial for the learner to adopt different prior weight distributions during learning, and shows that distribution-constrained learning outperforms unconstrained and sign-constrained learning.
Our theory and algorithm provide novel strategies for incorporating prior knowledge about weights into learning, and reveal a powerful connection between structure and function in neural networks.
Introduction
Learning and memory are thought to take place at the microscopic level by modifications of synaptic connections.
Unlike learning in artificial neural networks, synaptic plasticity in the brain operates under structural biological constraints.
Theoretical efforts to incorporate some of these constraints have focused largely on the degree of connectivity [17, 35] and the constraints on the sign of the synapses (Excitatory vs. Inhibitory) [4, 16], but few include additional features of synaptic weight distributions observed in the brain [11].
More generally, recent large-scale connectomic studies [36, 59, 61] are beginning to provide a wealth of structural information of neuronal circuits at an unprecedented scope and level of precision, which presents a remarkable opportunity for a more refined theoretical study of learning and memory that takes into account these hitherto unavailable structural information.
Perceptron [56] is arguably the simplest model of computation by single neuron and is the fundamental building block for many modern neural networks.
Despite the drastic oversimplification, studying the computational properties of (binary and analog) perceptron has been used extensively in computational neuroscience since its dawn, particularly in the cerebellum (as a model of sensory-motor association) but also in cerebral cortex (for generic associative memory functions) [43, 2, 16, 18, 15, 14].
Forming associations is considered an ‘atomic’ building block for generic cortical functions, and perceptron memory capacity sets a tight bound on the memory capacity in recurrently connected neuronal circuits with application to cortex and hippocampus [27, 55, 57].
Statistical mechanical analysis predicts that near capacity, an unconstrained perceptron classifying random input-output associations has normally distributed weights [29, 28, 21], see Fig.1(a).
In contrast, physiological experiments suggest that biological synapses do not change their excitatory/inhibitory identity during learning (but see recent [33]).
In order to take perceptron a step closer to biological realism, prior work has imposed sign constraints during learning [4, 16].
In this case, the predicted weight distribution is a delta-function centered at zero plus a half-normal distribution, see Fig.1(b).
However, a wide range of connectomic studies ranging from cortical circuits in animals [38, 32, 47, 73, 61, 40, 9], to human cerebral cortex [47, 61] have shown evidence of lognormally distributed synaptic connections.
As an example, Fig.1(c) shows the weight connection distribution in mouse primary auditory cortex (data adapted from [38]).
Possible reasons for the ubiquitous lognormal distributions range from biological structural/developmental constraints to computational benefits [63].
Various potential mechanisms for lognormal distributions has been proposed, from multiplicative gradient updates in feedforward networks[34, 40], to mixture of additive and multiplicative plasticity rules in spiking networks[30], but the majority of these proposals lead not just to lognormal distributions but also to sparsification in the weights. Instead of adding yet another explanation to the computational origin of lognormal distribution, here we take the observed weight distribution as a prior on the network structure, and ask for its computational consequences.
The goal of the paper is to present for the first time a quantitative and qualitative theory of neural network learning performance under non-Gaussian and general weight distributions (not limited to lognormal distributions).
In this paper, we combine two powerful tools: statistical mechanics and optimal transport theory, and present a theory of perceptron learning that incorporates the knowledge of both distribution and sign information as constraints, and gives accurate predictions for capacity and generalization error.
Interestingly, the theory predicts that the reduction in capacity due to the constrained weightdistribution is related to the Wasserstein distance between the cumulative distribution function of the constrained weights and that of the standard normal distribution.
Along with the theoretical framework, we also present a learning algorithm derived from information geometry that is capable of efficiently finding viable perceptron weights that satisfy desired distribution and sign constraints.
This paper is organized as follows: in Section 2.1 we derive the perceptron capacity for classifying random input-output associations using statistical mechanics, and illustrate our theory with a simple example.
In Section 3, we derive our learning algorithm using optimal transport theory, and show that distribution of weights found by the learning algorithm coincide with geodesic distributions on a Wasserstein statistical manifold, and therefore training can be interpreted as a geodesic flow.
In
Section 4 we analyze a parameterized family of biologically realistic weight distributions, and use our theory to predict the shape of the distribution with optimal parameters.
We map out the experimental parameter landscape for the estimated distribution of synaptic weights in mammalian cortex and show that our theory’s prediction for optimal distribution is close to the experimentally measured value.
In Section 5 we further develop a statistical mechanical theory for teacher-student perceptron rule learning and ask for the best way for the student to incorporate prior knowledge about the weight distribution of the rule (i.e., the teacher).
Our theory shows that it is beneficial for the learner to adopt different prior weight distributions during learning.
[Figure]
Capacity 2.1 Learning under weight distribution constraints
We begin by considering a canonical learning problem: classifying random input-output associations by a perceptron.
In biological memory systems, the heavily correlated sensory data is undergoing heavy preprocessing including massive decorrelations, and previous work on brain related perceptron modeling [27, 16, 57] assumes similarly unstructured data.
The data consists of pairs { ξ µ , ζ µ } Pµ =1 , where ξ µ is an N -dimentional random vector drawn i.i.d.
from a standard normal distribution, p ( ξ µi ) = N (0 , 1) , and ζ µ are random binary class labels with p ( ζ µ )
= 12 δ ( ζ µ + 1) + 12 δ ( ζ µ − 1) .
The goal is to find a hyperplane through the origin, described by a perceptron weight vector w ∈ R N , normalized to || w || 2 = N.
We call w a separating hyperplane when it correctly classifies all the examples with margin κ > 0 :
[Equation]
We are interested in solutions w to Eqn.1 that obey a prescribed distribution constraint, w i ∼ q ( w ) , where q is an arbitrary probability density function.
We further demand that ⟨ w 2 ⟩ q ( w ) = 1 to fix the overall scale of the distribution (since the task is invariant to the overall scale of w ).
Thus, the goal of learning is to find weights that satisfy 1 with the additional constraint that the empirical density function ˆ q ( w ) = 1 N (cid:80) Ni δ ( w − w i ) , formed by the learned weights is similar to q ( w ) , and more precisely that it converges to q ( w ) as N → ∞ (see Section 2.2 below).
Extension of this setup that includes an arbitrary number of populations each satisfying its own prescribed distribution constraints is discussed in Section 4 and in Appendix A.1.2.
Note that the sign constraint is a special case of this scenario with two synaptic populations: one excitatory and one inhibitory.
We further discuss the generalization of this setup to include biased inputs and sparse labels in Appendix A.1.3.
2.2 Statistical mechanical theory of capacity
We are interested in the thermodynamic limit where P, N → ∞ , but the load α
= PN stays O (1) .
This limit is amenable to mean-field analysis using statistical mechanics.
Following Gardner’s seminal work [29, 28], we consider the fraction V of viable weights that satisfies both Eqn.1 and the distribution constraint ˆ q = q, to all possible weights:
[Equation]
[Figure] In Eqn.2, we impose the distribution constraint ˆ q = q by demanding that in the thermodynamic limit, all Fourier modes of q and ˆ q are the same , i.e., that q ( k ) = (cid:1) dwe ikw q ( w )
=
ˆ q ( k ) = 1 N (cid:80)
Ni e ikw i , where in the last equality we have used the definition of empirical distribution.
We perform a quenched average over random patterns ξ µ and labels ζ µ .
This amounts to calculating ⟨ log V ⟩ , which can be done using the replica trick [29, 28].
We focus on solutions with maximum margin κ at a given load α , or equivalently, the maximum load capacity α c ( κ ) of separable patterns given margin κ .
We proceed by assuming replica symmetry in our mean field analysis, which in general might not hold because the constraint ˆ q = q is non-convex.
For all the results presented in the main text, replica symmetry solution is supported by numerical simulations.
In Appendix A.5 we explore the validity of replica symmetric solutions in the case of strongly bimodal distributions and show that they fail only very close to the binary (Ising) limit.
Detailed calculations of the mean-field theory are presented in Appendix A.1.1.
Our mean-field theory predicts that the reduction in capacity due to the distribution constraint is proportional to the
Jacobian of the transformation from w ∼ q ( w ) to a normally distributed variable x ( w ) ∼ N (0 , 1) , α c
[Equation]
= [Equation] [Equation] κ
[Equation]
(3) where α 0 ( κ ) = (cid:104) (cid:1) ∞ − κ Dt ( κ + t ) 2 (cid:105) − 1 is the capacity of an unconstrained perceptron, from Gardner theory [29, 28], and κ = 0 reduces to the classical result of α 0 (0)
= 2 .
The Jacobian factor, ⟨ dw/dx ⟩ x , can be written in terms of the constrained distribution’s cumulative distribution function (CDF), Q ( w ) , and the standard normal CDF P ( x ) = 12 (cid:104) 1 + Erf ( x [Equation] 2 ) (cid:105) , namely, [Equation]
[Equation]
Note that since the second moments are fixed to unity, 0 ≤ (cid:10) dwdx (cid:11) x ≤ 1 and it equals 1 iff p =
q.
2.3 Geometrical interpretation of capacity
The jacobian factor Eqn.4 can be rewritten as
[Equation]
where W k ( k = 2 in above) is the Wasserstein- k distance, given by
[Equation]
[In the following, we will make frequent use of both the probability density function (PDF), and the cumulative distribution function (CDF).
We distinguish them by using upper case letters for CDFs, and lower case letters for PDFs.]
The Wasserstein distance measures the dissimilarity between two probability distributions, and is the geodesic distance between points on the manifold of probability distributions [41, 25, 20].
Therefore, we can interpret Eqn.3 as predicting that the reduction in memory capacity tracks the geodesic distance we need travel from the standard normal distribution P to the target distribution Q (Fig.2(a)).
We demonstrate Eqn.3 and Eqn.5 with an instructive example.
Let’s consider a parameterized family of normal distributions, with the second moment fixed to 1: q ( w ) =
N ( √ 1 − σ 2 , σ 2 ) , see Fig.2(c).
At σ = 1 , q ( w ) is the standard normal distribution and we recover the unconstrained Gardner capacity α 0 ( κ = 0)
= 2 .
As σ → 0 , q ( w ) becomes a δ -function at 1 and α c ( κ ) → 0 (Fig.2(c)).
As evident in this simple example, perceptron capacity is strongly affected by its weight distribution.
Our theory enables prediction of the shape of the distribution with optimal parameters within a parameterized family of distributions.
We apply our theory to a family of biologically plausible distributions and compare our prediction with experimentally measured distributions in Section 4.
Optimal transport and the DisCo-SGD learning algorithm
Eqn.3 predicts the storage capacity for a perceptron with a given weight distribution, but it does not specify a learning algorithm for finding a solution to this non-convex learning problem.
Here we present a learning algorithm for perceptron learning with a given weight distribution constraint.
This algorithm will also serve to test our theoretical predictions.
For this purpose, we use optimal transport theory to develop an SGD-based algorithm that is able to find max-margin solutions that obey the prescribed distribution constraint.
Furthermore, we show that training can be interpreted as traveling along the geodesic connecting the current empirical distribution and the target distribution.
Stochastic gradient descent (SGD) on a cross-entropy loss has been shown to asymptotically converge to max-margin solutions on separable data [62, 50].
Given data { ξ µ , ζ µ } Pµ =1 , we use logistic regression to predict class labels from our perceptron weights, ˆ ζ µ = σ ( w t · ξ µ ) , where σ ( z ) = (1 + e − z ) − 1 and w t is the weight at the t -th update.
This defines an SGD update rule :
[Equation]
where the µ -summation goes from 1 to P for full-batch GD and goes from 1 to mini-batch size B for mini-batches SGD (see Appendix A.4 for more details).
The theory of optimal transport provides a principled way of transporting each individual weight w ti to a new value so that overall the new set of weights satisfies the prescribed target distribution.
In 1 -D, the optimal transport plan T ∗ has a closed-form solution in terms of the current CDF P and target CDF Q [65, 3]: T ∗ =
Q − 1 ◦ P , where ◦ denotes functional composition.
We demonstrate the optimal transport map in Fig.2(b) for the instructive example discussed in Section 2.3.
In order to apply T ∗ to transport our weights { w i } (omitting superscript t ), we form the empirical CDF ˆ Q ( w ) =[-1 N (cid:80) Ni =
-]
1 N (cid:80) N i =1 1 w i ≤ w , which counts how many weights w i are observed below value w .
Then the new set of weights { ˆ w i } satisfying target CDF Q can be written as
[Equation]
We illustrate Eqn.8 in action in Table 1(b). [Figure] [Figure]
However, performing such a one-step projection strongly interferes with the cross-entropy objective and numerically often results in solutions that do not perfectly classify the data.
Therefore, it would be beneficial to have an incremental update rule based on Eqn.8:
[Equation]
where we have used a different update time τ to differentiate with the cross-entropy update time t .
We present our complete algorithm in Table 1(a), which we named ‘Distribution-constrained SGD’ (DisCo-SGD) algorithm.
In the DisCo-SGD algorithm, we perform alternating updates on Eqn.7 and
Eqn.9, and identify δt and δτ as learning rates η 1 and η 2 .
Note that in logistic regression, the norm of the weight vector || w || is known to increase with training and the max-margin solution is only recovered at || w || → ∞ .
In contrast, imposing a distribution constraint fixes the norm.
Therefore, to allow a variable norm, in Table 1 we include a trainable parameter β in our algorithm to serve as the norm of the weight vector.
This algorithm allows us to reliably discover linearly separable solutions obeying the prescribed weight distribution Q .
Interestingly, Eqn.9 takes a similar form to geodesic flows in Wasserstein space.
Given samples { w i } drawn from the initial distribution P and { ˆ w i } drawn from the final distribution Q , samples { w τi } from intermediate distributions P τ along the geodesic can be calculated as w τ ( i ) =
(1 − τ ) w ( i ) + τ ˆ w ( i ) , where subscript ( i ) denotes ascending order (see more in Appendix A.2).
For intermediate perceptron [Figure] weights w t found by our algorithm, we can compute its empirical distribution p t and compare with theoretical distribution p τ along the geodesic with the same W 2 distance to the target distribution (see
Appendix A.2 for how to calculate p τ ).
In Fig.3(a), we show that indeed the empirical distributions p t agree with the geodesic distributions p τ at geodesic time τ ( t ) (Fig.3(a)).
The relation between the geodesic time τ and the SGD update time t is shown in Fig.3(b).
The interplay between the crossentropy objective and the distribution constraint is manifested in the rate at which the distribution moves along the geodesic between the initial distribution and the target one.
Biologically-realistic distribution (E/I balanced lognormals) and experimental landscape
In order to apply our theory to the more biologically-realistic cases, we generalize our theory from a single prescribed distribution to an arbitrary number of input subpopulations each obeys its own distribution.
We consider a perceptron that consists of M synaptic populations w m indexed by m , each constrained to satisfy its own weight distribution w mi ∼ q m ( w m ) .
We denote the overall weight vector as w ≡ { w m } Mm =1 ∈ R N × 1 , where the total number of weights is N = (cid:80) Mm =1 N m .
In this case, the capacity Eqn.3 is generalized to (See Appendix A.1.2 for detailed derivation):
[Equation]
where g m = N m /N is the fraction of weights in this population.
Eqn.
10 allows us to investigate the parameter space of capacity with biologically-realistic distributions and compare with the experimentally measured values.
In particular, we are interested the case with two synaptic populations that models the excitatory/inhibitory synpatic weights of a biological neuron, hence, m =
E, I .
We model the excitatory/inhibitory synaptic weights as drawn from two separate lognormal distributions ( g I = 1 − g E ): w Ei ∼
[Equation]
and w Ii
[Equation]
We also demand that the mean synaptic weights satisfy the E/I balance condition [69, 70, 67, 68, 57, 48, 18] g E (cid:10) w E (cid:11)
=
g I (cid:10) w I (cid:11) as is often observed in cortex connectomic experiments [5, 72, 51, 53, 8].
With the E/I balance condition and fixed second moment, the capacity is a function of the lognormal parameters σ E and σ I .
In Fig.4(a) we map out the 2d parameter space of σ E and σ I using Eqn.10, and find that the optimal choice of parameters which yields the maximum capacity solution is close to the experimentally measured values in a recent connectomic studies in mouse primary auditory cortex [38].
In order to test our theory’s validity on this estimated distribution of synaptic weights, we perform
DisCo-SGD simulation with model parameters σ E and σ I fixed to their experimentally measured values.
Both the capacity (Fig.4(b)), max-margin κ at different load (Fig.4(c)), and the empirical weights found by the algorithm (Fig.4(d)) are in good agreement with our theoretical prediction.
Generalization performance 5.1 Distribution-constrained learning as circuit inference
A central question in computational neuroscience is how underlying neural circuits determine its computation.
Recently, thanks to new parallelized functional recording technologies, simultaneous recordings of the activity of hundreds of neurons in response to an ensemble of inputs are possible [1, 12].
An interesting challenge is to infer the structural connectivity from the measured inputoutput activity patterns.
It is interesting to ask how are these stimuli-response relations related to the underlying structure of the circuit [54, 39].
In the following, we try to adress this circuit reconstruction task in a simple setup where a student perceptron tries to learn from a teacher perceptron [60, 22].
In this setup, the teacher is considered to be the underlying ground-truth neural circuit.
The student is attempting to infer the connection weights of this ground-truth circuit by observing a series of input-output relations generated by the teacher.
After learning is completed, one can assess the faithfulness of the inference by comparing the teacher and student.
The teacher-student setup is also a well-known ‘toy model’ for studying generalization performance [42, 37, 44].
In this case since the learning data are generated by the teacher, the overlap between teacher and student determines the generalization performance of the learning.
Here we ask to what extent prior knowledge of the teacher weight distribution helps in learning the rule and how this knowledge can be incorporated in learning.
A similar motivation may arise in other contexts, in which there is a prior knowledge about the weight distribution of an unknown target linear classifier.
Let’s consider the teacher perceptron, w t ∈ R N , drawn from some ground-truth distribution p t .
Given random inputs ξ µ with p ( ξ µi ) =
N (0 , 1) , we generate labels by ζ µ = sgn ( w t · ξ µ / ||
w t || + η µ ) , where η µ is input noise and η µ ∼ N (0 , σ 2 ) .
We task the student perceptron w s to find the max-margin linear classifier for data { ξ µ , ζ µ } pµ =1 : max κ : ζ µ w s · ξ µ ≥ κ || w s || .
Let’s define the teacher-student overlap as
[Equation]
(11) which is a measure the faithfulness of the circuit inference.
The student’s generalization error is then related to the overlap by ε g = 1 /π arccos (cid:0) R/ √ 1 + σ 2 (cid:1) [60, 22].
As a baseline, let’s first consider a totally uninformed student (without any structural knowledge of the teacher), learning from a teacher with a given (in particular non-Gaussian) weight distribution.
In this case, we can determine the overlap R (Eqn.11) as a function of load α by solving the replica symmetric mean field self-consistency equations as in [60, 22].
An example of such learning for a lognormal teacher distribution is shown in Fig.5(a) (‘unconstrained’) for the noiseless case ( σ = 0 ).
Note that in the presence of noise in the labels ( σ ̸ = 0) , α is bounded by α c ( σ ) , since max-margin learning of separable data is assumed.
The case with nonzero σ is presented in Appendix A.3.4.
In this unconstrained case, the student’s weight distribution evolves from a Gaussian for low α to one which increasingly resembles the teacher distribution for large α (Fig.5(b)).
Next, we consider a student with information about the signs of the individual teacher weights.
We can apply this knowledge as a constraint and demand that the signs of individual student weights agree with that of the teacher’s.
The additional sign-constraints require a modification of replica calculation in [60, 22], which we present in Appendix A.3.1.
Surprisingly, we find both analytically and numerically that if the teacher weights are not too sparse, the max-margin solution generalizes poorly: after a single step of learning (with random input vectors), the overlap, R , drops substantially from its initial value (see ‘sign-constrained’ in Fig.5(a)).
The source of the problem is that, due to the sign constraint, max-margin training with few examples yields a significant mismatch between the student and teacher weight distributions.
After only a few steps of learning, half of the student’s weights are set to zero, and the student’s distribution,
[Equation] , deviates significantly from the teacher’s distribution (see more in Appendix A.3.3).
The discrepancy between the teacher and student weight distributions therefore suggest that we should incorporate distribution-constraint into learning. [Figure]
5.2 Distribution-constrained learning outperforms unconstrained and sign-constrained learning
Let’s consider the case that the student weight are constrained to some prior distribution q s ( w s ) , while the teacher obeys a distribution p t ( w t ) , for an arbitrary pair q s , p t .
We can write down the
Gardner volume V g for generalization as in the capacity case (Eqn.2): V
[Equation]
To obtain ensemble average of system over different realizations of the training set, we perform the quenched average of log V g over the patterns ξ µ and teacher w t , and consider the thermodynamic limit of N, P → ∞ and α = PN stays O (1) .
We use the replica trick similar to [60, 22].
Overlap
R (Eqn.11) can be determined as a function of load α by solving the replica symmetric mean field self-consistency equations in Appendix A.3.2.
In this distribution-constrained setting, we can perform numerical simulations with DisCo-SGD algorithm (Table 1) to find such weights and compare with the predictions of our theory.
Now we ask if the student has a prior on the teacher’s weight distribution p t , whether incorporating this knowledge in training will improve generalization performance.
One might be tempted to conclude that the optimal prior distribution the student should adopt is always that of the teacher’s, i.e., q s = p t .
We call this learning paradigm ‘fixed prior’, and show that its generalization performance is better than that of the unconstrained and sign-constrained case (Fig.5(a)).
However, instead of using a fixed prior for the student, we can in fact choose the optimal prior distribution p ∗ s at different load α .
This presents a new learning paradigm we called ‘optimal prior’.
In Fig.5(a), we show that choosing optimal priors at different α achieves the overall best generalization performance compared with all other learning paradigms.
For a given parameterized family of distributions, our theory provides a way to analytically obtain the optimal prior p ∗ s as a function of α (Fig.5(c)).
Note that unlike the unconstrained case (Fig.5(b)), the optimal prior starts from a δ -function at 1 at zero α , and asymptotically approaches the teacher distribution p t as α → ∞ .
Summary and Discussion
We have developed a statistical mechanical framework that incorporates structural constraints (sign and weight distribution) into perceptron learning.
The synaptic weights in our perceptron learning satisfy two key biological constraints: (1) individual synaptic signs are not affected by the learning task (2) overall synaptic weights obey a prescribed distribution.
These constraints may arise also in neuromorphic devices [31, 66].
Under the replica-symmetry assumption, we derived a novel form of distribution-constrained perceptron storage capacity, which admits a simple geometric interpretation of the reduction in capacity in terms of the Wasserstein distance between the standard normal distribution and the imposed distribution.
To numerically test our analytic theory, we used tools from optimal transport and information geometry to develop an SGD-based algorithm, DisCo-SGD, in order to reliably find weights that satisfy such prescribed constraints and correctly classify the data, and showed that training with the algorithm can be interpreted as geodesic flows in the Wasserstein space of distributions.
It would be interesting to compare our theory and algorithm to [7, 58] where the Wasserstein distance is used as an objective for training generative models.
We applied our theory to the biologically realistic case of of excitatory/inhibitory lognormal distributions that are observed in the cortex, and found experimentally-measured parameters close to the optimal parameter values predicted by our theory.
We further studied input-output rule learning where the target rule is defined in terms of a weighted sum of the inputs, and asked to what extent prior knowledge of the target distribution may improve generalization performance.
Using the teacher-student perceptron learning setup, we showed analytically and numerically that distribution constrained learning substantially enhances the generalization performance.
In the context of circuit inference, distribution constrained learning provides a novel and reliable way to recover the underlying circuit structure from observed input-output neural activities.
In summary, our work provides new strategies of incorporating knowledge about weight distribution in neural learning and reveals a powerful connection between structure and function in neural networks.
Ongoing extensions of the present work include weight distribution constraints in recurrent and deep architectures as well as testing against additional connectomic databases. Acknowledgments
This paper is dedicated to the memory of Mrs. Lily Safra, a great supporter of brain research.
The authors would like to thank Madhu Advani, Haozhe Shan, and Julia Steinberg for very helpful discussions.
W.Z. acknowledges support from the Swartz Program in Theoretical Neuroscience at
Harvard and the NIH grant NINDS (1U19NS104653). B.S. acknowledges support from the Stanford Graduate Fellowship. D.D.L. acknowledges support from the NIH grant NINDS (1U19NS104653).
H.S. acknowledges support from the Swartz Program in Theoretical Neuroscience at Harvard, the
NIH grant NINDS (1U19NS104653), and the Gatsby Charitable Foundation.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] or [N/A] .
You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section 3.
• Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section 2.2 and Appendix A.5.
(c) Did you discuss any potential negative societal impacts of your work?
[No] We do not see any potential negative societal impacts of our work.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section 2.2.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A.1-A.3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See Section 3.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Section A.4.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] Our theory concerns the thermodynamic limit where the observables are dominated by their typical values and fluctuations are small.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[N/A] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Appendix Preliminaries
Throughout the appendix, we make frequent use of Gaussian integrals.
We introduce short-hand notations (cid:1)
[Equation]
2 π e − t 2 / 2 and H ( x ) ≡ (cid:1) ∞ x
Dt .
Also, when we do not specify the integration range it is understood that we are integrating from −∞ to ∞ .
A.1 Capacity supplemental materials A.1.1 Replica calculation of distribution-constrained capacity
In this section, we present the replica calculation of the distribution-constrained storage capacity of a perceptron.
As described in main text Eqn.2, we need to perform a quenched average ⟨·⟩ over the patterns ξ µ and labels ζ µ for log V , which can be carried out using the replica trick, ⟨ log V ⟩ = lim n → 0 ( ⟨ V n ⟩− 1) /n .
Following [29, 28], we consider first integer n , and at the end perform analytic continuation of n → 0 .
The replicated Gardner volume is:
[Equation]
Let’s rewrite the Heaviside step function using Fourier representation of the δ -function δ ( x ) = (cid:1)
[Equation]
ikx as (defining z µα = ζ
[Equation]
|
[Equation]
|| )
[Equation]
Note that now all the ξ µ , ζ µ dependence is in e − ix µα z µα .
We perform the average with respect to ξ µi ∼ p ( ξ µi ) =
N (0 , 1) and p ( ζ µ ) = 12 δ ( ζ µ + 1) + 12 δ ( ζ µ − 1) (also note that || w α || = √ N ):
[Equation]
Introducing the replica overlap parameter q αβ = 1 N (cid:80) i w αi w βi , and notice that the µ index gives P identical copies of the same integral.
We can suppress the µ indices and write
[Equation]
where
[Equation]
captures all the data dependence in the quenched free energy landscape, and therefore it is called the ‘energetic’ part of the free energy.
In contrast, the δ -functions in Eqn.13 are called ‘entropic’ part because they regulate what kind of weights are considered in the version space (space of viable weights).
The entropic part
[Equation]
Note that the normalization constraint δ ( || w α
||
− N ) is automatically satisfied by requiring q αα = 1 .
Using replica-symmetric ansatz:
ˆ q αβ = − i 2 (∆ˆ qδ αβ + ˆ q 1 ) , and q αβ = (1 − q )
δ αβ + q , we have
[Equation]
and
[Equation]
where in the last step HST denotes Hubbard-Stratonovich transformation
[Equation]
that we use to linearize the quadratic term at the cost of introducing an auxiliary Gaussian variable t to be averaged over later.
Recall that ˆ q
[Equation]
k
[Equation]
=
[Equation] ikw
[Equation]
, the distribution constraint becomes
[Equation]
Note that (cid:80)
i
[Equation]
) by inverse Fourier transform.
Next,
[Equation]
Now we can write down the full free energy.
We ignore overall constant coefficients such as 2 π ’s and i ’s in the integration measure, which become irrelevant upon taking the saddle-point approximation.
We also leave out the denominator of V , as it does not depend on data and is an overall constant.
Note that under the replica-symmetric ansatz the replica index α gives n identical copies of the same integral and thus the replica index α can be suppressed (same for synaptic index i ): [Equation]
[Equation]
where (please note that q is replica overlap, and q ( w ) is the imposed target distribution) [Equation]
[Equation]
Note that integrals in Eqn.23 can be evaluated using saddle-point approximation in the thermodynamic limit N → ∞ .
Redefining 2 πiλ ( − w ) − 12 ∆ˆ qw 2 → − λ ( w ) and − ˆ q 1 → ˆ q 1 , we have
[Equation]
We seek the saddle-point solution for G 0 with respect to the order parameters ∆ˆ q , λ ( w ) , and ˆ q 1 :
[Equation]
[Equation]
[Equation]
We observe that the saddle-point equation Eqn.26 fixes the second moment of the imposed distribution q ( w ) to 1 and therefore can be thought of as a second moment constraint.
G 0 now simplifies to
[Equation]
The remaining ˆ q 1 saddle-point equation is a bit more complicated,
[Equation]
Integration by parts for the second term in rhs:
[Equation]
) where in the last step we have defined an induced distribution f ( w ) =
Z ( t ) − 1 exp (cid:8) −
[Equation]
w ) + √ ˆ q 1 tw
[Equation]
Since the second moments are fixed to 1, we have
[Equation]
(31) which gives a nice interpretation of q in terms of the average overlap of w in the induced distribution f ( w ) .
Limit q →
We are interested in the critical load α c where the version space (space of viable weights) shrinks to a single point, i.e., there exists only one viable solution.
Since q measures the typical overlap between weight vectors in the version space, the uniqueness of the solution implies q → 1 at α c .
In this limit, the order parameters { ˆ q 1 , λ ( w ) } diverges and we need to re-derive the saddle point equations Eqn.
and Eqn.31 in terms of the undiverged order parameters { u, r ( w ) } :
[Equation]
Now G 0 becomes
[Equation]
33) and
Z
[Equation]
We can perform a saddle-point approximation for the w integral in Z ( t ) at the saddle value w such that r ′ ( w ) = ut :
[Equation]
Then
G [Equation] =
[Equation]
Let’s use integration by parts to rewrite
[Equation]
(37) where Q ( w ) is the CDF of the imposed distribution q ( w ) and P ( t )
= 12 (cid:104)
+ Erf ( t √ 2 ) (cid:105) is the normal
CDF.
Now the saddle-point equation
[Equation]
(38) determines w ( t ) implicitly.
The u equation gives
[Equation]
[Equation]
where in the last equality we have used integration by parts.
Using Eqn.38-39 G 0 is simplified to G 0 =
[Equation]
dt
[Equation]
t .
The energetic part
We would like to perform a similar procedure as shown above, to Eqn.17 using the replica-symmetric ansatz.
We observe that the effect of the distribution constraint is entirely captured in G 0 and therefore G 1 is unchanged compared with the standard Gardner calculation of perceptron capacity.
We reproduce the calculation here for completeness.
Under the replica-symmetric ansatz q αβ = (1 − q )
δ αβ + q , Eqn.17 becomes
[Equation]
where we have again used the Hubbard-Stratonovich transformation to linearize the quadratic piece.
Performing the Gaussian integrals in x α (define α = PN ),
[Equation]
At the limit n → 0 ,
[Equation]
Perform the Gaussian integral in ρ and define ˜ κ = κ + t √ q √ 1 − q , we have [Equation]
[Equation]
At the limit q → 1 , α → α c , (cid:1) ∞
−∞ Dt is dominated by (cid:1) ∞ − κ Dt , and H (˜ κ ) → 1 √ 2 π ˜ κ e − ˜ κ 2 / 2 .
The
O (cid:16) 11 − q (cid:17) (leading order) contribution gives
[Equation]
Let G
=
G 0 + G 1 .
As n → 0 , ⟨ V n ⟩ = e n ( NG ) → 1 + n ( NG ) , and ⟨ log V ⟩ = lim n → 0 ⟨ V n ⟩− 1 n =
NG .
Combining with Eqn.40 (relabel t ↔ x to distinguish between the two auxiliary Gaussian variables), we have
[Equation]
Capacity α c is reached when Eqn.13 goes to zero.
We arrive at the distribution-constrained capacity α c ( κ [Equation] =
[Equation] ( κ ) (cid:28) dw dx (cid:29) [Equation] x , (47) where α 0 ( κ ) =
(cid:104) (cid:1) ∞ − κ Dt ( κ + t ) 2 (cid:105) − 1 is the unconstrained capacity.
Instructive Examples
(1) Standard normal distribution w ∼ N (0 , 1)
In this case w = x and α c ( κ ) =
α 0 ( κ ) .
(2) Normal distribution with nonzero mean w ∼ N ( µ, σ 2 ) .
This is the example discussed in the main text Fig.1.
In this case w = µ + σx and µ 2 + σ 2 = 1 due to the second moment constraint Eqn.26.
Then α c ( κ ) =
σ 2 α 0 ( κ ) .
Lognormal
distribution w ∼
[Equation]
In this case w = e µ + σx where µ = − σ 2 .
α c ( κ ) =
σ 2 e − σ 2 α 0 ( κ ) .
Geometrical interpretation
Note that although the Jacobian factor (cid:10) dwdx (cid:11) x takes a simple form, in practice sometimes it might not be the most convenient form to use.
Integrating by parts ( p ( x ) = N (0 , 1) ),
[Equation]
Now define u = P ( x ) so that du = p ( x ) dx and w = Q − 1 ( P ( x ))
=
Q − 1 ( u ) , we can express the
Jacobian in terms of the CDFs
[Equation]
Furthermore,
[Equation]
where we have used second moments equal to 1 and the definition of the Wasserstein- k distance in the second equality.
Therefore, we have arrived at the geometric interpretation of the Jacobian term
[Equation]
A.1.2 Theory for an arbitrary number of synaptic subpopulations
In this section, we generalize our theory in the above section to the set up of a perceptron with
M synaptic populations indexed by m , w m , such that each w mi satisfies its own distributions constraints w mi ∼ q m ( w m ) .
We denote the overall weight vector as w ≡ { w m } Mm =1 ∈ R N × 1 , where the total number of weights is N = (cid:80) Mm =1 N m .
The replica overlap now becomes q αβ = 1 N (cid:80)
Mm (cid:80) N m i w mαi w mβi .
The distribution constraint becomes (see Eqn.21 for the case of M = 1 )
[Equation]
We introduce ˆ q αβ , λ m ( k ) to write the δ -functions into Fourier representations, and use replicasymmetric ansatz
ˆ q αβ = − i 2 (∆ˆ qδ αβ + ˆ q 1 ) , and q αβ = (1 − q ) δ αβ + q as before.
After similar manipulations that lead to Eqn.25, the entropic part of the free energy becomes ( g m = N m /N is the fraction of weights in m -th population) [Equation]
[Equation]
Now the second moment constraint
= ∂G 0 /∂
∆ˆ q (Eqn.26) becomes the weighted sum of second moments from each population:
[Equation]
We take the q → 1 limit as before:
[Equation]
Use saddle-point approximation for Z m ( t ) and integrate by parts as in Eqn.35-37, the entropic part becomes
[Equation] [Equation]
Now the saddle-point equation for order parameters r ′ m ( w m ) and u gives
[Equation]
Therefore,
[Equation]
The energetic part (Eqn.35) remains unchanged and thus (relabel t ↔ x )
[Equation]
E/I balanced lognormals
Now we specialize to the biologically realistic E/I balanced lognormal distributions described in the main text.
We are interested the case with two synaptic populations m =
E, I that models the excitatory/inhibitory synpatic weights of a biological neuron.
w Ei ∼
[Equation]
and w Ii ∼
[Equation]
Let’s denote the E/I fractions as g E = r
and g I = 1 − r .
The CDF of the lognormals are given by
[Equation]
The corresponding inverse CDF is
[Equation]
The capacity is therefore α c =
α(cid:34)(cid:88)
[Equation]
This model has five parameters { r, σ E , σ I , µ E , µ I } .
We use values of r reported in experiments (the ratio between of E. connections found and I. connections found).
We also have two constraints.
The E/I balanced constraint g E (cid:10) w E (cid:11)
[Equation]
and the second moment constraint 1 = (cid:80)
[Equation]
Therefore there are two free parameters left and we choose to express µ E and µ I in terms of the rest:
[Equation]
The parameter landscape is plotted against the two free parameters σ E and σ I .
Here we report comparisons across different experiments [38, 9, 32, 47, 64, 73] similar to main text Fig.4 (Fig.
(a) is included here for reference).
Note that despite the apparently different shape of distributions, all the experimentally measured parameter values are within the first quantile of the optimal values predicted by our theory.
A.1.3 Capacity for biased inputs and sparse label
In this section, we generalized our theory in Section A.1.1 to the set up of nonzero-mean input patterns ξ µ and sparse labels ζ µ : [Figure]
[Equation]
In this case, we need to include a bias in the perceptron ˆ ζ µ = sgn ( w [Equation] ∥ − b ) to be able to correctly classify patterns in general.
Note that m = 0 and f = 1 / 2 reduces to the case in Section A.1.1.
We observe due to the multiplicative relation between the Jacobian term and the original Gardner capacity in Eqn.47, entropic effects (such as distribution constraints and sign-constraints) factors with the energetic effects (such as the nonzero mean inputs and sparse labels), and they don’t interfere with each other.
Therefore, the calculations for nonzero mean inputs and sparse labels are identical with the original
Gardner case.
Here we only reproduce the calculation for completeness.
Readers already familiar with this calculation should skip this part.
The analog of Eqn.15 reads (define the local fields as h µi =
(cid:80)
α x µα w αi )
[Equation]
µα
[Equation]
where in the second equality we have carried out the Gaussian integral in ξ µ and in the third equality we introduced the order parameters
[Equation]
Now the full energetic term becomes
[Equation]
Now G 1 becomes
[Equation]
Note that the hat-variables ˆ M α conjugated with M α are in subleading order to ˆ q αβ in the thermodynamic limit, and therefore G 0 is unchanged.
Let v =
M − b/m , we have now the capacity α
[Equation]
where the order parameter v needs to be determined from the saddle-point equation
[Equation]
In Fig.7 we numerically solve α c ( κ ) for different values of m and f .
A.2 Optimal transport theory
In recent years, Wasserstein distances has found diverse applications in fields ranging from machine learning [7, 26, 49] to geophysics [23, 24, 19, 45, 46].
In optimal transport theory, the Wasserstein- k distance arise as the minimal cost one needs to pay in transporting one probability distribution to another, when the moving cost between probability masses are measured by the L k norm [71].
When one equips the probability density manifold with the Wasserstein- 2 distance as metric, it becomes the
Wasserstein space, a Riemannian manifold of real-valued distributions with a constant nonnegative sectional curvature [41, 25, 20].
Note that in our statistical mechanical theory main text Eqn.3-5, the Wasserstein- 2 distance naturally arises in the mean-field limit without assuming any a priori transportation cost.
Here we briefly review the theory of optimal transport.
Intuitively, optimal transport concerns the problem of finding the shortest path of morphing one distribution into another.
In the following, we will use the Monge formulation [65, 3].
Given probability distributions P and Q with supports X and Y , we say that T : X → Y is a transport map from P to Q if the push-forward of P through T, T # P , equals Q :
[Equation]
Eqn.72 can be understood as moving probability masses x ∈ X from distribution P to y ∈ Y according to transportation map T , such that upon completion the distribution over Y becomes Q .
We are interested in finding a transportation plan that minimizes the transportation cost as measured by some distance function d : X × Y → R :
[Equation]
The transportation plan that minimizes Eqn.73 is called the optimal transport plan T ∗ =
argmin T C ( T ; d ) .
When the distance function d is chosen to be the L k norm, the minimal cost becomes the Wasserstein- k distance:
[Equation]
In 1 -dimension, the Wasserstein- k distance has a closed form given by main text Eqn.6, and the optimal transport map has an explicit formula in terms of the CDFs: T ∗ =
Q − 1 ◦ P .
An example of the optimal transport map and how it moves probability masses between distributions is given in
Fig.8 for transport between p ( w ) =
N (0 , 1) and q ( w ) =
[Equation]
(cid:111) .
Note that in this case, the optimal transport plan is simply T ∗ ( x ) =
e µ + σx .
Now consider the manifold M of real-valued probability distributions, where points on this manifold are probability measures that admits a probability density function.
When endowed with the W k metric, ( M , W k ) becomes a metric space and is in particular a geodesic space [65, 3].
We can explicitly construct the geodesics connecting points on M .
We parameterize the geodesic by the geodesic time τ ∈ [0 , 1] .
Then given T ∗ an optimal transport plan, the intermediate probability distributions along the geodesic take the following form [65]:
[Equation]
[Figure] where Id is the identity map and P τ is a constant speed geodesic connecting P τ =
=
P and
P τ =1 = Q .
For the discrete case, we can describe the sample { w τi } from P τ in a simple manner in terms of the samples { w i } drawn from P and { ˆ w i } drawn from Q .
We can arrange the samples in the ascending order, or equivalently, forming their order statistics (cid:8) x ( i ) : x (1) ≤ ... ≤ x ( N ) (cid:9) , which can be thought of as atoms in a discrete measure.
Then in terms of the order statistics,
[Equation]
[Equation]
Upon infinitetesimal change in the geodesic time, τ → τ + δτ , the geodesic flow becomes
[Equation]
Specializing to the case discussed in main text Section 3, w ( i ) =
w τ =
( i ) is the initialization for the perceptron weight and therefore just a constant, we can promoted it w ( i ) →
w τ ( i ) to fix the overall scale in the perceptron weight, then we arrive at main text Eqn.9.
A.3 Generalization supplemental materials A.3.1 Replica calculation of generalization with sign-constraint
In this section, we calculate the sign-constraint teacher-student setup.
To ease notation, let’s denote the teacher perceptron w t ≡ w 0 and the (replicated) student perceptron w as ≡ w a .
Given random inputs ξ µ with p ( ξ µi ) =
N (0 , 1) , we generate labels by ζ µ = sgn ( w 0 · ξ µ / ||
w
||
+ η µ ) , where η µ is input noise and η µ ∼ N (0 , σ 2 ) .
Let’s denote the signs of the teacher perceptron as s i = sgn
( w 0 i ) .
The student perceptron’s weights are constrained to have the same sign as that of the teacher’s, so we insert Θ( s i w ai ) in the Gardner volume to enforce this constraint (we leave out the denominator part of V as it does not depend on data and is an overall constant):
[Equation]
We observe that upon redefining s i w ai → w ai , s i ξ µi → ξ µi , we can absorb the sign-constraints into the integration range of w from [ −∞ , + ∞ ] to [0 , ∞ ] :
[Equation]
Therefore, sign constraint amounts to restricting all the weights to be positive.
In the following, we denote
[Equation]
Let’s define the local fields as
[Equation]
We leave the average over teacher w 0 to the very end.
[Equation]
where in the last step we perform the average over noise η µ ∼ N (0 , σ 2 ) and patterns p ( ξ µi )
=
N (0 , 1) , and make use of the normalization conditions (cid:80) i ( w 0 i ) 2 = N and (cid:80) i ( w ai ) 2 = N .
Now let’s define order parameters
[Equation]
We introduce conjugate variables ˆ q ab and ˆ R a to write the δ -functions into its Fourier representations, and after some algebraic manipulations we can bring the Gardner volume into the following form ( α ≡ p/N ):
[Equation]
where ( ¯ h 0 µ = γh 0 µ ;
[Equation]
[Equation]
The energetic part G 1 is the same as the unconstrained case in [60, 22].
After standard manipulations, we have
[Equation]
Entropic part
In this subsection
[Equation] we perform the integrals in the entropic part [Equation] and we will see novel terms coming from the constraint on the student’s integration range.
We start by assuming a replica-symmetric solution for the auxiliary variables introduced in the Fourier decomposition of the δ -functions,
[Equation]
and q ab = (1 − q ) δ ab + q.
Then the entropic part,
[Equation]
where we have introduced Gaussian variable t to linearize quadratic term as usual.
Now the integral becomes n identical copies and we can drop the replica index a ,
[Equation]
We can bring the log term into the form of an induced distribution f ( w ) ,
[Equation]
Under saddle-point approximation, we obtain a set of mean field self-consistency equations for the order parameters:
[Equation]
q → 1 limit
In this limit the order parameter diverges
[Equation]
and we define the new set of undiverged order parameters as
[Equation]
Then
[Equation]
Then ⟨ w ⟩ f = 1∆ (cid:16) ˜ Rw
+ t ˜ q (cid:17) , and the integral over the auxiliary variable is dominated by values of t such that ˜ Rw
+ t ˜ q > 0 .
In the following, we denote (cid:10) [ g ( t )] + (cid:11) t as integrating over range of t such that g ( t ) > 0 .
Then the self-consistency equations Eqn.90 take a compact form (after rescaling order parameters ˜ R
→ ˜ R ∆ , ˜ q → ˜ q ∆ )
[Equation]
Eqn.90 becomes ( ˜ κ = κ/ (cid:112) 1 −
[Equation]
The free energy is (recall that γ = 1 / √ 1 + σ 2 )
G =
[Equation]
A.3.2 Replica calculation of generalization with distribution-constraint
In this subsection, we will consider the case where student weights are constrained to some prior distribution q s ( w s ) , while the teacher obeys a distribution p t ( w t ) , for an arbitrary pair q s , p t .
We can write down the Gardner volume V g for generalization as in the capacity case (main text Eqn.2):
V g = (cid:1)
[Equation]
We treat the distribution constraint q s ( w ) similar to Section A.1.1.
The entropic part of the free energy becomes
[Equation]
At the limit q → 1 , we make the following ansatz
[Equation]
Then
[Equation]
We can absorb 12 ∆ w 2 into the definition of r ( w ) , 12 ∆ w 2 + r ( w ) → r ( w ) , and 0 = ∂G 0 /∂ ∆ gives the second moment constraint, 1 = (cid:1) dwq [Equation] [Equation]
Then,
[Equation]
[Equation]
Next, we perform a saddle-point approximation on the log-term in G 0 ,
[Equation]
where w s is the saddle-point value for the weight [Equation] and is determined implicitly by
[Equation]
Note that r ′ ( w s ) is now an induced random variable from random variables w t and t .
For later convenience, we rescale r ′ ( w s ) to define a new random variable z ,
[Equation]
where we have also defined
[Equation]
The induced distribution on z is then
[Equation]
Now the entropic part becomes
[Equation]
Integrate by parts,
[Equation]
Now 0 = ∂G/∂r ′ ( w s ) gives
[Equation]
which implicitly determines w s ( z ) .
Next,
[Equation] [Equation]
[Equation]
The free energy then simplifies to
[Equation]
The energetic part as q → 1 becomes (same as the unconstrained and sign-constrained case)
[Equation]
The remaining two saddle point equations are (1) the vanishing log-Gardner volume and (2)
= ∂G/∂R :
[Equation] [Equation]
In summary, the order parameters { R, κ, u, ε } can be determined from a set of self-consistency equations:
[Equation]
where we have introduced ˜ κ = κ/ (cid:112) 1 − γ 2 R 2 , an auxiliary normal variable t ∼ N (0 , 1) , and an induced random variable z ≡ t + εw t with induced distribution
[Equation]
Note that w s ( z ) can be determined implicitly by equating the CDF of the induced variable z and the distribution that the student is constrained to:
[Equation]
Examples
(1) Lognormal distribution
In the following, we solve w s ( z ) explicitly from the CDF equation Q ( w s ) =
˜ P ( z ) .
For a lognormal teacher,
[Equation]
The second moment constraint implies µ = − σ 2 .
The induced CDF of z is
[Equation]
Let x =
(ln w − µ ) /σ,
[Equation]
Now the CDF of w s is
[Equation]
Therefore, equating ˜ P ( z ) and Q s ( w s ) :
[Equation]
We can solve for w s ( z ) by (recall z ≡ t
[Equation]
Or in terms of error functions
[Equation]
We can also calculate the initial overlap (before any learning):
[Equation]
) Uniform distribution
Assuming that both the teacher and the student have a uniform distribution in range [0 , σ ] .
The second moment constraint fixes σ = √ 3 .
We can solve (as in the lognormal example above),
[Equation]
Half-normal distribution
Assuming that both the teacher and the student has a half-normal distribution
[Equation]
The second moment constraint fixes σ = 1 , and
[Equation]
Arbitrary number of synaptic subpopulations
Just like in the case of Section A.1.2, we can generalize our theory above to incorporate distribution constraints with an arbitrary number of synaptic subpopulations.
Let’s consider a student perceptron with M synaptic populations indexed by m , w m , such that each w mi satisfies its own distributions constraints w m i ∼ Q m ( w m ) .
We denote the overall weight vector as w ≡ { w m } M m =1 ∈ R N × 1 .
The total number of weights is N = (cid:80)
Mm =1 N m , and we denote the fractions as g m = N m /N .
Since the derivation is similar to that of Section A.1.2 and Section A.3.2, we will only present the results here.
As before, the order parameters { R, κ, u, ε } can be determined from a set of self-consistency equations:
[Equation]
where ˜ κ = κ/ (cid:112) 1 − γ 2 R 2 , t ∼ N (0 , 1) .
and an induced random variable z ≡ t + εw t with induced distribution the same as Eqn.118.
Note that every w m ( z ) can be determined by equating the CDF of the induced variable z and the m -th distribution that w m ( z ) is constrained to:
[Equation]
A.3.3 Sparsification of weights in sign-constraint learning
For unconstrained weights, max-margin solutions are considered beneficial for generalization particularly for small size training sets.
As a first step toward biological plausibility, one can try to constraint the sign of individual weights during learning (e.g., excitatory or inhibitory).
In the generalization error setup, we can impose a constraint that the teacher and student have the same set of weight signs.
Surprisingly, we find both analytically and numerically that if the teacher weights are not too sparse, the max-margin solution generalizes poorly: after a single step of learning (with random input vectors), the overlap, R , drops substantially from its initial value R 0 (by a factor of √ 2 for a half-Gaussian teacher, see the blue curves in Fig.9(a).
We can verify this by calculating R 0 in two different ways.
As an example, in the following we consider the case where both the teacher and student have half-normal distributions.
(1) By definition, the overlap is R = w s · w t ∥ w s ∥∥ w t ∥ .
Since w s and w t are uncorrelated before learning ( α = 0 ), the initial overlap is then R 0 = ⟨ w s ⟩⟨ w t ⟩ || w s ||
w t
|| =
2 π ; (2) Take the α → 0 limit in Eqn.90 and Eqn.91 and calculate R 0+ = lim
α → 0+ R ( α ) = √ 2 π . Therefore, in this example R 0+ = R 0 / √ 2 .
The source of the problem is that due to the sign constraint, max-margin training with few examples yields a significant mismatch between the student and teacher weight distributions.
After only a few steps of learning, half of the student’s weights are set to zero, and the student’s distribution, p ( w s ) =
12 δ (0) + 1 √ 2 π exp {− w 2 s 4 } , deviates significantly from the teacher’s half-normal distribution (Fig.9(b)).
A.3.4 Noisy teacher
We generate examples { ξ µ , ζ µ } Pµ =1 from a teacher perceptron, w t ∈ R N : ζ µ = sgn ( w t · ξ µ /
|| w t
|| + η µ ) , where η µ is input noise and η µ ∼ N (0 , σ 2 ) .
In this subsection we present additional numerical results for the case when σ ̸ = 0 .
As in previous sections, we define the noise level parameter γ = 1 / √ 1 + σ 2 .
Our theory’s prediction is confirmed by numerical simulation for a wide range of teacher noise level γ and teacher weight distributions P t ( w t ) .
We find that distribution-constrained learning performs consistently better all the way up to capacity (capacity in this framework is due to teacher noise).
For illustration, in Fig.10 we show theory and simulation for fixed prior learning of three different teacher distributions: uniform, half-normal, and lognormal. [Figure] [Figure]
A.4 DisCo-SGD simulations Avoid vanishing gradients
Note that we often observe a vanishing gradient in DisCo-SGD when we choose a constant learning rate η 1 , and in such cases the algorithm tends to find poor margin κ which deviates from the maxmargin value predicted from the theory.
We find that scaling η 1 with the standard deviation of the gradient solves this problem:
[Equation]
where the standard deviation is computed across the synaptic index i and η 0 1 is a constant.
Mini-batches
For the capacity simulations, we always use full-batch in the SGD update, so it is in fact simply gradient descent.
However, in the case of generalization, we find that training with mini-batches improves the generalization performance, since it acts as an source of stochasticity during training.
In main text Fig.5 we use mini-batch size B = 0 .
8 P ( 80% of examples are used for each SGD update).
When we vary teacher’s noise level, we find that scaling B with γ improves the quality of the solutions, as measured by the generalization performance (or equivalently, the teacher-student overlap).
Generally, the more noisy the teacher is, the smaller the mini-batches should be.
This is because smaller mini-batch size corresponds to higher stochasticity, which helps overcoming higher teacher noise.
Parameters
All the capacity simulations are performed with the following parameters N = 1000 , η 01 = 0 .
01 , η 2 = 0 .
6 , t max = 10000 , where t max is the maximum number of iterations of the DisCo-SGD algorithm.
All results are averaged over 300 realizations.
In main text Fig.4, the experimental [38] parameters are g E = 45 .
8% , σ E = 0 .
833 , σ I = 0 .
In main text Fig.5(a): We show the teacher-student overlap as a function of α .
Dots are simulations performed with series of student distribution from σ s = 0 .
1 to σ s = 1 . 4 , where the teacher distribution sits in the middle of this range, σ t = 0 .
Each such simulation is performed with fixed σ s and varying load α ∈ [0 . 05 , 2 . 5] .
In main text Fig.5(b): we show the empirical weight distributions found by unconstrained perceptron learning for α ∈ [0 .
In main text Fig.5(c) we show optimal student distribution for α ∈ [0 .
Note that optimal prior learning approaches the teacher distribution much faster than unconstrained learning.
All the generalization DisCo-SGD simulations are performed with the same parameter as in the capacity DisCo-SGD simulations, but with two additional parameter: teacher’s noise level γ and
SGD mini-batch size B .
For the simulations in Fig.10 we use γ = 0 . 4 , B = 0 .
2 P ; γ = 0 .
55 , B = 0 .
[Table]
[Table]
[Table]
[Table]
[Table]
[Table] [Table]
A.5 Replica symmetry breaking A.5.1 Bimodal distributions
In deriving the capacity formula, we have assumed replica-symmetry (RS).
It is well-known that replica-symmetry breaking occurs in the Ising perceptron [52, 13], so it is natural to ask to what extent our theory holds when approaching the Ising limit.
Let’s consider a bimodal distribution with a mixture of two normal distributions with non-zero mean centered around zero,
[Equation]
The second moment constraint requires µ 2 + σ 2 = 1 .
We can gradually decrease the Gaussian width σ , or equivalently µ = √ 1 − σ 2 (which we call ‘separation’ in the following) and compare the capacity theoretically predicted by the RS theory and numerically found by the DisCo-SGD algorithm.
In Fig.11 we can see that the simulation agrees well with the RS theory until one gets very close to the Ising limit ( µ = 1 ).
To understand finite size effects, we extrapolate to the infinite size limit ( N → ∞ ) in Fig.12, and found that the deviation from RS theory has a sharp transition near µ = 1 , marking the breakdown of the RS theory.
Ising perceptron
It is also interesting to compare our distribution-constrained RS theory to the unconstrained RS theory.
In this Ising limit,
[Equation] [Figure] [Figure]
and CDF
[Equation]
Equating Q ( w ) with the normal CDF P ( x ) and solve for w ( x ) , we find w ( x ) = sgn
( x ) .
Then dw/dx = [-
-]2 δ ( x ) and (cid:10) dwdx (cid:11) x = 2 √ 2 π .
Therefore,
[Equation]
which is exactly the same as the prediction from the unconstrained RS theory [52, 13].
In contrast, the exact capacity of Ising perceptron with replica-symmetry breaking is α c ≈ 0 . 83 .
For comparison, we have included these values in Fig.12(a), as well as the capacity found by the state-of-the-art supervised learning algorithm
(Stochastic Belief Propagation, SBPI [10]) for Ising perceptron.
A.5.2 Sparse distributions
For a teacher with sparse distribution, p (
[Equation]
We found that the simulations start to deviate from the theory, and the reason might be due to replica symmetry breaking.
In Fig.13, we use the optimal prior learning paradigm similar to main text Fig.5.
We see that our RS theory no longer gives accurate prediction of overlap in this case. [Figure]