Synaptic Weight Distributions Depend on the Geometry of Plasticity
Authors: Roman Pogodin
Jonathan Cornford
Arna Ghosh
Gauthier Gidel
Guillaume Lajoie
Blake Aaron Richards
Venue: ICLR 2024 spotlight (2024)
Keywords: synaptic weight distributions
synaptic plasticity
biologically plausible learning
mirror descent
Primary area: applications to neuroscience & cognitive science
TL;DR: different algorithms result in different synaptic weight distributions, and the log-normal one is consistent with exponentiated gradient

Abstract: A growing literature in computational neuroscience leverages gradient descent and learning algorithms that approximate it to study synaptic plasticity in the brain. However, the vast majority of this work ignores a critical underlying assumption: the choice of distance for synaptic changes - i.e. the geometry of synaptic plasticity. Gradient descent assumes that the distance is Euclidean, but many other distances are possible, and there is no reason that biology necessarily uses Euclidean geometry. Here, using the theoretical tools provided by mirror descent, we show that the distribution of synaptic weights will depend on the geometry of synaptic plasticity. We use these results to show that experimentally-observed log-normal weight distributions found in several brain areas are not consistent with standard gradient descent (i.e. a Euclidean geometry), but rather with non-Euclidean distances. Finally, we show that it should be possible to experimentally test for different synaptic geometries by comparing synaptic weight distributions before and after learning. Overall, our work shows that the current paradigm in theoretical work on synaptic plasticity that assumes Euclidean synaptic geometry may be misguided and that it should be possible to experimentally determine the true geometry of synaptic plasticity in the brain.

Published as a conference paper at ICLR 2024
SYNAPTIC WEIGHT DISTRIBUTIONS DEPEND ON THE GEOMETRY OF PLASTICITY
Roman Pogodin'
Jonathan Cornford"
Arna Ghosh McGill & Mila
McGill & Mila
McGill & Mila
roman pomondpogodinae n@mi quebec
cornfor @@mi 1a. quebec
Gauthier Gidel
Guillaume Lajoie
Blake Aaron Richards McGill Mila & CIFAR!
Université de Montréal & Mila
Université de Montréal & Mila
ABSTRACT
A growing literature in computational neuroscience leverages gradient descent and learning algorithms that approximate it to study synaptic plasticity in the brain. However, the vast majority of this work ignores a critical underlying assumption: the choice of distance for synaptic changes i.e. the geometry of synaptic plas- ticity. Gradient descent assumes that the distance is Euclidean, but many other distances are possible, and there is no reason that biology necessarily uses Eu- clidean geometry. Here, using the theoretical tools provided by mirror descent, we show that the distribution of synaptic weights will depend on the geometry of synaptic plasticity. We use these results to show that experimentally--oseree log-normal weight distributions found in several brain areas are not consistent with standard gradient descent (i.e. a Euclidean geometry), but rather with non- Euclidean distances. Finally, we show that it should be possible experimentally test for different synaptic geometries by comparing synaptic weight distributions before and after learning. Overall, our work shows that the current paradigm theoretical work synaptic plasticity that assumes Euclidean synaptic geometry may be misguided and that it should be possible to experimentally determine the true geometry of synaptic plasticity in the brain.
INTRODUCTION
Many computational neuroscience studies use gradient descent to train models that are then com- pared the brain Schrimpf et al. (2018): Nayebi al. (2018); Yamins et al. (2014); Bakhtiari al. 220TT) Flesch et al. (2022), and many others explore how synaptic plasticity could approximate gradient descent Richards al. 2219 One aspect of this framework that is often not explicitly considered is that in order to follow the gradient of a loss in the synaptic weight space, one must have a means of measuring distance the synaptic space, i.e. of determining what constitutes a large versus small change in the weights Carlo Surace al. (2018). In other words, whenever we build a neural network model we are committing to a synaptic geometry. Standard gradient descent assumes Euclidean geometry, meaning that distance in synaptic space is equivalent to the L2-norm of weight changes.
There are however other distances that be used. For example, in natural gradient descent the Kullback- divergence of the network's output distributions is used as the distance |Martnn 2020). More broadly, the theory of mirror descent provides tools for analyzing and building algo- rithms that use different distances in parameter space Nemirovskij & Yudin (adin((1883 (1883); Beck & Teboulle (003). , which has proven useful in variety of applications (Shale--Swwarz et al. (2012); Bubeck ct al. Lattimore & Gyorgy (202I). However, within computational neuroscience the question of synaptic geometry is often overlooked: most models use Euclidean geometry without considering other options Carlo Surace et al. (0018). This assumption has no basis in neuroscience data, so how could we possibly determine the synaptic geometry used by the brain?
* EEqual contribation. DDett. CS and OR: Dept. Maths and Stats; Dept. of Neurology & Neurosurgery, School of Computer Science, Montreal Neurological Institute; "Learning in Machines and Brains Program.

---

