Learning to solve the credit assignment problem
Abstract
Backpropagation is driving today’s artiﬁcial neural networks.
However, despite extensive research, it remains unclear if the brain implements this algorithm.
Among neuroscientists, reinforcement learning (RL) algorithms are often seen as a realistic alternative.
However, the convergence rate of such learning scales poorly with the number of involved neurons.
Here we propose a hybrid learning approach, in which each neuron uses an RL-type strategy to learn how to approximate the gradients that backpropagation would provide.
We show that our approach learns to approximate the gradient, and can match the performance of gradient-based learning on fully connected and convolutional networks.
Learning feedback weights provides a biologically plausible mechanism of achieving good performance, without the need for precise, pre-speciﬁed learning rules.
Introduction
It is unknown how the brain solves the credit assignment problem when learning: how does each neuron know its role in a positive (or negative) outcome, and thus know how to change its activity to perform better next time?
Biologically plausible solutions to credit assignment include those based on reinforcement learning (RL) algorithms [4].
In these approaches a globally distributed reward signal provides feedback to all neurons in a network.
However these methods have not been demonstrated to operate at scale.
For instance, variance in the REINFORCE estimator scales with the number of units in the network.
This drives the hypothesis that learning in the brain must rely on additional structures beyond a global reward signal.
In artiﬁcial neural networks credit assignment is performed with gradient-based methods computed through backpropagation.
This is signiﬁcantly more efﬁcient than RL-based algorithms.
However there are well known problems with implementing backpropagation in biologically realistic neural networks.
For instance backpropagation requires a feedback structure with the same weights as the feedforward network to communicate gradients (so-called weight transport).
Yet such structures are not observed in neural circuits.
Despite this, backpropagation is the only method known to solve learning problems at scale.
Thus modiﬁcations or approximations to backpropagation that are more plausible have been the focus of signiﬁcant recent attention [8, 3].
Notably, it turns out that weight transport can be avoided by using ﬁxed, random feedback weights, through a phenomenon called feedback alignment [8].
However feedback alignment does not work in larger, more complicated network architectures (such as convolutional networks).
Here we propose to use an RL algorithm to train a feedback system to enable learning.
We propose to use a REINFORCE-style perturbation approach to train a feedback signal to approximate what would have been provided by backpropagation.
We demonstrate that our model learns as well as regular backpropagation in small models, overcomes the limitations of ﬁxed random feedback weights ("feedback alignment") on more complicated feedforward networks, and can be utilized in convolutional networks.
Our method illustrates a biologically realistic way the brain could perform gradient descent-like learning.
Method
Let an N hidden-layer network be given by ˆ y
= f ( x ) ∈ R p , composed of a set of layer-wise summation and non-linear activations
h i =
f i ( h i − 1 ) =
σ (cid:0) W i
h
i − 1 (cid:1) , for hidden layer states h i ∈ R n i , non-linearity σ and with input
h
= x and output h
N +1 = ˆy .
Deﬁne L as the loss function L ( x , y ) , where the data ( x , y ) ∈ D are drawn from a distribution ρ .
Our aim is then to minimize: E ρ [ L ( x , y )] .
Backpropagation computes the error signal ˜
e i in a top-down fashion:
[Equation]
Let the loss gradient term be denoted as λ i =
∂
L ∂ h
i =
( W i +1 )
T e i +1 .
Here we replace λ i with an approximation, with its own parameters to be learned:
λ i
≈ g
( h i , ˜ e i +1 ; B ) , for parameters B .
We will use ˜ e i to denote the gradient signal backpropagated through the synthetic gradients, and e i for the true gradients.
To estimate B we use stochasticity inherent to biological neural networks.
For each input each unit produces a noisy response: h it = σ (cid:0)(cid:80)
k W i · k h
i − 1 t (cid:1)
+ c
h ξ it , for independent Gaussian noise
ξ i ∼ ν
= N (0 , I ) with standard deviation c h > 0 .
This then generates a noisy loss ˜
L ( x , y , ξ ) and a baseline loss L ( x , y ) =
˜ L ( x , y , 0) .
We will use the noisy response to estimate gradients, that then allow us to optimize the baseline L .
This is achieved by linearizing the loss: ˜ L ≈ L + ∂ L
∂h ij c h ξ ij .
This gives E (( ˜ L − L ) c
h ξ ij | x , y ) ≈ c 2 h ∂
L ∂h
ij | x , y , with expectation taken over the noise distribution ν ( ξ ) .
Thus we obtain an estimator of the loss gradient: ˆ λ i :=
( ˜ L ( x , y , ξ ) − L ( x , y ))
ξ i c h .
Applications 3.1 Fully connected networks solving MNIST
To demonstrate the method can be used to solve simple supervised learning problems we use node perturbation with a four-layer network and MSE loss to solve MNIST (Fig. 1).
We approximate loss gradients as follows:
g ( h i , ˜
e i +1 ; B ) =
( B i +1 )
T ˜ e i +1 .
The feedback parameters B
i +1 are estimated by solving the least squares problem: ˆ B i +1 =
argmin B E (cid:107) B T ˜ e i +
− ˆ λ
i (cid:107) 22 , where ˆ λ is the perturbationbased estimator derived above.
B is updated with each mini-batch using stochastic gradient-descent to minimize this loss.
1 Updates to W i are made using the synthetic gradients ∆ W i =
η ˜ e i h
i − 1 , for learning rate η .
We observed that the system is able to provide a close correspondence between the feedforward and feedback matrices in both layers of the network (Fig. 1a).
The relative error between
B i and W i is lower than what is observed for feedback alignment, suggesting that this co-adaptation of W i and B i is indeed beneﬁcial.
We observe that the alignment (the angle between the estimated gradient and the true gradient, proportional to e T WB T ˜ e ) is lower for node perturbation than for feedback alignment (Fig. 1b).
Recent studies have shown that sign congruence of the feedforward and feedback matrices is all that is required to achieve good performance [10].
Here the sign congruence is also higher in node perturbation (Fig. 1c).
Finally, the learning performance of node perturbation is comparable to backpropagation (Fig. 1d) – achieving close to 3% test error.
These suggest node perturbation for learning feedback weights can be used in deep networks. [Figure] l a y e r [Figure]
[Figure] 3.2 Auto-encoding MNIST
A known shortcoming of feedback alignment is in auto-encoding networks with tight bottleneck layers [8].
To see if our method has the same shortcoming we examine a simple auto-encoding network with MNIST input data (size 784-200-2-200-784, MSE loss).
We also compare the method to the ‘matching’ learning rule [9], in which updates to B match updates to W .
As expected, feedback alignment performs poorly.
Node perturbation actually performs better than backpropagation, and comparable to ADAM (Fig. 2a).
In fact ADAM begins to overﬁt, while node perturbation does not.
The matched learning rule performs similarly to backpropagation.
These results are surprising at ﬁrst glance.
Perhaps, similar to feedback alignment, learning feedback weights strikes the right balance between providing a useful signal to learn, and constraining the updates to be sufﬁciently aligned with B , acting as a type of regularization [8].
The noise added when estimating the feedback weights may also serve to regularize the latent representation, as, indeed, the latent space learnt by node perturbation shows a more evenly distributed separation of digits.
While, in contrast, the representations learnt by backprop and ADAM show more structure, and feedback alignment does not learn a useful representation at all (Fig. 2b,c).
These results show that node perturbation is able to successfully communicate error signals through thin layers of a network as needed.
3.3 Convolutional neural networks solving CIFAR
Finally we test the method on a convolutional neural network (CNN) solving CIFAR10.
The CNN has the architecture Conv(3x3, 1x1, 32), MaxPool(3x3, 2x2), Conv(5x5, 1x1, 128), MaxPool(3x3, 2x2),
Conv(5x5, 1x1, 256), MaxPool(3x3, 2x2), FC 2048, FC 2048, Softmax(10), with hyperparameters found through random search.
For this network we learn feedback weights direct from the output layer to each earlier layer: g ( h i , ˜
e i +1 ; θ i ) =
( B i +1 ) T ˜
e N (similar to ‘direct feedback alignment’).
Here this was solved by gradient-descent.
We obtain a test accuracy of 75.2%.
When compared with ﬁxed feedback weights (test accuracy of 72.5%) and backpropagation (test accuracy of 77.2%), we see it is advantageous to learn feedback weights. This shows the method can be used in a CNN, and can solve challenging computer vision problems without weight transport.
Discussion
Here we implement a perturbation-based synthetic gradient method to train neural networks.
We show that this hybrid approach can be used in both fully connected and convolutional networks.
By removing both the symmetric feedforward, feedback weight requirement imposed by backpropagation this approach is a step towards more biologically-plausible deep learning.
In contrast to many perturbation-based methods, this hybrid approach can solve large-scale problems.
We thus believe this approach can provide powerful and biologically plausible learning algorithms.
While previous research has provided some insight and theory for how feedback alignment works [8, 3, 2] the effect remains somewhat mysterious, and not applicable in some network architectures.
Recent studies have shown that some of these weaknesses can be addressed by instead imposing sign congruent feedforward and feedback matrices [10].
Yet what mechanism may produce congruence in biological networks is unknown.
Here we show that the shortcomings of feedback alignment can be addressed in another way: the system can learn to adjust weights as needed to provide a useful error signal.
Our work is closely related to Akrout et al 2019 [1], which also uses perturbations to learn feedback weights.
However our approach does not divide learning into two phases, and training of the feedback weights does not occur in a layer-wise fashion.
Here we tested our method in an idealized setting, however it is consistent with neurobiology in two important ways.
First, it involves the separate learning of feedforward and feedback weights.
This is possible in cortical networks where complex feedback connections exist between layers, and where pyramidal cells have apical and basal compartments that allow for separate integration of feedback and feedforward signals [5].
Second, noisy perturbations are common in neural learning models.
There are many mechanisms by which noise can be measured or approximated [4, 7], or neurons could use a learning rule that does not require knowing the noise [6].
While our model involves the subtraction of a baseline loss to reduce the variance of the estimator, this does not affect the expected value of the estimator; technically the baseline could be removed or approximated [7].
Thus we believe our approach could be implemented in neural circuits.
There is a large space of plausible learning rules that can learn feedback signals in order to more efﬁciently learn.
These promise to inform both models of learning in the brain and learning algorithms in artiﬁcial networks.
Here we take an early step in this direction. Supplementary Material A Proofs
We review the key components of the model.
Data ( x , y ) ∈ D are drawn from a distribution ρ .
The loss function is linearized:
[Equation]
such that
[Equation]
with expectation taken over the noise distribution ν ( ξ ) .
This suggests a good estimator of the loss gradient is
[Equation]
Let ˜ e i be the error signal computed by backpropagating the synthetic gradients:
[Equation]
Then parameters B i +1 are estimated by solving the least squares problem:
[Equation]
{+Under what conditions can we show that ˆ B i +
+}
→ W i +1 (with enough data)?
One way to ﬁnd an answer is to deﬁne the synthetic gradient in terms of the system without noise added.
Then
B T ˜ e is deterministic with respect to x , y and, assuming ˜ L has a convergent power series around ξ = 0 , we can write
[Equation]
Taken together these suggest we can prove ˆ B
{+i +
+}
→ W i +1 in the same way we prove consistency of the linear least squares estimator.
For this to work we must show the expectation of the Taylor series approximation (2) is well behaved.
That is, we must show the expected remainder term of the expansion:
[Equation]
is ﬁnite and goes to zero as c h → 0 .
This requires some additional assumptions on the problem.
We make the following assumptions:
• A1: the noise ξ is subgaussian, • A2: the loss function L ( x , y ) is analytic on D , • A3: the error matrices ˜e n ( ˜e n ) T are full rank, for 1 ≤ n ≤ N + 1 , • A4: the mean of the remainder and error terms is bounded:
[Equation]
for 1 ≤ n ≤ N .
Consider ﬁrst convergence of the ﬁnal layer feedback matrix, B N +1 .
In the ﬁnal layer it is true that e N +1 =
˜
e N +1 .
Theorem 1.
Assume A1-4.
For g FA ( h i , ˜
e i +1 ; B i +1 )
=
B i +1 ˜
e i +1 , then the least squares estimator
[Equation]
5) solves (4) and converges to the true feedback matrix, in the sense that:
[Equation]
Proof.
Let L ( m ) ij :=
∂ m L ∂h imj .
We ﬁrst show that, under A1-2, the conditional expectation of the estimator (5) converges to the gradient L (1) Nj as c h → 0 .
For each ˆ λ Nj , by A2, we have the following series expanded around ξ = 0 :
[Equation]
Taking a conditional expectation gives:
[Equation]
We must show the remainder term
[Equation]
goes to zero as c h → 0 .
This is true provided each moment E (( ξ Nj ) m | x , y ) is sufﬁciently well-behaved.
Using
Jensen’s inequality and the triangle inequality in the ﬁrst line, we have that
[Equation]
[monotone convergence]
[Equation]
[Equation]
With this in place, we have that the problem (4) is close to a linear least squares problem, since
[Equation]
with residual η N = ˆ λ N − E (ˆ λ N | x , y ) .
The residual satisﬁes
[Equation]
This follows since e N +1 is deﬁned in relation to the baseline loss, not the stochastic loss, meaning it is measurable with respect to ( x , y ) and can be moved into the conditional expectation.
From (7) and A3, we have that the least squares estimator (5) satisﬁes
[Equation]
Thus, using the continuous mapping theorem
[Equation]
A4 and Eq.
[Equation]
[Equation]
Then we have:
[Equation]
We can use Theorem 1 to establish convergence over the rest of the layers of the network when the activation function is the identity.
Theorem 2.
Assume A1-4.
For g FA ( h i , ˜
e i +1 ; B i +1 )
=
B i +1 ˜
e i +1 and σ ( x ) = x , the least squares estimator
[Equation]
solves (4) and converges to the true feedback matrix, in the sense that:
[Equation]
Proof.
Deﬁne
[Equation]
assuming this limit exists.
From Theorem 1 the top layer estimate ˆ B N +1 converges in probability to ˜ W N +1 ( c ) .
We can then use induction to establish that ˆ B j in the remaining layers also converges in probability to ˜ W j ( c ) .
That is, assume that ˆ B j converge in probability to ˜ W j ( c ) in higher layers N + 1 ≥ j > n .
Then we must establish that ˆ B n also converges in probability.
To proceed it is useful to also deﬁne
[Equation]
as the error signal backpropagated through the converged (but biased) weight matrices ˜ W ( c ) .
Again it is true that ˜˜ e N +1 =
e N +1 .
As in Theorem 1, the least squares estimator has the form:
[Equation]
Thus, again by the continuous mapping theorem:
[Equation]
In this case continuity again allows us to separate convergence of each term in the product: plim
[Equation]
using the weak law of large numbers in the ﬁrst term, and the induction assumption for the remaining terms.
In the same way
[Equation]
Note that the induction assumption also implies lim c → 0 ˜˜ e n ( c ) =
e n .
Thus, putting it together, by A3, A4 and the same reasoning as in Theorem 1 we have the result:
[Equation]
Corollary 1.
Assume A1-4.
For g DFA ( h i , ˜ e N +1 ; B i +1 )
=
B i +1 ˜
e N +1 and σ ( x ) = x , the least squares estimator
[Equation]
solves (4) and converges to the true feedback matrix, in the sense that:
[Equation]
Proof.
For a deep linear network notice that the node perturbation estimator can be expressed as:
[Equation]
where the ﬁrst term represents the true gradient
[Equation]
given by the simple linear backpropagation
[Equation]
the second and third terms are the remainder and a noise term
[Equation]
as in Theorem
[Equation]
Deﬁne
[Equation]
Then following the same reasoning as the proof of Theorem
[Equation]
we have:
[Equation]
Then we have:
[Equation]
Discussion of assumptions
It is worth making the following points on each of the assumptions:
• A1.
In the paper we assume ξ is Gaussian.
Here we prove the more general result of convergence for any subgaussian random variable.
• A2.
In practice this may be a fairly restrictive assumption, since it precludes using relu non-linearities.
Other common choices, such as hyperbolic tangent and sigmoid non-linearities with an analytic cost function do satisfy this assumption, however.
• A3.
It is hard
to establish general conditions under which ˜ e n (˜ e n ) T will be full rank.
While it may be a reasonable assumption in some cases.
Extensions of Theorem 2 to a non-linear network may be possible.
However, the method of proof used here is not immediately applicable because the continuous mapping theorem can not be applied in such a straightforward fashion as in Equation (10).
In the non-linear case the resulting sums over all observations are neither independent or identically distributed, which makes applying any law of large numbers complicated.