Title: SPIKE-BASED CAUSAL INFERENCE FOR WEIGHT ALIGNMENT

Abstract: In artificial neural networks trained with gradient descent, the weights used for processing stimuli are also used during backward passes to calculate gradients. For the real brain to approximate gradients, gradient information would have to be propagated separately, such that one set of synaptic weights is used for processing and another set is used for backward passes. This produces the so-called "weight transport problem" for biological models of learning, where the backward weights used to calculate gradients need to mirror the forward weights used to process stimuli. This weight transport problem has been considered so hard that popular proposals for biological learning assume that the backward weights are simply random, as in the feedback alignment algorithm. However, such random weights do not appear to work well for large networks. Here we show how the discontinuity introduced in a spiking system can lead to a solution to this problem. The resulting algorithm is a special case of an estimator used for causal inference in econometrics, regression discontinuity design. We show empirically that this algorithm rapidly makes the backward weights approximate the forward weights. As the backward weights become correct, this improves learning performance over feedback alignment on tasks such as Fashion-MNIST, SVHN, CIFAR-10 and VOC. Our results demonstrate that a simple learning rule in a spiking network can allow neurons to produce the right backward connections and thus solve the weight transport problem.

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. W ji = W ij (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, R SA , as one possible cost function that, when minimized, leads to weight symmetry:

R SA := W -Y T 2 F (1) = W 2 F + Y 2 F -2tr(W Y )

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, R SA .

• 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.

4.3. LIF DYNAMICS:

During the RDD feedback training phase, each unit in the network is simulated as a leaky integrateand-fire neuron. Spiking inputs from the previous layer arrive at feedforward synapses, where they are convolved with a temporal exponential kernel to simulate post-synaptic spike responses p = [p 1 , p 2 , ..., p m ] (see Appendix A.1). The neurons can also receive driving input pi , instead of synaptic inputs. The total feedforward input to neuron i is thus defined as: is the maximum input drive during this window of time. Bottom: Feedback received at a synapse, q k . q pre k is the feedback signal at the start of an RDD window, while q avg k is the average of the feedback signal during the time window. D. Samples of ∆q avg k vs. u max i are used to update a piece-wise linear function of u max i , and the causal effect β ik is defined as the difference of the left and right limits of the function at the spiking threshold.

I i := ω pi if pi > 0 m j=1 W ij p j otherwise (2)

where W ij is the feedforward weight to neuron i from neuron j in the previous layer, and ω is a hyperparameter. The voltage of the neuron, v i , evolves as:

EQUATION

where g L and g D are leak and dendritic conductance constants, respectively. The input drive to the neuron, u i , is similarly modeled:

EQUATION

If the voltage v i passes a spiking threshold θ, the neuron spikes and the voltage is reset to a value v reset = -1 (Figure 2C ). Note that the input drive does not reset. This helps us to perform regressions both above and below the spike threshold.

In addition to feedforward inputs, spiking inputs from the downstream layer arrive at feedback synapses, where they create post-synaptic spike responses q = [q 1 , q 2 , ..., q n ]. These responses are used in the causal effect estimation (see below).

4.4. RDD ALGORITHM:

Whenever the voltage approaches the threshold θ (ie. |v i -θ| < α where α is a constant), an RDD window is initiated, lasting T = 30 ms in simulated time (Figure 2C ). At the end of this time window, at each feedback synapse, the maximum input drive during the RDD window, u max i , and the average change in feedback from downstream neuron k during the RDD window, ∆q avg k , are recorded. ∆q avg k is defined as the difference between the average feedback received during the RDD window, q avg k , and the feedback at the start of the RDD window, q pre k :

∆q avg k := q avg k -q pre k (5)

Importantly, u max i provides a measure of how strongly neuron i was driven by its inputs (and whether or not it passed the spiking threshold θ), while ∆q avg k is a measure of how the input received as feedback from neuron k changed after neuron i was driven close to its spiking threshold. These two values are then used to fit a piece-wise linear model of ∆q avg k as a function of u max i (Figure 2D ). This piece-wise linear model is defined as:

f ik (x) := c 1 ik x + c 2 ik if x < θ c 3 ik x + c 4 ik if x ≥ θ (6)

The parameters c 1 ik , c 2 ik , c 3 ik and c 4 ik are updated to perform linear regression using gradient descent:

EQUATION

An estimate of the causal effect of neuron i spiking on the activity of neuron k, β ik , is then defined as the difference in the two sides of the piece-wise linear function at the spiking threshold:

EQUATION

Finally, the weight at the feedback synapse, Y ik , is updated to be a scaled version of β ik :

EQUATION

)

where γ is a hyperparameter and σ 2 β is the standard deviation of β values for all feedback synapses in the layer. This ensures that the scale of the full set of feedback weights between two layers in the network remains stable during training.

5. RESULTS:



5.1. ALIGNMENT OF FEEDBACK AND FEEDFORWARD WEIGHTS:

To measure how well the causal effect estimate at each feedback synapse, β ik , and thus the feedback weight Y ik , reflects the reciprocal feedforward weight W ki , we can measure the percentage of feedback weights that have the same sign as the reciprocal feedforward weights (Figure 3A ). When training on CIFAR-10 with no RDD feedback training phase (ie. feedback weights remain fixed throughout training), the feedback alignment effect somewhat increases the sign alignment during training, but it is ineffective at aligning the signs of weights in earlier layers in the network. Compared to feedback alignment, the addition of an RDD feedback training phase greatly increases the sign aligmnent between feedback and feedforward weights for all layers in the network, especially at earlier layers. In addition, the RDD algorithm increases sign alignment throughout the hierarchy more than the current state-of-the-art algorithm for weight alignment introduced recently by Akrout et al. Akrout et al. (2019) (Figure 3A ). Furthermore, RDD feedback training changes feedback weights to not only match the sign but also the magnitude of the reciprocal feedforward weights (Figure 3B ), which makes it better for weight alignment than hard-wired sign symmetry (Moskovitz et al., 2018) .

5.2. DESCENDING THE SYMMETRIC ALIGNMENT COST FUNCTION:

The symmetric alignment cost function (Kunin et al., 2019) (Equation 1) can be broken down as:

EQUATION

where we define R decay and R self as:

EQUATION

R decay is simply a weight regularization term that can be minimized using techniques like weight decay. R self , in contrast, measures how well aligned in direction the two matrices are. Our learning rule for feedback weights minimizes the R self term for weights throughout the network (Figure 4 ). By comparison, feedback alignment decreases R self to a smaller extent, and its ability to do so diminishes at earlier layers in the network. This helps to explain why our algorithm induces weight alignment, and can improve training performance (see below).

5.3. PERFORMANCE ON FASHION-MNIST, SVHN, CIFAR-10 AND VOC:

We trained the same network architecture (see Appendix A.3) on the Fashion-MNIST, SVHN, CIFAR-10 and VOC datasets using standard autograd techniques (backprop), feedback alignment and our RDD feedback training phase. RDD feedback training substantially improved the network's performance over feedback alignment, and led to backprop-level accuracy on the train and test sets (Figure 5 ). 

6. DISCUSSION:

In order to understand how the brain learns complex tasks that require coordinated plasticity across many layers of synaptic connections, it is important to consider the weight transport problem. Here, we presented an algorithm for updating feedback weights in a network of spiking neurons that takes advantage of the spiking discontinuity to estimate the causal effect between two neurons (Figure 2 ). We showed that this algorithm enforces weight alignment (Figure 3 ), and identified a loss function, R self , that is minimized by our algorithm (Figure 4 ). Finally, we demonstrated that our algorithm allows deep neural networks to achieve better learning performance than feedback alignment on Fashion-MNIST and CIFAR-10 (Figure 5 ). These results demonstrate the potential power of RDD as a means for solving the weight transport problem in biologically plausible deep learning models.

One aspect of our algorithm that is still biologically implausible is that it does not adhere to Dale's principle, which states that a neuron performs the same action on all of its target cells (Strata & Harvey) . This means that a neuron's outgoing connections cannot include both positive and negative weights. However, even under this constraint, a neuron can have an excitatory effect on one downstream target and an inhibitory effect on another, by activating intermediary inhibitory interneurons.

Because our algorithm provides a causal estimate of one neuron's impact on another, theoretically, it could capture such polysynaptic effects. Therefore, this algorithm is in theory compatible with Dale's principle. Future work should test the effects of this algorithm when implemented in a network of neurons that are explicitly excitatory or inhibitory.

A APPENDIX:

A.1 LIF NEURON SIMULATION DETAILS Post-synaptic spike responses at feedforward synapses, p, were calculated from pre-synaptic binary spikes using an exponential kernel function κ:

EQUATION

)

where tjk is the k th spike time of input neuron j and κ is given by:

EQUATION

where τ s = 0.003 s and τ L = 0.01 s represent short and long time constants, and Θ is the Heaviside step function. Post-synaptic spike responses at feedback synapses, q, were computed in the same way.

A.2 RDD FEEDBACK TRAINING IMPLEMENTATION:

A.2.1 WEIGHT SCALING Weights were shared between the convolutional network and the network of LIF neurons, but feedforward weights in the LIF network were scaled versions of the convolutional network weights:

EQUATION

where W Conv is a feedforward weight matrix in the convolutional network, W LIF is the corresponding weight matrix in the LIF network, m is the number of units in the upstream layer (ie. the number of columns in W Conv ), σ 2 W Conv is the standard deviation of W Conv and ψ is a hyperparameter. This rescaling ensures that spike rates in the LIF network stay within an optimal range for the RDD algorithm to converge quickly, even if the scale of the feedforward weights in the convolutional network changes during training. This avoids situations where the scale of feedforward weights is so small that little or no spiking occurs in the LIF neurons.

A.2.2 FEEDBACK TRAINING PARADIGM:

The RDD feedback training paradigm is implemented as follows. We start by providing driving input to the first layer in the network of LIF neurons. To create this driving input, we choose a subset of 20% of the neurons in that layer, and create a unique input spike train for each of these neurons using a Poisson process with a rate of 200 Hz. All other neurons in the layer receive no driving input. Every 100 ms, a new set of neurons to receive driving input is randomly chosen. After 30 s, this layer stops receiving driving input, and the process repeats for the next layer in the network.