Published as a conference paper at ICLR 2024
Here, using tools from mirror descent Nemirovskij & Yudin (1983): Beck & Teboulle (200)), show that synaptic geometry can be determined by observing the distribution of synaptic weight changes during learning. Specifically, we prove that in situations where synaptic changes are rel- atively small the distribution of synaptic weights depends on the synaptic geometry (with mild as- sumptions about the loss function and the dataset). We use this result to show that the geometry defines a dual space in which the total synaptic changes are Gaussian. As a result, if one can find dual space in which experimentally observed synaptic changes are Gaussian, then one knows the synaptic geometry. Applying this framework to existing neural data, which suggests that synap- tic weights are log-normally distributed [Song et al. (005);; Loewenstein et al. (2011); Melander et al. (2021); Buzsáki & Mizuseki (2014). we conclude that the brain is unlikely to use a Euclidean synaptic geometry. Moreover, we show how to use our findings to make experimental predictions. In particular, we show that it should be possible to use experimentally observed weight distributions before and after learning to rule out different candidate geometries. Altogether, our work provides a novel theoretical insight for reasoning about the learning algorithms of the brain.
1.1 RELATED WORK
There is large and growing literature on approximating gradient-based learning in the brain Lillicrap et al. (2016); Liao et al. (2016); AArout et al. ( Podlaski & Machens (2020); Clark et al. (2021). The vast majority of this work assumes Euclidean synaptic geometry Carlo Surace iet al. (018)) Notably, even if the brain does not estimate gradients directly, as long as synaptic weight updates relatively small. then the brain's learning algorithm must be non-orthogonal to some gradient in expectation |Richards & Kording (2023). As such, our work is relevant to neuroscience regardless of the specific learning algorithm used by the brain.
Our work draws strongly from the rich and long-standing literature on mirror descent, which was originally introduced 40 years ago for convex optimization Nemirovskij & Yudin (1983). Recent
years have seen a lot of work in this area [Beck & Teboulle (2003); Duchi et al. 2010); Bubeck
al. (013); Ghai et al. 22020: Lattimore & Gyorgy| (2021), especially in online optimization
settings such as bandit algorithms Shalev-Shwart et al. (2012); Lattimore & Szepesvári (2020). More recently, researchers have started apply mirror descent to deep networks and have used it to
try to develop better performing algorithms than gradient descent Azizan & Hassibi (2018); Azizan jet al. (202T) This work is related to natural gradient descent which also explores non- non-Euclidean
geometries Amari [1985 1998); Ollivier et al. (2017).
Finally, our work is relevant to experimental neuroscience literature on synaptic weight distributions. Using variety of techniques, including patch clamping SSng et al. (2005) and fluorescent imaging Melander et al. [202[); [oewenstein et al. (201T); Vardalaki et al. (2022), neuroscientists have ex- plored the distributions of synaptic strengths across variety of brain regions and species. This work has generally reported log- normal distributions in synaptic weights Song et al. (2005); Loewenstein et al. (2011): Melander al. (002), though, recent study observed a more complicated bimodal distribution in log-space Dorkenwald et al.] (2022). Moreover, the parameters of log-normal dis- tributions observed in primary auditory cortex Levy Reyes (2012) are close to optimal in terms of perceptron capacity Zhong el al. (2022). Our work connects this experimental literature to our theoretical understanding of learning the brain it makes it possible to test theories of synaptic geometry using weight distribution data.
2 MIRROR DESCENT FRAMEWORK
In order to derive our core results we will rely on tools from mirror descent. To introduce it, we first revisit gradient descent. Assume we have loss function I(w). In gradient descent, we choose the next point in weight space wf+1 from the current point w" by minimizing a linearized version of ((w) with a penalty for taking large steps in weight space (where the strength of the penalty is controlled by the learning rate 17). Importantly, penalizing large steps in weight space necessitates a distance function. If we choose the squared 7.2 n.r-norm as our distance function we obtain:
w ++1 = arg min g(w, w'), 9(w, wt) = 1(w')= + VI( w') + 2n
w
w
(1)
w
2

---

Published as a conference paper at ICLR 2024
This choice of distance function results in the standard gradient descent update when we solve the unconstrained optimization problem in Eq. (
w w' Vl(w"))
(2)
In mirror descent, we consider more general set of distance functions provided by the Bregman divergence D& w, w) Bregman (1977). For a strietly convex function, called a potential, o( ((w):
Do(w, w) = ((w) - ((w) - vØ(w) (w -W)
(3)
Re-writing Eq. ① into this more general form we get:
w arg min go (w, (w, w'), go(w, w') ((w")= + vi(w") (w - w') + - (w, (w, w²). (4)
w
We can express the update in closed form using the gradient of the potential:
primal
dual space
Ve(ww')
Vo(w#+)) Ve(w") Vl(w"). As o is strictly convex, V is invertible, so:
(5)
w 11
w V (7 17 v((w"))
(6)
V
w Vi nommal grlcient
VI: grodion karning @
Ve(w (w("") V mimo Sesceet
The update moves w from the "primal" space to the "dual" via Ve(w), performs regular gradient descent update in the dual space, and then projects back to the primal space (Fig. [
Figure 1: Mirror descent dynamics.
We consider p-norms < w) = Il w l| and negative entropy
o(w)=E, wy log |w:]. The former for p= recovers standard gradient descent. The latter leads to the exponentiated gradient algorithm Kivinen & Warmuth 1997) (O denotes element-wis product):
w w+1 w e(-n/V((") o sign
Fig. 2 shows how different potentials change problem geometry and, therefore, solutions.
Equation 55) already predicts main result: if updates in the dual space approximately inde- pendent, or contain noise, over time the sum of the updates will look Gaussian by the central limit theorem (since they linearly in the dual space). In fact, Loewenstein al. (201T) showed that synaptic weights (of rodent auditory cortex) stay log-normal over time (implying Gaussian changes in the log space) and follow multiplicative dynamics in synaptic weight space. This is consistent with the negative entropy potential and in particular Eq. 77). This data, however, wasn't collected when training mice on specific task. is not clear if Geaning-driven dynamics follows the same update geometry. Here, we develop theory for distinguishing synaptic geometries during training.
Eq. (6) allows us to understand synaptic weight updates via two independent terms: a credit sig- nal and an intrinsic synaptic geometry. The credit signal is the gradient Vi(w") or a more biologically plausible approximation. The synaptic geometry is determined by the potential ø; it captures how the credit signal Vi(w') is used to update synaptic weights. Note however, we are not necessarily propos- ing that this separation reflected in underly- ing neurobiological processes, just that it can be used as model for understanding weight updates. Most studies of biologically plausi- ble deep learning [illicrap et al. (2016); Liao et al. | (2016); Akrout et al. (2019): Podlaski & Machens (202)): Clark et al. (2021) use Eu- clidean distance by default, effectively ignoring the second term (although see Carlo Surace et al. centrate on the second term, and derive results that independent of the first term. Therefore, and crucially, previous theoretical studies that derive biologically plausible estimates of credit signals
requtive entropy NNE)
I(w1, W22
requtivo ertiopy NE I(w1, W2 2-namm
Z-aam
88
64
60
0.2
24
VoNN (1h)
V ONN (w₂)
Figure 2: Eq. (4) dynamics for I(w1, w2) 1(1 Wr+
Wh J twi+ Blue: negative entropy (NE, Eq. 7D: green: gradient descent/2-norm (Eq, Z);; orange: 3-norm. Left: loss surface (and level sets shown below it) and dynamics in the regu- lar (W1, w22 coordinates. Right: same, but in the (VONE(W)). (ÓNE(W?)) coordinates.
(018)) and Schwarz et al. (021 l).. Here we con-
fully compatible with our work.
3

---

Published as a conference paper at ICLR 2024
2.1 IMPLICIT BIAS MIRROR DESCENT
Gradient descent in overparametrized linear regression finds the minimum norm solution: a set of weights closest to the initial weights w according to the 2-norm distance Gunasekar et al. Zhang et al. (202T). This bias towards small changes is referred to as the implicil bias of the algorithm. Recently, similar result was obtained for mirror descent, wherein the implicit bias will depend on the potential o GGunasekar et al. (2018). Here we discuss this result and its applicability to gradient- based learning in deep networks.
Linear regression. To begin, consider linear problem: predict y" from x" as 0" (x" w. For
N points and D-dimensional weights, we write this as y = Xw. Gunasekar et al. (2018) showed that for losses ((. y) with unique finite root (I(T, - 0 iff y y), the mirror descent solution w (assuming exists) is the closest w to the initial weights w w.r.t. the Bregman divergence Dø:
W
arg min Do Do.w.w"). y=x
(8)
An example of such setup is the MSE loss (( y) 1 y) and overparametrized model with 2 N. For a fixed y, this gives rise to an anti anti-Hebbian gradient (0 y)x.
Deep networks. A crucial condition is hidden in the proof of the above result: the gradient updates VI(ō", yy") have to span the space of x". For a linear model 7 = x w, this is naturally satisfied as VI(ō", y) "'("".") x" For deep networks, Azizan et al. (2021) showed a similar to Eq. (8) dz result for solutions close to the initial weights, and empirically noted that different potentials result in different final weights. Thus, we can make a similar assumption: if the solution is close to the linearize given deep network o (w,,) x) around the initial weights w":
initial weights, then
flin x, w²) = f(w", x) + f x (w-w").
This function is linear in the weights w, but not in the inputs x. Intuitively, if w doesn't change a lot from w a linearized network should behave similarly to the original one with respect weight changes. The linear approximation moves us back to a setting akin to linear regression except that the gradients now span ff(w", xx) rather than x, Thus, for linearized networks, Eq. ( becomes:
W
arg min Do(w,w") W: f
(10)
One may wonder whether it is appropriate to assume that a solution exists near the initial weights when considering real brains. There are two reasons that we argue this is appropriate in our context. First, in biology, animals are never true "blank slates". they instead come with both life experience and evolutionary priors. As such, it is not unreasonable to think that in real brains the solution in synaptic weight space often sits close to the initial weights, and moreover, there would be strong evolutionary pressure for this to be so. Second, neural tangent kernel (NTK; Jacot et al. (0188) theory shows that, with some assumptions, infinite width networks are identical to their linearized versions, and for finite width networks, the linear approximation gets better width increases Lee let al. (2019) (although learning dynamics don't always follow NTK theory, see e.g. [Bordelon & Pehlevan (2022)) Thus, for very large networks (such as the mammalian brain) it is not unreason- able to think that a linear approximation may be appropriate.
Below, we will develop theory for Eq. 8 (linear regression), and then experimentally show that the results hold for fine-tuning of deep networks (which are close to linearized networks in Eq. (10))
3 WEIGHT DISTRIBUTIONS IN MIRROR DESCENT
Our goal in this section to derive a solution for the distribution of the final synaptic weights s as a function of the potential ø. To do this, we begin by noting that Eq. (8) can be solved using Lagrange multipliers À E N that enforce y = Xw; the Lagrangian is: (ww = Do wo) + (y X T Solve for the minimum of L by differentiating it w.n.t. w and setting it to 0:
(11)
vo(w)- - vo(wo) x T x= = 0. i L(w, X)
(12)
8
4

---

Published as a conference paper at ICLR 2024
Since VΦ- is the inverse of Vé, obtain: - V0-1(Vo(()) (Vo(w") (ooww)+X']] x)
(13)
where À will be chosen to satisfy = Xw
To obtain closed-form solution, mapping with the potential is invertible, so VØ- (Vel ( w
linearize Eq. 13 around Vø(w") using the fact that the
w V (Vo(w))) + 22 (Vo(w")) x w H H
(14)
where we denoted Ha-1 = 22h- (Vib( (ww(()))..))))))>)))) For potentials that couple weights together, the Hessian will be non-diagonal. However, we use "local" potentials in which each entry of Vo(w) depends only on W; (i.e. Φ(w) E; f(w;) for some function f). For such potentials, the Hessian becomes diagonal. Since the mapping w.r.t. the potential is invertible, we can use the inverse
function theorem to compute the Hessian HA 1: Ha v2 (z) (V²o(w")) for Vo(w>), assuming e is twice continuously differentiable and H is non-singular.
Since H is ultimately function of ø. we have the structure of our solution. However, we need to solve for A in order to satisfy y = X Assuming that H&-1 positive-definite and that X E R.ND is rank N (D 2 N) for the sake of invertibility, À can be approximated by Ā as:
>= = (XH - Xw-)
(15)
With Eq. [15). we can now state the main result. Intuitively, what this result will show is that (1) we can know the shape of the distribution of the final weights, (2) that shape is independent of the loss function and the dataset, but not the initial weights. More formally, this is stated as:
Theorem 1 (Informal). Consider N i...d. samples y" ,",x", such that: x" RR" are zero-mean and bounded; pairwise correlations Cıj = =Ex;x and d3 = =ov((<);) (x7))) between entries of single x" decay quickly enough so >5= (cill const and E5=1 1dy const for all i; y" (x") T ww"; the teacher weights w" and the initial weights w have zero-mean, O(1/D) variance, i..d entries, and finite 8th moment: V26 -ww") has finite 1st and 2nd moments.
Then for 7= (XH+-IXT) (yxww"), individual entries of X T converge (in distribution) to Gaussian with a constant variance o as D. N + 2o with N- o (/5+66)))) (for any > 0):
Dh
x xT)),
P N((,,0)
(16)
VN
where h = E ( (a scaling factor that is identical for all entries, which is not equivalent Hg-1), and o3 depends on the distributions of inputs and initial weights. The whole vector X TT X converges to a Gaussian process (over discrete indices) with weak correlations for distant points.
Proof sketch. First, we show that scaled (XH2 .xT) behaves like an identity for large N, D. Then, we show that (X T (y- (y-xx Xww)), a sum of N exchangeable variables (i.e. any permutation of these points has the same distribution) that satisfy conditions for central limit theorem for such variables. Finally, we convert this result into a convergence to Gaussian process.
We postpone the full proof, along with version of the theorem for generic labels (without the teacher weights w*), to Appendix [ To unpack this theorem a bit more for the reader, combined with Eq. (3)). the theorem shows that given some initial weights, w the distribution of w will depend on that initial distribution plus term that converges to a Gaussian (in the dual space), as long as the loss has a unique finite root (defined before Eq. (D) and the dataset correlations are small (Cij, dy in Theorem Ip. As such, the distribution of the solution weights will depend on the initial weights and the potential, but not on the loss or data. To clarify the assumptions used, we need the pairwise correlations Cjj. d'y between distant inputs to be small. This is a common feature in natural sensory data Ruderraa (1994) and neural activity Schulz et al. (2015). Finally, we note that the N scaling with D results from generic large deviations bound and could likely be improved under some assumptions. In Section 4 we experimentally verify that = D and = D0.75 result in Gaussian behavior in the tested settings.
5

---

Published as a conference paper at ICLR 2024
Practical usage of our result. Since we expect the weight change in the dual space to look Gaus- sian, we also expect the weight change in the dual space of a wrong potential to not be Gaussian. If we denote the true potential Vo f and another potential Vo' = our result states that the change in the dual space is Gaussian: /( ) f(w") = where & is Gaussian. If we use the wrong potential, i.e. f', get f" is approximately an identity, so we would see a (slightly less than for f) Gaussian change. If not, the change will be non-Gaussian due the nonlinear (-)) transformation. Therefore we will be able to determine the potential used for training by finding the one with the most Gaussian change.
(f ( If and are similar, then we (+)
Applicability other learning rules. Our result in Theorem [ is made possible by two compo- nents: the structure of mirror descent solutions (see Eq. (8) and the Gaussian behavior of large sums. However, the mirror descent framework can be applied to any learning rule, if we replace the gradi- ent w.r.t. the loss in Eq. (5) with generic error term g': Vø(('=I) small and approximately independent (e.g. exchangeable) error terms g' their total contribution to the weight change would also follow the central limit theorem, resulting in small Gaussian weight change. However, for generic error terms we cannot rely on Eq. (8) and can only provide generic conditions for Guassian behavior. Therefore, while our work leverages gradient- optimization, the intuition we present applies more broadly.
Vo(w") - -".g. This way, for
Theorem applicability in lazy and rich regimes. The variance of the Gaussian term in Theorem depends on the choice of the potential (via h). This means that the magnitude of the change from w to w will depend
ø. Moreover, if the "learned" term X X in Eq. (13) is not small compared to Fe(w"), ). our theory will not be applicable since the lincarization will not be valid. Therefore, the applicability of our theorem is related to the question of "rich" versus "lazy" regimes of learning. Our theory is valid only in the lazy regime, in which learning is mostly linear Chizat et al. (2019).
However, whether we are in the rich or lazy regime turns out to depend on the potential. Assume for simplicity that all initial weights are the same and positive: wy w a/ D. The standard lazy regime corresponds to large weights with a a= 1, and the standard rich regime corresponds to small weights with a 1/v/D. We can show (see Appendix A. for a derivation) that for p-norms, the dual weights are asymptotically larger than the weight changes as long as >> /N/D (for N data points and D weights). For negative entropy, this bound is more loose: > (/ log D). Therefore, for the same weight initialization and dataset size. negative entropy would typically produce smaller updates. For the standard for deep networks initialization with a our theory should be applicable to datasets with N < D) (e.g. N= D..5).
4 EXPERIMENTS
Here empirically verify our theory under conditions relevant for neuroscientific experiments. We use PyTorch Paszke et al. (2019) and FFCV library for fast data loading Leclere et al. (2022). The experiments were performed on a local cluster with A100 NVIDIA GPUs. Experimental details provided in Appendix|] B Code is available at ggthubbonnooonngggoooiimmrmeeere
4.1 LINEAR REGRESSION
We begin by testing whether the theorem holds in the case of linear regression, where we know that Eq. (8) holds. However, for experimental applicability, we want verify that the histogram of & over entries i is Gaussian even when we draw it from a single network (rather than several initializations of a network), which would be the condition that holds in real neural experiments. As well, we want to verify that the change of weights &
is smaller than Ve((w!).
For each potential, we draw 2; ~ (/ + (-0.5 0.5) such that the correlation structure is as follows: Ex, x =0 (-1))==l/-++ii--j for 1, 2. Note that is more relaxed condition than the correlation assumption used in proving Theorem [ This is, therefore, a test of whether the theorem applies more broadly. In our experiments here, the initial weights are drawn from un; ~ ( 1/D) for width D, and labels are drawn as y = +1 (equal probability). We then optimize (y E; I, for networks of different widths D and N = DO.5 D0.75 Strictly speaking only N ((D1((++)) satisfies Theorem so again, we are testing whether the theorem holds empirically
when relax our assumptions. We measure two things: the magnitude of weight changes in the
6

---

Published as a conference paper at ICLR 2024
CDF difference
magnitude of change
B
normal initialization
log-normal initialization
ACDF
MeM 1.0
ACDF
4CDF
0.2
D D25
N-04.5
0.7
Gaussian for
NE
DD.
3-nomm
0.1
0.
0..
0.3
0.2
0.2
0.1
0.1
0.0
0.0
102
10.
10
10
2.5
number of weights
number of weights
evaluated potential
evallated potential
Figure 3: A. Linear regression solutions for negative entropy (NE) and fast correlation decay (c= 2). Left: integral of absolute CDF difference (ACDF) between normalized uncentered weights and
(V001) 1). Right: magnitude of weight changes relative to the initial weights in the dual space (Aø).
Solid line: median over 30 seeds; shaded area: 5/95% percentiles; pink: N DD.5, blue: N = D0.76
B. ACDF for Gaussian (left) and log-normal (right) weight initializations and a Gaussian addition w.r.t. o evaluated on another potential o' (e.g. blue lines are sampled for NE but evaluated on every potential). Solid line: mean over 30 seeds; shaded areas: mean + standard deviation.
dual space Aδ(w) = / (w (i.e. & normalized relative to the initial weights) and ACDF dt (CDD CDF (t) (i.e. the difference between the empirical cumulative density function, (C) and the standard normal (DF
First, we find that ACDF always relatively small (S 0.2) and converges towards zero as the number of weights increases (Fig. BA for negative entropy with entropp to the appendix; see Fig. [6. Hence, we find that the weight changes in the dual space behave like a Gaussian when we know what potential was used for training, confirming Theorem [] Second, we find that the magnitude of weight changes in the dual space A w) is much smaller than 1 for large widths, justifying the linearization Eq.
2. The full results postponed
4.2 ROBUSTNESS TO POTENTIAL CHANGE
We also need test what happens if we don't 't know the true potential used for training, here denoted 6, which is the case for neural data. As discussed after Theorem | for d' close to the true potential we should see Gaussian changes, and for ø' far from Φ - Gaussian ones. We test negative entropy and 2/3-norms as é, and p-norm with E [1.1,3] 3] and negative entropy as ø'; w is drawn from either Gaussian or log-normal distributions with approximately the same variance. We draw a Gaussian E to compute w using & from & = V -) - Fe(w"). We find that the "empirical" change for d', & Vo' "WW'"""V""'""""
is indeed Gaussian only for small variations from the true potential Φ (Fig. BB. In particular, the distinction between negative entropy and other, non- multiplicative updates, is very pronounced For For á o= 3-norm and log-normal initialization, the range
for a Gaussian for almost reached p = 2.3 (Fig. 3B, bottom, orange line). Thus, we hypothesize potential. o' that is similar to the true potential Ø, get nearly Gaussian &s. In contrast, if we have the wrong form for the potential (e.g. the potential the negative entropy but we hypothesize 3-norm potential), then we get clearly non-Gaussian distribution. As such, we use the distribution of Ey to test hypotheses about synaptic geometry.
4.3 FINETUNING OF DEEP NETWORKS
We expect our theory to work for networks that behave similarly to their linearization during learn- ing. One such example is a pretrained trained network that's being finetuned on new data. If the new data comes from a similar distribution as the pretraining data, then the linearized network should approximate its non- oon-einerr counterpart well with respect to weight changes Mohamadi et al. (022); Ren et al. (2023). This is also the scenario we expect to see in neural data if we are train- ing an animal on an ethologically relevant task that matches its prior life experience and innate capabilities to a reasonable degree.
We use normalized hut uncentered weight distributions, since we predict the change & to be zero-mean. Negative entropy is placed at 1, as for weights that sum to 1, p- noroms equivalent to the shifted and scaled negative Tsallis entropy, which tends to negative entropy as p 1.

---

Published as a conference paper at ICLR 2024
&CDF
negative entropy
ACDF
2-norm
evaluation on different potential
ShuffiNNee Effirretteet
trained on:
negative enrropy (NE)
2-nom
2 norm
RT
ACDF x BtufkoNet
by
enthet
03
as
0.05
(34
K34
0A-
0.00
16
1ST
nutrbee of weights
AGDF
3-norm
wei Ax
normalized (uncenterad] weights &CDF
ERocenthe b8
Roohort
Resherdd
OfRect RT
RO 8
0.15
RT
bo
53
RCDF 0.155 &CDF 0.086
NRAT
0.6-
0.10
b1
0.4
0.00
0.2
0.00
20.5
10
157
numboefoo
0.0
1.5
2.5
NF
15
25
NR 15 7
NF 1.5 25 x-axiz: ovaluated on NE/ p-norm
3 N5 2 25
weights
Figure 4: Finetuning on 10 randomly sampled ImageNet validation subsets (N = D0.5 data points). A. Integral of absolute CDF difference (ACDF) between normalized uncentered weights and (0, 1) for networks trained with different potentials. Circle: individual value; bar: mean over seeds. Pink box (bottom right): examples of weight change histograms (pink) plotted against A( 1) (black). B. Same A., but ACDF is calculated w.r.t. other potentials.
Finetuning deep networks still breaks several assumption of Theorem D the initial weights are not ii.d., the network is not linearized (although potentially close to it), and the data correlations are unknown. Moreover, the cross-entrope loss we use is technically not a loss with a unique finite root. However, the readout layers in the deep networks that we test have fewer units than the number of ImageNer Deng et al. (2009) classes, so the activity in the readout layer is not linearly separable and the weights do not tend to infinity. Yet, they still large and change significantly during training, so we exclude the readout layer when assessing weight distributions (along with biases and BatchNorm parameters). We use networks pretrained on ImageNet, and finetune them 100% accuracy on subset of ImageNet validation set. The pretrained networks have not seen the validation data, but they already perform reasonably well on it, so there's no distribution shift in the data that would force the networks to change their weights lot. We used N D data points for D weights (N D0.75 exceeded dataset size) and four architecture types: ShuffleNet v2 Ma et al. (2018) (x0.5/x1), EfficientNet Tan & Le (2019) (b0-b3), ResNet He et al.] (2016) (18/34), and CORnet Kubilius et al.) (2019) (S/RT) chosen to span a wide range of parameter counts (0.3M to 53.4M). CORnets were also chosen since they mimic the primate ventral stream and its recurrent connections.
For all three tested potentials, weight changes are close to a Gaussian distribution (ACDF < 0.2; Fig. 4A). ShuffleNet v2s and ResNets reach consistently more Gaussian solutions than EfficientNets and CORnets, despite ShuffleNet v2s having fewer parameters. This suggests that some trained ar- chitectures are more "linear", which may be due to the architecture itself, the way the networks were trained. The magnitude of changes in the dual space was typically smaller than 0.2 (see Ap- pendix B Just like for the toy task in Fig- 38, if we trained with specific potential, Φ. we only observed Gaussian changes when the hypothesized potential, d', was close to 6 (Fig. 4B). The only exceptions were 3-norm- trained EfficientNets 0,1.3 and CORnet-S, for which the 2-norm solution was slightly more Gaussian than the 3-norm solution even when & was 3-norm. However, this dif- ference was small, and the networks had overall worse fit than other architectures. In all cases, using negative entropy as ø' provided the best fit for negative entropy-traince networks and the worst fit for other potentials. This is important because negative entropy results in multiplicative weights up- dates, while p-norms result in additive updates, which represent very distinct hypotheses for synaptic plasticity mechanisms. Taken together. these empirical results suggest that Theorem Tholds more generally and can be used to infer the distribution of weight changes beyond linear regression.
[
4.4 ESTIMATING SYNAPTIC GEOMETRY EXPERIMENTALLY
Previously we showed that Theorem I holds when finetuning pretrained deep networks, despite non-i.i.d. initial weights and an unknown input correlation structure. As such, these finetuning experiments indicate our theory is applicable to neuroscience experiments when measure the distribution of weights before and after learning, and use the histogram of weight changes to estimate the synaptic geometry. In this section we apply this technique to experimental neuroscience data.
8

---

Published as a conference paper at ICLR 2024
A
initial weights for NE
final weights
B
initial weights for 3-norm
final weights
1.2
Init weghts data
final weights data
init weights data
finsl weights tu data
1.0
0.8-
p-0.77
0.8
p-0.5
p=0.23
02
10
10"2
spine volume
10
10P
10
10"2
spine volume
10-1
10P
10-3
10 spine volume
10-1
105
10
10
spine volume
12-1
10P
Figure 5: Observed vs, modeled synaptic distributions. Pink: spine volume (um²) from Dorkenwald et al. (2022); vertical bars: point mass initializations. A. Left: initial weights for negative entropy (NE); right: final weight for a Gaussian change in the dual space. B. Same as A, but for 3-norm.
Specifically, we find that if one knows the initial weights, w then it is possible to distinguish be- tween different potentials. But if the initial weights are unknown, then multiple potential functions can fit the data. To show this, we use data from a recent experimental study that measured ana- logue of synaptic weights (synaptic spine size) using electron microscopy Dorkenwald et al. (2022) (pink histogram in Fig. 5). In the log space, Dorkenwald et al. (2022) found that the distribution was well modelled by a mixture of two Gaussians with approximately the same variance (pink histogram
in Fig. 5 instead of one Gaussian as has been previously reported Loewenstein al. (201 1b. First, we show that the experimental data is consistent with the negative entropy potential (Fig. ( when using the following initialization: consider elements of w equal to constant 141 with probability p. and to μ12 with probability 1 1-p. Then, V Vé(w>> = Vo(w") + will be mixture of two Gaussians with the same variance but different means, and Fig. [A shows the that fits the experimental data (parameters from Dorkenwald et al. (2022) with equalized variance, see Appendix B.44) However, the data also consistent with the 3-norm potential for different initialization: if w is a mixture of constant and log-normal, then adding an appropriate & in the dual space also results in weight distribution that resemble a mixture of two log-normals and appears to the experimental data well (Fig. B.. Thus, if our goal is to determine the synaptic geometry of the brain, then it is important to estimate the synaptic weight distributions before and after training. Nevertheless, with this data we can rule out Euclidean synaptic geometry, and if we do have access to w then results show that it is indeed possible to experimentally estimate the potential function.
5 DISCUSSION
We presented a mirror descent-based theory of what synaptic weight distributions can tell us about the underlying synaptic geometry of a network. For range of loss functions and under mild as- sumptions on the data, we showed that weight distributions depend on the synaptic geometry, but not on the loss the training data. Experimentally, we showed that our theory applies to finetuned deep networks across different architectures (including recurrent ones) and network sizes. Thus, our theory would likely apply as well to the hierarchical, recurrent architectures seen in the brain. Our work predicts that if we know synaptic weights before and after learning, we can find the underlying synaptic geometry by finding a transformation in which the weight change is Gaussian.
It is important to note that by adopting the mirror descent framework we made the assumption that the brain would seek to achieve the best performance increases for the least amount of synaptic change possible. But, these results could be extended to learning algorithms that are not explicitly derived from that principle, such as three-factor Hebbian learning FFrmmaux & Gerstner (2016); Kušmierz et al. (2017); Pogodiin & Latham (2020). For a single layer, three-factor updates between y and follow e y x for some error signal e. As long as dynamics leads to a solution, the analysis should be similar to the general mirror descent
since weights will span (as required by the theory). For multi-layered networks, all we would need to assume is that the input to each layer does not change too much over time, similar to the linearization argument made for deep networks.
More broadly, our work opens new experimental avenues for understanding synaptic plasticity in the brain. Our approach isolates synaptic geometry from other components of learning, such as losses, error signals, error pathways, etc. Therefore, our findings make it practical to experimentally determine the brain's 's synaptic geometry.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported by the following sources. BAR: NSERC (Discovery Grant RGPIN-2020-
05105, RGPIN-201 8-04821; Discovery Accelerator Supplement: RGPAS-2020-00031 Arthur B.
McDonald Fellowship: 566355-2022), Healthy Brains, Healthy Lives (New Investigator Award:
2b-NISU-8), and CIFAR (Canada AI Chair; Learning Machine and Brains Fellowship). GL:
NSERC (Discovery Grant
Canada- CIFAR AI Chair, Canada Research Chair
in Neural Computations and Interfacing, and Samsung Electromics Co., Ldt. GG: Canada CIFAR Al Chair. JC: IVAOO Postdoctoral Fellowship and the Canada First Research Excellence Fund. AG:
Vanier Canada Graduate scholarship.
This research was enabled in part by support provided by (Calcul Québec) (https:/ www. calculquebec. wwweebecccaee/////// and Compute Canada (www womputecanadeeeoooccoocccoc ca). The authors acknowledge the material support of NVIDIA in the form of computational resources.
REFERENCES
Mohamed Akrout, Collin Wilson, Peter Humphreys, Timothy Lillicrap. and Douglas B T'Teed. Deep learning without weight transport. Advances in neural information processing systems, 32, 2019.
Shun-ichi Amari. cnfferentialgeemmeerree methods in statistics. Lecture Notes on Statistics, 1985.
Shun-Ichi Amari. Natural gradient works efficiently learning. Neural computation, 1998.
Navid Azizan and Babak Hassibi. Stochastic gradient/mmrror descent: Minimax optimality and implicit regularization. arXiv preprint arXiv: 0005200000000.....0 2018.
Navid Azizan, Sahin Lale, and Babak Hassibi. Stochastic mirror descent on overparameterized
nonlinear models. IEEE Transactions on Neural Networks and Learning Systems, 2021.
Shahab Bakhtiari, Patrick Mineault, Timothy Lillicrap. Christopher Pack, and Blake Richards. The functional specialization of visual cortex emerges from training parallel pathways with self- supervised predictive learning. Advances Neural Information Processing Systems, 34:25164- 25178, 2021.
Amir Beck and Marc Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Reseurch Letters, 31(3):67-175, 2003.
Patrick Billingsley. Convergence of probability measures. John Wiley Sons, 1999.
JR Blum, Chernoff, Rosenblatt, and H Teicher. Central limit theorems for interchangeable processes. Canadian Journal of Mathematics, 00:2222-29, 1958.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. arXiv preprint arXiv:2205.0966 09655, 2022.
Lev Bregman. The relaxation method of finding the common point convex sets and its applica- tion to the solution of problems in convex programming. USSR computational mathematics and mathematical physics, 7(3):200-217, 1967.
Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Trends@ in Machine Learning, 8(3-4):211-357,001
Foundations and
György Buzsáki and Kenji Mizuseki. The log-dynamic brain: how skewed distributions affect network operations. Nature Reviews Neuroscience, 15(4):264 278, 2014.
Simone Carlo Surace, Jean-Pascal Pfister, Wulfram Gerstner, and Johanni Brea. On the choice of metric in gradient-bart based theories of brain function. arXiv e-prints, pp. 18Xiv-
2015, 2018.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the JEEE/CVF conference on computer vision and pattern recognition, pp. 15750-15758, 2021.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
10

---

Published as a conference paper at ICLR 2024
David Clark, LF Abbott, and Yue Chung. Credit assignment through broadcasting global error vector. Advances in Neural Information Processing Systems, 14:10053-10061 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255, leee, 2009.
Sven Dorkenwald, Nicholas L Turner, Thomas Macrina, Kisuk Lec, Ran Lu, Jingpeng Wu, Agnes Bodor, Adam A Bleckert, Derrick Brittain, Nico Kemnitz, al. Binary and analog variation of synapses between cortical pyramidal neurons. Elife, 11:e76120, 2022.
John Duchi, Shai Shalev-Shwaed Shwartz, Yoram Singer, and Ambuj Tewari. Composite objective mirror descent. In COLT, volume 10, pp. 14-26. Citeseer, 2010.
Timo Flesch, Keno Juechems, Tsvetomira Dumbalska, Andrew Saxe, and Christopher Summerfield. Orthogonal representations for robust depeext-eepnneent task performance brains and neural networks. Neuron, 10(7):1258-12700 2022.
Nicolas Frémaux and Wulfram Gerstner. Neuromodulated bikke--iming-deeeddeee plasticity, and theory of three-factor learning rules. Frontiers in neural circuits, 9:85, 2016.
Udaya Ghai, Elad Hazan, and Yoram Singer. Exponentiated gradient meets gradient descent. In Algorithmic learning theory, pp. 386-407. PMLR, 2020.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Sre- bro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pp. 1832- 1841. PMLR, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Jyrki Kivinen and Manfred K Warmuth. Exponentiated gradient versus gradient descent for linear predictors. information and computation, 132(1):1-63, 1997.
Jonas Kubilius, Martin Schrimpf, Ha Hong. Najib J. Majaj, Rishi Rajalingham, Elias B. Issa, Kohitij Kar, Pouya Bashivan, Jonathan Prescott-Royy Kailyn Schmidt, Aran Nayebi, Daniel Bear, Daniel L. K. Yamins, and James DiCarlo. Brain-Like Object Recognition with High-Performing Shal- low Recurrent ANNs. In Neurul Information Processing Systems (NewrIPS). Curran Associates, Inc., 2019.
Lukasz Kuśmierz, Takuya Isomura, and Taro Toyoizumi. Learning with three factors: modulating hebbian plasticity with errors. Current opinion in neurobiology, 46: 160-777 2017.
Tor Lattimore and Andras Gyorgy. Mirror descent and the information ratio. In Conference on Learning Theory, pp. 2965-2992. PMLR, 2021.
Tor Lattimore and Csaba Szepesvári. Bandit algorithms. Cambridge University Press, 2020.
Guillaume Leclere, Andrew Ilyas, Logan Engstrom, Sung Min Park, Hadi Salman, and Aleksander Madry. FFCV: Accelerating training by removing data bottlenecks. htttp:://githti com/
fibffcv/ffcw/ 2022. commit 720cc2e.
Yann LeCun, Corinna Cortes, Chris Burges, et al. Mnist handwritten digit database, 2010.
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.

---

Published as a conference paper at ICLR 2024
Robert B Levy and Alex D Reyes. Spatial profile of excitatory and inhibitory synaptic connectivity in mouse primary auditory cortex. Journal of Neuroscience, 32(16):5609-556 5619, 2012.
Qianli Liao, Joel Leibo, and Tomaso Poggio. How important is weight symmetry in backpropaga- tion? In Proceedings of the AAAI Conference on Artificial Intelligence, volume 30, 2016.
Timothy P Lillicrap. Daniel Cownden, Douglas B Tweed, and Colin J Akerman. Random synaptic feedback weights support error backpropagation for deep learning. Nature communications, 7(1): 13276, 2016.
Yonatan Loewenstein, Annerose Kuras, and Simon Rumpel. Multiplicative dynamics underlie the emergence of the log-normal distribution of spine sizes in the neocortex in vivo. Journal of Neuroscience, 3 31(26)99481-944 2011.
Ningning Ma, Xiangyu Zhang, Hai-Tao Zheng, and Jian Sun. Shufflenet v2: Practical guidelines for efficient architecture design. In Proceedings of the European conference computer vision (ECCV), pp. 116-131, 2018.
James Martens. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research, 21(1):5776-585 5851, 2020.
Alexander G de Matthews, Mark Rowland, Jiri Hron, Richard Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. arXiv preprint arXiv: (0ti:/880..1222 11271, 2018.
Joshua B Melander, Aran Nayebi, Bart C Jongbloets, Dale A Fortin, Maozhen Qin, Surya Ganguli, Tianyi Mao, and Haining Zhong. Distinet in vivo dynamics of excitatory synapses onto cortical pyramidal neurons and parvalbumin-positive interneurons. Cell reports, 37(6):109972, 2021.
Mohamad Amin Mohamadi, Wonho Bac, and Danica J Sutherland. Making look-ahead active learn- ing strategies feasible with neural tangent kernels. arXiv preprint arXiv:2206. 12569, 2022.
Aran Nayebi, Daniel Bear, Jonas Kubilius, Kohitij Kar, Surya Ganguli, David Sussillo, James DiCarlo, and Daniel L. Yamins. Task-driven convolutional recurrent models of the visual system. Advances in neural information processing systems, 31, 2018
Arkadij Semenovič Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. Wiley-Interscience, 1983.
Yann Ollivier, Ludovic Arnold, Anne Auger, and Nikolaus Hansen. Information- op-
timization algorithms: A unifying picture via invariance principles. The Journal of Machine Learning Research, 2017.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances neural information processing systems, 32, 2019.
Bill Podlaski and Christian K Machens. Biological credit assignment through dynamic inversion of feedforward networks. Advances in Neural Information Processing Systems, 13:10065-10076, 2020.
Roman Pogodin and Peter Latham. Kernelized information bottleneck leads to biologically plausible 3-factor hebbian learning in deep networks. Advances in Neural Information Processing Systems, 33:7296-7307, 2020.
Yi Ren, Shangmin Guo, Wonho Bae, and Danica J Sutherland. How to prepare your task head for finetuning. arXiv preprint 0rXiv:2302.05779, 2023.
Blake A Richards, Timothy P Lillicrap. Philippe Beaudoin, Yoshua Bengio, Rafal Bogacz, Amelia Christensen, Claudia Clopath, Rui Ponte Costa, Archy de Berker, Surya Ganguli, et al. A deep learning framework for neuroscience. Nature neuroscience, 22(11):1611-177 2019.
Blake Aaron Richards and Konrad Paul Kording. The study of plasticity has always been about gradients. The Journal of Physiology, 2023.
12

---

Published as a conference paper at ICLR 2024
Daniel L Ruderman. The statistics of natural images. Network: computation in neural systems, (4):517, 1994.
Martin Schrimpf, Jonas Kubilius, Ha Hong, Najib J Majaj, Rishi Rajalingham, Elias B Issa, Ko- hitij Kar, Pouya Bashivan, Jonathan Prescott-Roy, Franziska Geiger, al. Brain-score: Which artificial neural network for object recognition is most brain-like? BioRxiv, pp. 407007, 2018.
David PA Schulz, Maneesh Sahani, and Matteo Carandini. Five key factors determining pairwise correlations in visual cortex. Journal of neurophysiology. 114(2):1022-14) 2015.
Jonathan Schwarz, Siddhant Jayakumar, Razvan Pascanu, Peter E Latham, and Yee Tch. Power- propagation: A sparsity inducing weight reparameterisation. Advances Neural Information Processing Systems, 34:28889-28903, 2021.
Shai Shalev-Shwarzz et al. Online learning and online convex optimization. Foundations and Trends@ in Machine Learning, 4(2):107-194, 2012.
Sen Song, Per Jesper Sjöström, Markus Reigl, Sacha Nelson, and Dmitri B Chklovskii. Highly nonrandom features of synaptic connectivity in local cortical circuits. PLoS biology, 3(3):668, 2005.
Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural net- works. In International conference machine learning, pp. 6105-6114, PMLR. 2019.
Dimitra Vardalaki, Kwanghun Chung, and Mark T Harnett. Filopodia are structural substrate for silent synapses in adult neocortex. Nature, pp. 1-5, 2022.
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
Daniel LK Yamins, Ha Hong, Charles F Cadicu, Ethan Solomon, Darren Seibert, and James DiCarlo. Performance-optimizes hierarchical models predict neural responses in higher visual cortex. Proceedings of the National Academy of Sciences, (21))))8619 8624,86244 2014.
Chiyuan Zhang, Samy Bengio, Moritz Hardt. Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
Weishun Zhong, Ben Sorscher, Daniel D Lee, and Haim Sompolinsky. A theory of learning with constrained weight-distribution. arXiv preprint urXiv:2206.0em 08933, 2022.
13

---

Published as a conference paper at ICLR 2024
APPENDICES
A WEIGHT DISTRIBUTION CONVERGENCE
We start by giving intuitive outline of our proof of Gaussian convergence of weight changes. As for any Gaussian convergence, we'll end up using the central limit theorem (although an unusual version of it) w.r.t. the dataset size. With this in mind, the steps are:
Setup First, we decide how we take the large data and network limits. We have a vector of weight changes whose size grows to infinity. To discuss what converges to, we consider this vector to be a part of a stochastic process (so, an already anfinit-diimesionn vector). We then consider the convergence of this stochastic process some limiting process (in our case, to a Gaussian process). This allows us to approximate a sample of weight changes as a sample from that limiting process.
Lemma
Next, we show that the inverse matrix in Eq. (5) used in the weight change expression behaves like an identity matrix. we just took the limit w.r.t the network size, this lemma would be trivial since off-diagonal elements are products of independent vectors. But we also increase the dataset size (and hence the matrix itself), so we need to be more careful. This where we get (most of) the requirements for how the dataset size should scale with network size.
Lemma 2
shows that the deviations from identity resulting from Lemma don 't affect the final weight change expression when we take the limits.
Lemma 3 shows that if you replace that inverse matrix with an identity, you can apply the exchange- able central limit theorem to the weight change expression. At this point, it's applied for any finite slice of the stochastic process we defined. The exchangeable central limit theo- rem has much more restrictive moment conditions than the standard CLT, since it doesn't t require independence,
have carefully check all of them.
Theorem Finally, we combine the lemmas and show that any finite slice of that stochastic process converges to Gaussian random variable. Due to the properties of Gaussian processes, we can conclude that the whole process converges to a Gaussian process.
Setup We want to show that the weight change looks Gaussian in the large width limit. Since the number of weights goes infinity, "looks Gaussian" should translate to convergence to a Gaussian process. Our approach will be similar to Matthews et al. (2018), which showed convergence of wide networks to Gaussian processes. However, they worked with dinite-diensionall inputs and multiple hidden layers; for our purposes we'll need to have an infinit-dimmnssion input and a single layer.
We'll be working with countable stochastic processes defined over Σ>o- We need three stochastic processes: the input data X, the initial weights W and the "teacher" weights W
For width D, the label yo and the model at initialization Th are defined as
D J Σ d=1
yn Σ
Xa Wa.
Xa WT.
(17)
For N sample points of dimension D, we stack them into an N x D matrix X and denote the D-dimensional samples of weights as DD DD The weight change we're re interested in then
_D 1
(X))
( P (X))T) x wDD)
(18)
VD
-XD) T (x D HD (XD)T)- (yD y2))
(19)
Note the explicit 1/VD scaling we omitted in the main text for convenience and N-dimensional /DD/ (we dropped
vectors of labels yD, yb. Here H is a diagonal matrix with had = the o- subscript used in the main text for readability).
Invertibility assumption. Throughout the proofs we assume that X H (XP)T is almost surely invertible. This is a very mild assumption, although it can break if X sampled from discrete
14

---

Published as a conference paper at ICLR 2024
distribution, if Xa have strong correlations along d or if the initial weights not almost surely non-zero.
<P is D-dimensional vector, but we understand it subset of the following stochastic process:
E (x) IX"HO (X*)T)= (b)-
(20)
This inftnite-width expansion similar to the one used in Matthews et al. 2018 It's a well- defined stochastic process since N is finite, meaning that E is weighted sum of N stochas- tic processes. However, it formalizes convergence of imite-dimensional vector & to something infinite-dimensional. That is, we'll work with convergence of the process =D limit process E in distribution
First, we show that an appropriately scaled matrix inside the inverse behaves like an identity:
Lemma 1. Assume that the input data points Xa have mean and unit variance. Additionally assume that uniformly for any d, there exists a constant e such that
20 Σ Cov (X2. x3 cc
Σ (EXx sc,
(21)
Also for i.i.d. initial weights, define moments (assuming they exist) of had = (WO//// as
h2(D) = Var -
(22)
Then for N, D large enough (so the r.h.s. is smaller than 4).
(DM(0) X2H
( (HPP))) - In
N4 c
h2 1
8 8
(23)
D
h
2
with probability at least 1 p.
Proof. Denote
1 X (XPT, Ag = Dh = Dh
Σ X X hadd
(24)
dropping the explicit D-dependence in h1 and h2 for convenience.
Since all Xa are independent from WO. E A.j = ōij. Variance can be bounded. For = j.
Var(Au)
22 Var Var((X)) haa) + D2 h2
Σ Cov((X)) had, (X2) haw)
(25)
d.d'
Σ h2 + (h, + h2) Var((X))
h Cov ((X)), (x))))
Σ
(26)
h2
D2 D2
d
h2(1 + Var((X
c
C
l12
lz
c+1 1
( h2
1
+
1
(27)
D
D
h
Dh
D
h
where in the last line we used the second part of Eq.
21
For + j.
D Σ
Var(Ag)
Σ Var (X haa) + D2
Cov(X' had, X'x X3 hao) (28)
12
d,d'
(0 + h2) (Var(X)
(Cov(X di x2
h2
(29)
D2 h2
D2
D
h
d.d
'See Eq. 10 in Matthews et al. processes.
2018 for the discussion on what "in distribution" means for stochastic
15

---

Published as a conference paper at ICLR 2024
now using the first part of Eq. (21 for the last inequality.
Using Chebyshev's inequality and the union bound, we can bound
max AAs - desl 2 ) $ 6 D
N2 c c+1
h2 l1
P
1+
=
p.
(30)
Therefore,
N c D
hg h
A I & ||A Ivll
1+
(31)
with probability at least
By Lemma 4.1.5 of Vershynin (2018), if ||A Ivll2 N then all singular values o, of A lie between (1 e N and (1 + e eN for N large enough so e < 1. Therefore, we also have
bound on the singular values of A Additionally taking parameters big enough for N I 0.5, have
IN max -1) S
N 222 (1 - N (1 N) S SEN,,
1
(32)
which completes the proof.
Now need to show that the inverse behaves like an identity in the overall expression well:
Lemma 2. the assumptions of Lemma |
additionally assume that W and W have fi- nite mean and variance independent of and D. that hz(D)/h, (D)² is independent of D. and Ed' E X X S c. uniformly over d. Then for any sealar weights os and a finite index set
A C Z as D, N + 00 with N = ((D2/(6+)) for any 8 > O,
Σ a p)) (A-+- -Iv) (YD - yb) 7p 0. PEA
(33)
Remark. For both cross-entropy and p-norm potentials, V is a power of W/VD, and so h2(D)/ (D)) depends only on the moments of W, but not on D, therefore satisfying the assumption.
Proof. Again denoting the matrix inside the inverse as A (see Eq. via Eauchy-Schwarz
24
we first bound the following
2 S
Σ ai (X (A-1- (yD (YD vb)
Σ ai (X
IyD ybll |A-1- -8--INN IN
(34)
First term. Using
and that Xp variances are one, and also that a,, A are fixed,
Σ a X p)
N E
Σ Xp PEA
SNE
Σ a
Σ
(x!) = O(N). (35)
PEA
A
PEA
Second term. Using that weights and inputs are independent, and that the sum of input data covari- ances bounded by e by assumption,
16

---

Published as a conference paper at ICLR 2024
D
2 (W2 (Wa- w.)) )²
N
E YD yB112
E
(36)
D
d=1
N D
Σ (W2 W2)
=
(37)
D
d,d'zd
N
+
D (R - -2))² Var((W] - W2))) R (x))²
(38)
S (B - 2 + N
- -
= - O(N).
(39)
Third term. Here we will use the probabilistic bound from Lemma
Overall, for any 8 > 0,
2 (( 11-1-11 -Iw) (yn - -²/) 26)
1
P
(40)
N
IEA
< IP
Las ai (X::) iEd
Σ
> >N ||y - -12112 > NI |A- > NI+5
(41)
E EpEE Op (X p)
E IlyD y-bll
P
(42)
N1+
N2+
NI+8
00
1
+ 64
No+ō
c+1
(1+ h2 h?
(43)
N8/2
D
where we used Eq. 34 in the first line, the union bound with Markov's 's inequality in the second and the bounds on all terms in the third one (for N large enough so e/N < 2).
As long as 5 = DDD and h2/ is constant, the probability goes to zero with N, D.
The previous result suggests we can replace the inverse matrix in = (Eq. 20 with an identity. This will require an additional step, but now we can show that with the identity, for any finite index set A of points in =o and arbitrary weights a, the A, a-projection of =" solution converges to a Gaussian.
Lemma 3. Consider a finite index set A and an associated vector of weights a. If in addition to condition of Lemmas and 2 we assume weights have funite 8th moment, X2 are uniformly bounded, and that Xy ybb (for point i and index p) converges in probability random variable. Then the A, a-projection defined as,
N
Σ 1=1
RND
RN = -
- -b),,
(44)
P.A.
converges to NO, oy a AAa) in distribution, where (EA) = E =E and o = E(W w as long D increases monotonically with N such that N/D
Proof. We assumed that D is monotonically increasing function of N. Therefore, get a trian- gular array of Ry; with infinitely exchangeable rows. That is, for each N we have a process RN: with i = 1, ... such that any permutation of indices doesn't 't change the joint distribution. This is due to joint weight dependence among labels. Sums of infinitely exchangeable sequences behave similarly to sums of independent sequences when it comes to limiting distributions. In particular,
we will use the classic central limit theorem result of Blum et al. (9588..
To apply it, we need compute several moments of RNG-
17

---

Published as a conference paper at ICLR 2024
First moment. Since B(W; W2) 0, we have E RN, O.
Second moment, Again using E (Wa - W = and the fact that all weights are i..d., so E ( W ) = Öij, get
R E ER
1
E ap ap E xxxxxxxxxxxx (W3 - w.)
(45)
D
pp
1
Σ
ap Op
E X X
(W3 w2)²
(46)
D
pvdd
= This term is O(1) due to the 1/D scaling. Denoting ow = E (Wa w2) and App
D
Ea E X (X2)² / /D
(and the corresponding matrix r),, we can re-write the variance
x)) ow a T
(47)
Absolute 3rd and 4th moments. For convenience, we first find the 4th moment. Since the inputs bounded and independent from the weights, and dropping the summation over all pt and d, for readability,
I [Iw w -
E R
D2
Σ
II
Oppp
E
II
X p:
X
E
0(1)
(48)
This follows from E (II X2. X
(1) due to bounded inputs, and due centered and i.i.d.
weights E (ITL,, w if any d, differs from the rest (i.e. only O(D2) out of D non-zero; these are the ones with d1 = d2 and d3 = d, and permutations of those conditions).
For the absolute 3rd moment, we can use Cauchy-Schwarz to find its scaling:
E |RN | |Ryy
E RN 2 PE O(1),
(49)
since both terms in the square root O()). Covariance. For two different points i and j. defining the covariance opd = E. E.
E RN,
1
Σ
ap E X X (Wi
W -W2)
(50)
D
dd'
1
ap
E X Xix - w29
oi
Σ
apap Opdowa. (51)
D
D
pp'd
pp'd
Since we required Es pd S c, Ea Padoy S La Thus, Ry,RNN O(1DD) o(1/N).
oz S a by Cauchh-Schwarc
Covariance of Ry,,. Again, take + and adapt the calculation of the 4th moment:
( x
E ER R
ape II A=1
E
II A=1
x
X
I (WI, - w.2,)
D2
A==
(52)
Due to i.i.d. weights the whole sum has only O(D2) terms and therefore the whole expression is 0(1). However, require more precise control over the sum.
The summands split up into three cases (with a potential O(D) overlap): (1) d1 = d2 = dg = d,, (2) di dy, de d4 dj dg, d2 dg. and (3) d1 d2, ds d4. The first case has only D terms, each (11), in the sum over dg. The second case is a bit more tricky. Assuming dj = dg w.l.o.g., by Cauchy-Schwand and by the assumptions correlations (note that the weights are i.i.d.,
18

---

Published as a conference paper at ICLR 2024
so we can take the expectation w.r.t. the weights out of the sum),
Σ (X
X
X
X
(x
x
X
X
E ((W2, w2, P - -2
(53)
pa
d1.d2
Σ
(X
X
p2)
(X
X ps
E (Xa, x '' ((W2, - W )(
w22))
pi
d1,d2
(54)
Sc B /EDVE B 2' ((W) wP) W22)2)
(55)
Since other terms in the full sum are (11), the contribution of the second case is O(D).
The last case is what we had for covariance. Since d1 = d2 and dg = d4.
Σ E
X P2 ) (X) X (Xa,) )2) E ((W - W2, (W W2))
(56)
dg.dg
D(D 11D-1y>pp
1y
Trap σw + D Tara E (Wi - w
D D
(57)
Now we can combine all three terms, using that the weights have finite second and 4th moments.
We will also use that for d, d2 and d3 = d4, we can ignore the second term with E (WT - wP) since it's an O(D) contribution to an O(D²) sum. Therefore,
1
E R3g =
1 D22
II
pp
D(D - 1D)D-1))/pp.... Pppps + O(D))
(58)
Plk
= - a p 2 +0 (1) - 0
(59)
D
Asymptotic of the second moment. To find the limit of the variances, we need to compute
E R
Σ apay ppdd
w2)
(60)
This variable is uniformly integrable since its 4th moment exists (almost identical computation to R%y, and is uniformly bounded for all N (using Billingsley (1999). immediately after Theorem 3.5
for 2). Therefore, limy E. R2y = E limy R2, (Theorem 3.5 in Billingsley (1999)). The latter is easy to compute, since
E (WG - -Wy Dow.
(61)
Var E
(X5) (W2 -wge)
Eoo ((X)) w (X)) (W3 WD)2) (62)
d
d.d'
=O(D) Σ Cov ((X)) wd),, (xz) 22)
(63)
d.d'
= O(D) + oi
Cov ((x)), (x² O(D),
(64)
d, td
where in the last line we again used the assumption on summable covariances of (X))².
Therefore, by Chebyshev's 's inequality to (W))))( - WDy (note the returned D/D) converges to ON in probability. Therefore, by Slutsky's theorem Ryy converges to oi Eyp apap' X X Since the expectation of X X is ppp,, , we can denote ΣA as a matrix with entries 0² on,, such that
R2 = oi a T - of a TT AAA.
(65)
Since we assumed that Rxs converges to random variable, and R2, is uniformly integrable (since its 2nd moment exists), we can claim that E R2 oby ΣA E R.a
19

---

Published as a conference paper at ICLR 2024
For Rig. repeat the calculation. It's also uniformly integrable since its 8th moment exists: E II-=1 (W2, - W leaves out D terms out of Ds, the weights have a finite 8th moment, and the inputs are bounded. Therefore,
E Ri, Pi - ow [aT Exa)²
(66)
CLT for exchangeable arrays. Assume that (D) O(1) (defined in Eq. 47 has non-zero limit, and D is large enough for (D) > O. Then GNv: = RNv/on(D)) is zero mean, unit variance
and has finite third moment. It also satisfies the following:
1
1
E GNNIGN2 = 0
(67)
D
N
lim E G G 1, N
(68)
E E|GNi 0(1) ol o(vN).
(69)
Therefore, Gy: satisfies all conditions of Theorem in Blum et al. (1958), and
GNI a 2 (0,
(70)
N
Therefore, also have
N
RN
- d N( T EAa)
(71)
N
(Note that in the case a' T A 0, we don't need apply CLT since the sum converges zero in probability by Chebyshev's 's inequality.)
Now we're ready to prove the main result (see the more informal statement in the main text, Theo- rem D
Theorem 1. In the assumptions of Lemmas to the scaled sequence of stochastic processes
=" = (X) H" (X) (yD - yb)) converges to a Gaussian process:
ED DEV (00/ ED +a GP(O, E),
(72)
VN
where the covariance matrix over Xa and ow E (Wa - W
Remark. Lemma I used the following correlation assumptions: D < c and Ed= Cov (X3, x < c. Lemma 2 also used Ed= E X. Xar S c. In terms of Theorem in main text, the second condition corresponds to dy.- The first and third conditions are joined into the Cig condition for convenience since it implies both conditions used here.
Proof. First, for an A, a-projection of the scaled stochastic process, we split the solution into two terms (A is defined in Eq. (4
Σ PEA
ap
ED
Σ (X
p)
A lyD ybb)
(73)
VN
PEA
o x (yD yb)
-yn-yl Σ Cap T (A- - yb). (74) PEA
VN
VN
PEA
By Lemma the second term converges to zero probability. By Slutsky's theorem, since the first term converges a Gaussian in distribution, and the second one to zero in probability, the whole expression converges to Gaussian.
We proved convergence of the first term to a Gaussian in Lemma 3
Convergence a Gaussian process. We've ve shown that for any scalar weights a and a finite index set A C E>o, the projection of the sequence of the scaled stochastic processes ED w.r.t A, a converges to a Gaussian in distribution (with parameters found in Lemma 3). Therefore, ED converges Gaussian process in distribution by Lemma 6 of Matthews al, (0018).
n
20

---

Published as a conference paper at ICLR 2024
The teacher weights setup with yD Here we provide similar result for a generic y without detailed proof:
Ed=1 X Wa Wa simplifies the proofs, but it is not necessary.
Theorem 2. In the assumptions of Lemmas mean, bounded, and independent of D, X and W and also that W are zero-mean. Then the
additionally assume that labels y² zero-
scaled sequence of stochastic processes ED (X)
H TX (y (y-b)) converges to
Gaussian process:
DEV
d GP(O, (o + ) ED
(75)
VN
where Σ is the covariance matrix over X and weo E (W2) o Ey²
Proof. The label setup doesn't affect Lemma For Lemmas and 3 due to the independence assumption for the labels and the weights, the proofs of the lemmas split into moment calculations for W and y. The former is identical to the original calculation for W W since the weights are zero-mean. The latter doesn't involve sums over repeating coordinates, simplifying moment computations. The overlap similar the computation for W With this change in the lemmas, the rest of the proofs is the same as for Theorem
A.1 RICH AND LAZY REGIMES
Assuming for simplicity that all initial weights are the same and positive (w = w a/ (e/DD) and that of is equal to 1, then for & ~ (O,, 1), Eq. 13 becomes
VN
Volw ~ Vo(w")
2-2-I1(00) 5.
(76)
For p-norms, <(w) = ; Vo(w = =w -20-1(00 = (20)--2- Therefore,
a=11 (p - 1)VNa-2 a=- DP/2
Vo(w>) D(p-1)/2
(aVD - (p 1) -w-11) 1)
(77)
DP/2
As long as a > D/DD, the initial weights (first term) will be much larger than the weight change (second term), and we will be in the lazy regime. This additionally sets limit on the dataset size. Put another way, for of data points N has to be much smaller than D for the linearization to make sense (in addition to the more technical assumption on the N scaling in Theorem D
1 and the standard "lazy" initialization with O / I) weights, the number
For negative entropy, ((w) = w log Vo(w>) = 1 + log w 2 (w=0) wº. Therefore,
1
NN
V ( x 1+ log a 2 log D+
aVD* E.
(78)
Unlike for p-norms, the initial weights' contribution grows infinity even tends to zero. As result, we require only > VN/(D log D) keep learning in the lazy regime. However, for 1/VD (i.e. O(1/D) weights, the standard rich regime scaling) negative entropy still results in large weight changes since N grows with D.
Overall, different potentials result in different asymptotic boundaries of the lazy regime. Negative entropy allows smaller initialization schemes while maintaining the lazy regime.
B
EXPERIMENTAL DETAILS
B. LINEAR REGRESSION
For experiments in Section we optimized models using mirror descent, where the gradients were computed in the standard SGD way with momentum of 0.9 and no weight decay (with mixed precision). Learning rate was changed during learning according to the cosine annealing schedule. The initial learning rate was chosen single seed via grid search 16 points (OogaO-space
21

---

Published as a conference paper at ICLR 2024
negative entropy
AGDF
Adlw
AGDF
AdWw
0=2
0=1
00
06
2-norm
Ad(wl
AGDF
Adjol
0=2
0=2
C=1
C=1
05
00
08
0G
3-norm
ADDF
1/0 Ad(e!
AGDF
1.0 Adfe/
2
c-2
0=1
C=1
611
05
0.1
0.5
6.0
10
127
0.0
10F
13"
10"
06
707
10P
10
0.0
107
10
10"
number of weights
rumber of wnights
number weights
number of weights
Figure 6: Linear regression solutions for negative entropy, 2-norm and 3-norm potentials. First and third columns: integral of absolute CDF difference (ACDF) between normalized uncentered weights and (0, 1); second and third columns: magnitude of weight changes relative to the initial weights in the dual space (Aø). First two columns: quickly decaying correlation between 2; (c 2); last two columns: slowly decaying correlation and 5/95% percentiles (shaded area), N = D0.5 (orange) and N D0.75 (blue).
1). All plots: 30 seeds, median value (solid line)
from 1e-7 to le-1) and 100 epochs. the minimum MSE loss was larger than 1e-3, we increased the number of epochs by 100 and repeated the grid search (up to 500 epochs; all grid searches converged to MSE smaller than 1e-3). The widths were chosen as 10 points in a log log10-space from 1e2 to 2c4. The full results are shown in Fig. 6
B.2 POTENTIAL ROBUSTNESS
In Section 4.2 both figures in Fig. 33 were obtained in the following way. For D 100 and 30 seeds, we sampled a 1000-dimensional weight sample w" from either N (0, 1/D) or a log-normal with log v 22D, o² = log 2 multiplied by a +1 with equal probability (so it has the ap- proximately the same variance the Gaussian). We then computed the final weights w for the potential o as w - (V((w") + 0.1 (6(L/VD)
the Gaussian change had smaller magnitude than a weight with standard deviation of 1 1/VD.
B.3 FINETUNING
In Section 4.3 for IV chosen datapoints, the dataset was randomly sampled from the 50k available points for each seed (out of 10). Networks were trained on the cross-entropy loss using stochastic mirror descent; the gradients had momentum of 0.9 but no weight decay, batch size 256. Learning rate was changed during learning according to the cosine annealing schedule. The initial learning rate was chosen on log10 grid of 5 points from le-5 to le-1. The initial number of epochs was 30, but increased by 30 up to times if the accuracy was less than 100%. All resulting accuracies from this search were 2 99%. (Note that this is the desired behavior, since we're testing the models that fully learned the train set.) The dataset was not augmented; all images were center cropped to have a resolution of 224 pixels and normalized by the standard ImageNet mean/variance. We trained networks with mixed precision and FFCV Leclerc al. (2022) for the latter had convert ImageNet to the FFCV format (see the accompanying code).
The magnitude of weight changes was usually smaller than 0.2, although for 2-norm for ResNet18/34 8/34 and CORnet-S, and 3-norm for EfficientNets 10, it reached 0.3-0.5. For negative entropy, the magnitude always stayed smaller than 0.02.
22

---

Published as a conference paper at ICLR 2024
A 4CDF
RNN. 100 hidden units
ACDF
RNN, 1000 hidden units
B ACDF
SimSiam ResNet50
0.7-
trained on: 2 ioom
entroyy NE
0.6
0.6
0.5
0.4
0.4
0.3
0.2-
0.2
0.1-
0.0
0.0
NE
1.5
2.5
NE
1.5
1.5
3
X-axis: evaluated on NE p-norm
Figure 7: Same setup as Fig- 43. A. Recurrent neural networks (RNNs) with 100/1000 hidden units trained and then finetuned on row-wise sequential MNIST. B. SimSiam ResNet50 finetuned with the same self-supervised loss.
We have conducted experiments with recurrent neural networks trained on row-wise sequential MNIST LeCun et al. (2010). The networks were trained on the train set, and then finetuned on a subset of the test set (same procedure as for deep networks) for N D0.6 (the number of weights scales quadratically with the hidden size, so N equals the number of hidden units). The Gaussian fits, and Gon-Gaussia weight changes measured with a wrong potential were remarkably sim- ilar to theoretical predictions and behavior of deep networks (Fig. 7 Note that for networks trained with p-norms but evaluated with negative entropy, the weight flips its sign during training immediately conclude that negative entropy not used for training. Because we'r re fine- tuning already trained networks, most weights don't flip signs. We plot the distribution of changes in the dual space for all weights for simplicity (we do it in the code since the dual weights signed).
We also evaluated a self-supervised loss from SimSiam Chen & He (202)): for a pre-trained ResNet50 (see the full architecture in Chen & He 2. we finetuned it with the same self- supervised loss / = 0.5 z P2 + 0.5 0.5 P1 for positive pairs 1/2 passed through decoders with or without stop-gradient (z and p). We trained the network for 50 epochs and used the learning rate that minimized the loss the best, since we don't evaluate accuracy. We evaluated weight changes in the encoder (a ResNet50). The results were very similar to networks finetuned with a supervised loss, confirming our results for supervised experiments (Fig. 7
B.4 COMPARISON TO SYNAPTIC DATA
we used the data from Dorkenwald (2022). We used the the accompanying dataset specifically the "spine_vol_um3"" data available here The dataset contained 1961 spine volume recordings.
In Section 4.4
For negative entropy fitting (Fig. AA). we sampled the initial weights as 10 with probability p 0.77 and 0.77 with p 0.23. In the dual space, we added centered Gaussian with 10 10-0.23 The means and probabilities exactly match the ones fitted in Dorkenwald et al. (2022) (Table 2); the standard deviations were fitted as 10-0.24 and 10 10-0.22 we used a log-average since the standard deviations were very similar and we expect the Gaussian parameters to be the same across weights. For 3-norm (Fig. 5). we obtained a qualitative fit. With equal probability, the initial weights were either le-2 or log-normal with μ -2.7 and 0.9. In the dual space, the Gaussian change was zero-mean with 0.0012. A small part of the final weights became negative with this sampling procedure, which we treated as synapses that were reduced to zero weight during learning and therefore excluded them from the plots.
The 3-norm approximately fits the data because for a small constant initial weight c. the final weight will be approximately log- lormal: log (Vo-I(Ve (Vo(c) + {)) z log (c + o()) = E/V C/V2)(c). The second part of the mixture is already log-normal and is barely changed by the Gaussian change, resulting in mixture of two (approximately) log-normals.
23

---

Published as a conference paper at ICLR 2024
C NOTES ON MIRROR DESCENT
For negative entropy, we assumed that the weights do not change signs. This way, |e| log [o| is strictly convex for either w w < 0 and the gradient of potential is invertible. This implied by the exponentiated gradient update Eq. 7 which indeed preserves the signs of the weights.
Powerpropagation Schwarz al. 202 such that gradient descent is done over H. The original weights are therefore updated for loss L as (assuming all terms remain positive for simplicity):
reparametrized network's weights w
0 00
w++ -(0- 0' - a Owf 0g
ØL Ow
ØL Ow' (0") - a (0 "I 2 w
(79)
00
Ô Ow"
w' a² (w") 2(a-1) 0
-0 a² )0t)2-)
(80)
d w
where we linearized the update in the first line and then used the relation between 0 and w. If treat (w*) 2(a-1) as the Hessian of o in mirror descent (see Eq. (4)), powerpropagattoon can be
viewed as an approximation mirror descent with p-norms, where p 2/a. For a 2, we can
also treat powerpropagation as an approximation to exponentiated gradient (Eq. 7
24


Reviews:
[{"review_id": "izzQy52qF0", "reviewer": "Reviewer_jXnh", "summary": "This paper develops a theory, based on mirror descent, of the distribution of synaptic weights changes. In particular, this distribution depends on the geometry of synaptic plasticity. They test theory predictions which are largely verified.", "questions": "See weaknesses", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I’m a big fan of the ambition in this paper. It’s original, the results are good, and this line of work could have big implications. That being said, I found the paper really lacking in clarity (see below).", "weaknesses": "The big assumption that the brain is doing gradient descent…\n\n“Notably, even if the brain does not estimate gradients directly, as long as synaptic weight updates are relatively small, then the brain’s learning algorithm must be non-orthogonal to some gradient in expectation” What’s the actual citation for this?? (i.e. not a review)\n\nI found the paper pretty dense, and not easy to follow what was going on where, and I always had to keep lots of things in mind at any moment. I’m not sure exactly what to suggest, but considerable rewriting / putting things in appendices / focussing on intuition would be a good idea.\n\nCouldn’t parse Fig 3. Needed more help in the caption / annotation..\n\n“Nevertheless, with this data we can rule out a Euclidean synaptic geometry, and if we do have access to w0, then our results show that it is indeed possible to experimentally estimate the potential function.” Where do you show this? It’s not in Fig 5.\n\nI didn’t follow the explanation before eqn 14. Could do with some clarifying.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C9Ga3FagmR", "reviewer": "Reviewer_DFTN", "summary": "The paper uses the mirror descent framework to derive results for the distribution of weights in converged networks. The paper shows that under a linearity assumption (and an L2 loss), the final distribution of weights is Gaussian, and this result can be used to analyse weight updates in trained deep nets and also synaptic updates in brain regions.", "questions": "1. How can we guarantee that the loss function / potential function separation truly represents a separation between extrinsic and intrinsic factors in the weight changes?\n\n2. How to evaluate whether the linearity / Gaussian assumption has been violated, in cases where the loss and potential function are unknown?\n\n3. How to make the \"geometry of synaptic weight changes\" more interpretable for neuroscientific insight?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "Strengths:\n1. The paper is, for the most part clearly written, and understandable. The mathematical sections in the main paper are also quite accessible and easy to follow.\n\n2. The idea for using the mirror descent to analyse synaptic weight changes is an elegant one, and the simplifying the analysis using the linearity assumption indeed seems to make the problem more tractable for use on networks applied to data.\n\n3. The experiments in section 4 are fairly thorough in validating the theory presented in section 3.", "weaknesses": "1. The distinction between \"loss function\" and \"potential\": The paper claims in section 2 that the mirror descent framework splits a synaptic weight update gradient into two terms: one dependent on \"external errors\" i.e. the loss function $l$, and one \"intrinsic to the synapse\" i.e. the potential $\\phi$. However, I am skeptical that this is generally true: to the best of my understanding, we can say that gradients with respect to the potential $\\phi$, $\\Delta\\phi$ are independent of gradients wrt the the loss $\\Delta l$. However, given that both the loss function and potential are functions of the synaptic weights, and that the loss function is either chosen by the practioner or unknown (in the case of the brain), unless we _explicitly_ choose $l$ to represent _only_ errors extrinsic to the synapse, it is not clear to me how we can guarantee the above statement. Relatedly, it is also not clear to me whether we can always pick a loss function that only captures non-synapse-related changes in weights, and if we can guarantee that we have managed this in every use case.\n\n2. Finally, the paper claims in section 1 to \"make experimental predictions\" and provide \"new theoretical insights...about learning algorithms in the brain\". In section 4.4, while it is clear from the results and the analysis that a non-Gaussian distribution of weights under the linearity assumption might indicate a non-Euclidean synaptic geometry, it is not clear how we can interpret this for insights about how learning happens in the brain, and for potential plasticity mechanisms. It is also not clear how extensible the analogy is from deep recurrent networks is to the brain -- synaptic weight changes may not be analogous to network weight updates via gradient descent, and therefore any theoretical results based on gradient descent may not port easily to neuroscientific insights.\n\nThe weaknesses taken together seem to make the paper fall short of the claims in the introduction. Without the interpretability, and further experimentation with neuroscience data, it is not clear to me what the framework adds in terms of neuroscientific insight. While it might be useful for analysing deep net behaviour, again, it is not clear how this can be used to improve deep net training either.\n\nMinor point:\nTerminology -- the terms \"loss function\", \"distance\" and \"geometry\" are used interchangeably without clarification, throughout the paper, and particularly in sections 1-2.1. It is not until section 2.1 that the distinction between the three terms and what they indicate becomes somewhat clear. This hinders readability and generally makes it very hard to grasp the premise of the paper or its implications without these terms being clearly explained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "X0fR2LMHrx", "reviewer": "Reviewer_GuwS", "summary": "This paper provides a theoretical framework for understanding the geometry of synaptic changes based on the synaptic weight distribution. In particular, they show that the synaptic weight distribution in real brains are inconsistent with vanilla gradient descent, which assumes Euclidean distances in weight space. Based on the Mirror descent framework, they derive a theorem that shows that the distribution of weights after learning depends on the initial distribution plus a Gaussian in the dual space for large models and data. The authors then show that the theorem holds in a linear regression setting, and that hypothesizing the wrong potential yields non Gaussian weight distribution. Then, the authors show that their theorem holds well despite unsatisfied assumptions in the case of classic deep learning convnets fine-tuned on ImageNet data. Finally, they consider biological data and conclude that both the distribution before and after learning are needed to conclude about the geometry of synaptic changes, but that vanilla SGD can be ruled out.", "questions": "- Section 4.4: I find the choices of the initial $w_0$ surprising to demonstrate the fits, isn't $w_0$ supposed to be the distribution before learning? If yes then it is surprising that it needs to be a mix of two constants since it is likely not the case in brains. Could the authors elaborate on that?\n\n- Can you elaborate on the link between geometry of plasticity and the locality of the learning rule? I would expect non Euclidean geometry to be non local since the metric tensor needs to be inverted, however Eq 7 does look local. \n\n- Isn't the use of optimizers such as Adam a way to have a different geometry of plasticity in ANNs?\n\n\nMinor: \n\nFig 3B: caption is not coherent: top/bottom should be left/right.\n\nFig 4: I see the histograms in pink and not blue as written in the caption.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The strengths of the paper are:\n\n**Originality:** The approach is original and principled, and it can make verifiable predictions about the biology. It is nice that the theory is robust to changes in the potential.\n\n**Quality:** The quality of the text and figures is high.\n\n**Clarity:** The paper stays very clear despite the theorem being a tad math-heavy.\n\n**Significance:** I think tackling the question of what geometry is followed by synapses in the brain is significant. It is true that most computational neuroscience work assume vanilla SGD so the question whether the resulting weight distributions are coherent with biological data is an important point.", "weaknesses": "The main weakness of the paper is that even though the theory is agnostic to the loss function and dataset, it seems that it is not agnostic to the architecture given section 4.3. So it would maybe make sense to also try more plausible architectures than ANNs to test the theory like continuous Hopfield networks or Spiking networks trained with surrogate gradients for instance, or at least vanilla RNNs, since the brain is highly recurrent. This would provide a sense of how important is the architecture vs the learning rule.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XVDPQVLycx", "reviewer": "Reviewer_mbP3", "summary": "In this work, the authors use mirror descent principles to derive the distribution of final synaptic weights, under certain assumptions. Theoretical findings demonstrate that this distribution is determined by mirror descent potentials. Analyzing synaptic weight distributions before and after training is crucial for understanding brain mechanisms. The paper applies mirror descent theory to distinguish learning rules in the brain, modeling it under chosen synaptic geometry. Drawing from this theory, the paper shows that, under specific assumptions, weight changes in dual space follow a Gaussian distribution. This insight aids in inferring synaptic geometry from the weight distribution. The authors validate this approach through experiments on artificial neural networks and demonstrate its applicability to real neural data.", "questions": "see Weaknesses", "limitations": "", "rating": 1, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "The application of concepts from mirror descent to underscore the significance of selecting the appropriate distance function is a commendable idea.", "weaknesses": "The main drawback is the multitude of assumptions made in the article, leading to a highly robust conclusion. I've outlined what I consider to be the most unreasonable five assumptions:\n\n1. Our brain is a feedforward network with no feedback and lacks any dynamic processes.\n2. During learning, only the neurons in the last layer of our brain update.\n3. Our brain should exclusively perform a supervised learning task.\n4. Our brain uses gradient descent algorithms.\n5. The optimal values for our brain's neurons for a given problem are unique and deterministic.\n\nEach of these five assumptions is overly restrictive, and some are evidently incorrect, conflicting with known experimental evidence. \n\nMoreover, the conclusion itself also appears highly unreasonable: the distribution of synaptic strength in our brain neurons, as well as the brain's structure and the environmental context (task execution), appear to be unrelated. The distribution of synaptic strength is solely related to geometry.\n\n And the conclusion itself lacks significance. What matters is, assuming the article's conclusion is correct, why we would choose this specific geometry, what advantages it offers, and how we arrive at such a choice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)