Spike-based causal inference for weight alignment

1 INTRODUCTION . Any learning system that makes small changes to its parameters will only improve if the changes are correlated to the gradient of the loss function . Given that people and animals can also show clear behavioral improvements on specific tasks ( Shadmehr et al. , 2010 ) , however the brain determines its synaptic updates , on average , the changes in must also correlate with the gradients of some loss function related to the task ( Raman et al. , 2019 ) . As such , the brain may have some way of calculating at least an estimator of gradients . To-date , the bulk of models for how the brain may estimate gradients are framed in terms of setting up a system where there are both bottom-up , feedforward and top-down , feedback connections . The feedback connections are used for propagating activity that can be used to estimate a gradient ( Williams , 1992 ; Lillicrap et al. , 2016 ; Akrout et al. , 2019 ; Roelfsema & Ooyen , 2005 ; Lee et al. , 2015 ; Scellier & Bengio , 2017 ; Sacramento et al. , 2018 ) . In all such models , the gradient estimator is less biased the more the feedback connections mirror the feedforward weights . For example , in the REINFORCE algorithm ( Williams , 1992 ) , and related algorithms like AGREL ( Roelfsema & Ooyen , 2005 ) , learning is optimal when the feedforward and feedback connections are perfectly symmetric , such that for any two neurons i and j the synaptic weight from i to j equals the weight from j to i , e.g . Wji = Wij ( Figure 1 ) . Some algorithms simply assume weight symmetry , such as Equilibrium Propagation ( Scellier & Bengio , 2017 ) . The requirement for synaptic weight symmetry is sometimes referred to as the “ weight transport problem ” , since it seems to mandate that the values of the feedforward synaptic weights are somehow transported into the feedback weights , which is not biologically realistic ( Crick , 1989-01-12 ; Grossberg , 1987 ) . Solving the weight transport problem is crucial to biologically realistic gradient estimation algorithms ( Lillicrap et al. , 2016 ) , and is thus an important topic of study . Several solutions to the weight transport problem have been proposed for biological models , including hard-wired sign symmetry ( Moskovitz et al. , 2018 ) , random fixed feedback weights ( Lillicrap et al. , 2016 ) , and learning to make the feedback weights symmetric ( Lee et al. , 2015 ; Sacramento et al. , 2018 ; Akrout et al. , 2019 ; Kolen & Pollack , 1994 ) . Learning to make the weights symmetric is promising because it is both more biologically feasible than hard-wired sign symmetry ( Moskovitz et al. , 2018 ) and it leads to less bias in the gradient estimator ( and thereby , better training results ) than using fixed random feedback weights ( Bartunov et al. , 2018 ; Akrout et al. , 2019 ) . However , of the current proposals for learning weight symmetry some do not actually work well in practice ( Bartunov et al. , 2018 ) and others still rely on some biologically unrealistic assumptions , including scalar value activation functions ( as opposed to all-or-none spikes ) and separate error feedback pathways with one-to-one matching between processing neurons for the forward pass and error propagation neurons for the backward pass Akrout et al . ( 2019 ) ; Sacramento et al . ( 2018 ) . Interestingly , learning weight symmetry is implicitly a causal inference problem—the feedback weights need to represent the causal influence of the upstream neuron on its downstream partners . As such , we may look to the causal infererence literature to develop better , more biologically realistic algorithms for learning weight symmetry . In econometrics , which focuses on quasi-experiments , researchers have developed various means of estimating causality without the need to actually randomize and control the variables in question Angrist & Pischke ( 2008 ) ; Marinescu et al . ( 2018 ) . Among such quasi-experimental methods , regression discontinuity design ( RDD ) is particularly promising . It uses the discontinuity introduced by a threshold to estimate causal effects . For example , RDD can be used to estimate the causal impact of getting into a particular school ( which is a discontinuous , all-or-none variable ) on later earning power . RDD is also potentially promising for estimating causal impact in biological neural networks , because real neurons communicate with discontinuous , all-or-none spikes . Indeed , it has been shown that the RDD approach can produce unbiased estimators of causal effects in a system of spiking neurons Lansdell & Kording ( 2019 ) . Given that learning weight symmetry is fundamentally a causal estimation problem , we hypothesized that RDD could be used to solve the weight transport problem in biologically realistic , spiking neural networks . Here , we present a learning rule for feedback synaptic weights that is a special case of the RDD algorithm previously developed for spiking neural networks ( Lansdell & Kording , 2019 ) . Our algorithm takes advantage of a neuron ’ s spiking discontinuity to infer the causal effect of its spiking on the activity of downstream neurons . Since this causal effect is proportional to the feedforward synaptic weight between the two neurons , by estimating it , feedback synapses can align their weights to be symmetric with the reciprocal feedforward weights , thereby overcoming the weight transport problem . We demonstrate that this leads to the reduction of a cost function which measures the weight symmetry ( or the lack thereof ) , that it can lead to better weight symmetry in spiking neural networks than other algorithms for weight alignment ( Akrout et al. , 2019 ) and it leads to better learning in deep neural networks in comparison to the use of fixed feedback weights ( Lillicrap et al. , 2016 ) . Altogether , these results demonstrate a novel algorithm for solving the weight transport problem that takes advantage of discontinuous spiking , and which could be used in future models of biologically plausible gradient estimation . 2 RELATED WORK . Previous work has shown that even when feedback weights in a neural network are initialized randomly and remain fixed throughout training , the feedforward weights learn to partially align themselves to the feedback weights , an algorithm known as feedback alignment ( Lillicrap et al. , 2016 ) . While feedback alignment is successful at matching the learning performance of true gradient descent in relatively shallow networks , it does not scale well to deeper networks and performs poorly on difficult computer vision tasks ( Bartunov et al. , 2018 ) . The gap in learning performance between feedback alignment and gradient descent can be overcome if feedback weights are continually updated to match the sign of the reciprocal feedforward weights ( Moskovitz et al. , 2018 ) . Furthermore , learning the feedback weights in order to make them more symmetric to the feedforward weights has been shown to improve learning over feedback alignment ( Akrout et al. , 2019 ) . To understand the underlying dynamics of learning weight symmetry , Kunin et al . ( 2019 ) define the symmetric alignment cost function , RSA , as one possible cost function that , when minimized , leads to weight symmetry : RSA : = ‖W − Y T ‖2F ( 1 ) = ‖W‖2F + ‖Y ‖2F − 2tr ( WY ) where W are feedforward weights and Y are feedback weights . The first two terms are simply weight regularization terms that can be minimized using techniques like weight decay . But , the third term is the critical one for ensuring weight alignment . In this paper we present a biologically plausible method of minimizing the third term . This method is based on the work of Lansdell & Kording ( 2019 ) , who demonstrated that neurons can estimate their causal effect on a global reward signal using the discontinuity introduced by spiking . This is accomplished using RDD , wherein a piecewise linear model is fit around a discontinuity , and the differences in the regression intercepts indicates the causal impact of the discontinuous variable . In Lansdell & Kording ( 2019 ) , neurons learn a piece-wise linear model of a reward signal as a function of their input drive , and estimate the causal effect of spiking by looking at the discontinuity at the spike threshold . Here , we modify this technique to perform causal inference on the effect of spiking on downstream neurons , rather than a reward signal . We leverage this to develop a learning rule for feedback weights that induces weight symmetry and improves training . 3 OUR CONTRIBUTIONS . The primary contributions of this paper are as follows : • We demonstrate that spiking neurons can accurately estimate the causal effect of their spiking on downstream neurons by using a piece-wise linear model of the feedback as a function of the input drive to the neuron . • We present a learning rule for feedback weights that uses the causal effect estimator to encourage weight symmetry . We show that when feedback weights update using this algorithm it minimizes the symmetric alignment cost function , RSA . • We demonstrate that this learning weight symmetry rule improves training and test accuracy over feedback alignment , approaching gradient-descent-level performance on Fashion-MNIST , SVHN , CIFAR-10 and VOC in deeper networks . 4 METHODS . 4.1 GENERAL APPROACH . In this work , we utilize a spiking neural network model for aligning feedforward and feedback weights . However , due to the intense computational demands of spiking neural networks , we only use spikes for the RDD algorithm . We then use the feedback weights learned by the RDD algorithm for training a non-spiking convolutional neural network . We do this because the goal of our work here is to develop an algorithm for aligning feedback weights in spiking networks , not for training feedforward weights in spiking networks on other tasks . Hence , in the interest of computational expediency , we only used spiking neurons when learning to align the weights . Additional details on this procedure are given below . 4.2 RDD FEEDBACK TRAINING PHASE . At the start of every training epoch of a convolutional neural network , we use an RDD feedback weight training phase , during which all fully-connected sets of feedback weights in the network are updated . To perform these updates , we simulate a separate network of leaky integrate-and-fire ( LIF ) neurons . LIF neurons incorporate key elements of real neurons such as voltages , spiking thresholds and refractory periods . Each epoch , we begin by training the feedback weights in the LIF network . These weights are then transferred to the convolutional network , which is used for training the feedforward weights . The new feedforward weights are then transferred to the LIF net , and another feedback training phase with the LIF net starts the next epoch ( Figure 2A ) . During the feedback training phase , the LIF network undergoes a training phase lasting 90 s of simulated time ( 30 s per set of feedback weights ) ( Figure 2B ) . We find that the spiking network used for RDD feedback training and the convolutional neural network are very closely matched in the activity of the units ( Figure S1 ) , which gives us confidence that this approach of using a separate non-spiking network for training the feedforward weights is legitimate . During the feedback training phase , a small subset of neurons in the first layer receive driving input that causes them to spike , while other neurons in this layer receive no input ( see Appendix A.2 ) . The subset of neurons that receive driving input is randomly selected every 100 ms of simulated time . This continues for 30 s in simulated time , after which the same process occurs for the subsequent hidden layers in the network . This protocol enforces sparse , de-correlated firing patterns that improve the causal inference procedure of RDD .

Summary: Strong paper in the direction of a more biologically plausible solution for the weight transport problem, where the forward and the backward weights need to be aligned. Earlier work for feedback alignment has included methods such as hard-coding sign symmetry. In this method, the authors show that a piece-wise linear model of the feedback as a function of the input given to a neuron can estimate the causal effect of a spike on downstream neurons. The authors propose a learning rule based on regression discontinuity design (RDD) and show that this leads to stronger alignment of weights (especially in earlier layers) compared to previous methods. The causal effect is measured directly from the discontinuity introduced while spiking - the difference between the outputs of the estimated piece-wise linear model at the point of discontinuity is used as the feedback.