A.3 NETWORK AND TRAINING DETAILS:

The network architectures used to train on Fashion-MNIST and CIFAR-10 are described in Table 1 .

Inputs were randomly cropped and flipped during training, and batch normalization was used at each layer. Networks were trained using a minibatch size of 32.

A.4 AKROUT ET AL. (2019) ALGORITHM IMPLEMENTATION:

In experiments that compared sign alignment using our RDD algorithm with the Akrout et al. ( 2019) algorithm, we kept the same RDD feedback training paradigm (ie. layers were sequentially driven, and a small subset of neurons in each layer was active at once). However, rather than updating feedback weights using RDD, we recorded the mean firing rates of the active neurons in the upstream layer, r l , and the mean firing rates in the downstream layer, r l+1 . We then used the following feedback weight update rule:

EQUATION

)

where Y are the feedback weights between layers l + 1 and l, and η and λ WD are learning rate and weight decay hyperparameters, respectively. 

:



Figures and Tables:

Figure 1: Illustration of weight symmetry in a neural network with feedforward and feedback connections. Processing of inputs to outputs is mirrored by backward flow of gradient information. Gradient estimation is best when feedback synapses have symmetric weights to feedforward synapses (illustrated with colored circles).

Figure 2: A. Layers of the convolutional network trained on CIFAR-10 and the corresponding network of LIF neurons that undergoes RDD feedback training. Fully-connected feedforward weights (blue) and feedback weights (red) are shared between the two networks. Every training epoch consists of an RDD feedback training phase where feedback weights in the LIF net are updated (bold red arrows) and transferred to the convolutional net, and a regular training phase where feedforward weights in the convolutional net are updated (bold blue arrows) and transferred back to the LIF net. B. RDD feedback training protocol. Every 30 s, a different layer in the LIF network receives driving input and updates its feedback weights (red) using the RDD algorithm. C. Top: Sample voltage (vi, solid line) and input drive (ui, dashed line) traces. Whenever vi approaches the spiking threshold, an RDD window lasting T ms is triggered. u max i

Figure 3: A. Evolution of sign alignment (the percent of feedforward and feedback weights that have the same sign) for each fully-connected layer in the network when trained on CIFAR-10 using RDD feedback training (blue), using the algorithm proposed by Akrout et al. (2019) (purple), and using feedback alignment (red). B. Feedforward vs. feedback weights for each fully-connected layer at the end of training, with RDD feedback training (blue), the Akrout algorithm (purple), and feedback alignment (red).

Figure 4: Evolution of Rself for each fully-connected layer in the network when trained on CIFAR-10 using RDD feedback training (solid lines) and using feedback alignment (dashed lines). RDD feedback training dramatically decreases this loss compared to feedback alignment, especially in earlier layers.

Figure 5: Comparison of Fashion-MNIST, SVHN, CIFAR-10 and VOC train error (top row) and test error (bottom row). RDD feedback training substantially improves test error performance over feedback alignment in both learning tasks.

Figure S1: Comparison of average spike rates in the fully-connected layers of the LIF network vs. activities of the same layers in the convolutional network, when both sets of layers were fed the same input. Spike rates in the LIF network are largely correlated with activities of units in the convolutional network.

× 5, 64 ReLU Conv2D 5 × 5, 64 ReLU Conv2D 5 × 5, 64 ReLU 2 MaxPool 2 × 2, stride 2 MaxPool 2 × 2, stride 2 MaxPool 2 × 2, stride 2 3 Conv2D 5 × 5, 64 ReLU Conv2D 5 × 5, 64 ReLU Conv2D 5 × 5, 64 ReLU Network architectures used to train on Fashion-MNIST, SVHN, CIFAR-10 and VOC.

Table content: <table><tr><td>Layer</td><td>Fashion-MNIST</td><td>SVHN &amp; CIFAR-10</td><td>VOC</td></tr><tr><td>Input</td><td>28 × 28 × 1</td><td>32 × 32 × 3</td><td>32 × 32 × 3</td></tr><tr><td colspan="2">1 Conv2D 5 4 MaxPool 2 × 2, stride 2</td><td>MaxPool 2 × 2, stride 2</td><td>MaxPool 2 × 2, stride 2</td></tr><tr><td>5</td><td>FC 384 ReLU</td><td>FC 384 ReLU</td><td>FC 384 ReLU</td></tr><tr><td>6</td><td>FC 192 ReLU</td><td>FC 192 ReLU</td><td>FC 192 ReLU</td></tr><tr><td>7</td><td>FC 10 ReLU</td><td>FC 10 ReLU</td><td>FC 21 ReLU</td></tr></table>

Peer review (iclr2020):
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. Compared to feedback alignment, RDD-based pre-training demonstrates stronger alignment between forward and backward weights and better performance on CIFAR-10 and Fashion-MNIST datasets. Overall, the paper is very well written and addresses an important problem. The theoretical foundation, to my knowledge, is well studied.