Structural Credit Assignment in Neural Networks using Reinforcement Learning

1 Introduction . Training neural networks involves structural credit assignment : attributing credit ( or blame ) to nodes in the network for correct ( or incorrect ) predictions . The output from a node early in the network impacts all the outputs downstream and finally the prediction outputted at the end of the network . Our goal is to adjust the weights that produced the output for this node , so that the prediction would have been more accurate . The most widely used solution for the structural credit assignment problem is backpropagation [ 44 ] , namely gradient descent on the loss for the outputs . Moving beyond backprop provides more flexibility in training neural networks . Backprop requires differentiability of activations and losses for the network , as well as synchronicity for computing the gradient and updating the weights . To update a node internal to the network , a full feedforward and backward pass needs to be computed , with global gradient information sweeping backwards from the output . Ideally , for online agents operating real-time , with computational constraints , we would have nodes that update each step , locally and asynchronously . To make progress towards this lofty goal , we revisit an old idea : treating each node as an agent . Work in reinforcement learning ( RL ) , including ideas like eligibility traces , were in fact inspired by Klopf [ 23 ] and the hedonistic neuron . It is not surprising that the idea of using an RL agent for each node is 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . found in early work , including the original REINFORCE algorithm [ 61 ] , which is a policy gradient approach using sampled returns . Most work treating each node as an agent uses the REINFORCE update , often with baselines for variance reduction , including work on learning with spiking neurons [ 14 ] and CoAgent Networks ( CoANs ) [ 56 ; 55 ; 24 ] . The work in CoANs 1 ) nicely formalizes the idea of a collection of agents—each agent corresponding to a node or subset of nodes—cooperating to maximize return and 2 ) provides a general theorem on the validity of using the REINFORCE update . For this reason , we adopt their terminology and use CoANs to refer to networks composed of agents . More recently , other algorithmic ideas from reinforcement learning , beyond REINFORCE , have begun to affect training of ( stochastic ) neural networks . The ideas of critics and baselines , which reduce the variance of policy gradient updates , have been well-developed for stochastic computation graphs [ 60 ] . This work provides a unification of gradient derivations , but as yet not an investigation into practical algorithms for structural credit assignment in neural networks . Other work on learning under stochastic neurons has typically used REINFORCE as a basic method , and explored other heuristics to improve learning , such as straight-through estimators [ 8 ] , rather than improved RL approaches . Other work on credit assignment is loosely inspired by the idea of bootstrapping in RL , including synthetic gradients [ 20 ; 26 ] and fixed-point propagation [ 36 ] . Overall , however , the broader space of RL algorithms has not been leveraged to learn CoANs . One reason for this omission could be that the structural credit assignment problem within the neural network has not been clearly defined as an RL problem ; rather , it was simply intuitive to use REINFORCE approaches for each node . Even the theory from the original CoANs work focused on the return in the environment—since CoANs were used to solve a reinforcement learning problem— and did not explicitly formalize the structural credit assignment problem within the network . Another reason could be that many straightforward ideas are not effective , as we show in this work . To facilitate the use of RL algorithms , we first formalize the structural credit assignment problem as a finite horizon RL problem . We show local policy gradient updates provide an unbiased estimate of the joint gradient for structural credit assignment , ensuring REINFORCE is a sound approach . We then discuss key ideas from RL—namely exploration and off-policy learning—that can be leveraged to improve learning in CoANs . We show that REINFORCE can train multi-layered networks , but faces issues with suboptimality due to coagents learning under nondeterminism of fellow coagents . We provide an in-depth study highlighting this problem and measuring the entropy of different parts of the network . This in-depth study motivates the difficulties in using the common on-policy approaches , and we discuss and show how off-policy learning is a more promising direction . Finally , we discuss the advantages of CoANs when moving away from the standard iid learning setting , showing it can perform better than backprop on a continual learning problem with a highly correlated dataset . 1.1 Other Related Work . The literature on approaches to structural credit assignment is vast , with much of it using ideas different from reinforcement learning . One category of approaches uses local updates to make activations similar to a target vector of activations , such as target propagation [ 7 ; 28 ] , the method of auxiliary variables [ 9 ] and fixed point propagation [ 36 ] . Kickback approximates the backprop update for ReLu networks , using an approximation to the gradient that allows for local updates [ 5 ] . Feedback alignment [ 31 ; 38 ] involves using random weights , instead of the actual weights in the next layer , that avoids symmetric propagation that is thought to be biologically implausible . Weight perturbation and node perturbation approaches have been used to estimate gradients , with node perturbations emerging as the preferred approach [ 46 ] . Fiete and Seung [ 14 ] showed that their node perturbation algorithm actually includes REINFORCE as a special case , linking these two classes of methods . However , the connection only exists for Gaussian noise perturbations and REINFORCE ; the connection is lost for other perturbations as well as for other RL algorithms . Other work has focused on changes in weights across time . Spiking neural networks and the associated spike-timing-dependent plasticity ( STDP ) learning rule [ 32 ] adjusts weights based on the relative timing of activations for nearby nodes . Equilibrium propagation [ 47 ] involves a phase of propagation in the network until reaching a low energy state , followed by a learning update . This work showed similarities to STDP , Contrastive Hebbian Learning [ 35 ] and Contrastive Divergence [ 18 ] . Reinforcement learning updates have been used for spiking neural networks , called Rewardmodulated STDP . These updates use a global reward but with local update rules [ 29 ; 30 ] , with some interesting insights that perturbations can be beneficial to induce exploration [ 30 ] . This area has focused on delayed reward , namely assigning credit from node changes across time and across multiple updates . The local updates use node perturbation with eligibility traces to link perturbations on this step , to accuracy ( rewards ) at later time steps [ 30 ; 34 ] . A more recent algorithm , called cross propagation [ 58 ] , explicitly adjusts weights back-in-time , to account for accuracy on this step , similarly to some meta-learning strategies but completely online . There has also been some work using REINFORCE to learn activation paths through a network [ 13 ] and learning when to activate parts of the network [ 8 ; 6 ] . Other work has connected structural and temporal credit assignment , but in the opposite direction from this work : specifying temporal credit assignment as a structural credit assignment problem [ 4 ] . Once we use RL agents as nodes , which have stochastic policies , there is a clear connection to the work on stochastic neural networks . Much of this work has looked at networks with stochastic binary activations [ 37 ; 8 ; 42 ; 33 ] , though the wider literature on stochastic computation graphs encompasses a broad range of stochastic neural networks [ 52 ; 48 ; 60 ] . Early work considered an EM-style algorithm [ 52 ] and a simple alternative , called the straight-through estimator [ 8 ] , that directly passes the gradient back through the node to the weights that created the activation . The straight-through estimator has recently been shown to be a valid estimator for stochastic binary networks [ 49 ] , and a lower-variance update has been proposed [ 17 ] . Multi-agent reinforcement learning approaches tackle a similar problem , in the cooperative setting . The coagents in the neural network can be seen as a group of agents cooperating to produce accurate predictions , though with the notable difference that there is an ordering to the actions taken by coagents . A common approach in this area has been to use reinforcement learning agents , and consider mechanisms for coordination without centralization . Some approaches have been to carefully define rewards for each agent [ 63 ; 62 ] ; to use a single global reward plus some noise [ 11 ] ; to use independent Q-learners [ 51 ] ; or to estimate a global critic [ 39 ] potentially with local policy updates [ 15 ; 43 ] . When using independent Q-learners for each agent , the other agents are treated as part of the environment ; consequently , the environment appears non-stationary . Hyper Q-learning [ 54 ] reduces the impact of this non-stationarity by estimating other agent ’ s policies . These strategies do not directly extend to CoANs , but the connection to the cooperative multi-agent reinforcement learning problem could provide fruitful avenues for improved algorithms . Finally , there have been several works empirically investigating alternative update strategies and architectures . Spiking neural networks generally have lower accuracy than standard deep neural networks , but a recent study has shown that with advances in hardware and algorithms to train spiking neural networks , this gap has become smaller [ 53 ] . Stochastic neural networks , particularly with binary activations , have been shown to be difficult to train , but with some improvements to the gradient estimator , can have significant advantages , including providing a level of regularization [ 42 ] . In this work , we aim to provide a more comprehensive empirical investigation into the use of reinforcement learning approaches to train CoANs . 2 Structural Credit Assignment as a Finite Horizon RL Problem . In this section , we describe how to formalize the credit assignment problem in a feedforward neural network as a finite-horizon RL problem . We start in the simplest setting , where we have a fully connected feedforward neural network composed of k layers of size n. We assume we have inputs x ∈ X , prediction targets y ∈ Y , hidden layer activations aj = f ( zj ) for pre-activations zj with activation function f . The basic idea for the finite horizon formulation is that the horizon is the number of layers , and on each step the input state for the agent is the activations from the previous step and the output is the activations for this layer or the final output prediction . Consider the following agent-environment interaction . On the first step , given the sampled input x ∈ X , the initial observation is o0 = [ x , 0 ] where the 0 in the observation indicates the step-index in the finite horizon problem . The agent observes input o0 and then takes action a0 ∈ A that is a vector of activations ( or pre-activations z0 ) —these actions can be discrete or continuous . The next observation is deterministically o1 = [ a0 , 1 ] ( or o1 = [ f ( a0 ) , 1 ] = [ f ( z0 ) , 1 ] ) . Then the agent inputs o1 and outputs the activations for the next layer a1 and obtains next observation o2 = [ a1 , 2 ] ( or o2 = [ f ( a1 ) , 2 ] ) . This transition is Markov , because given o1 and a1 , the next observation o2 is independent of o0 . This interaction continues for k steps , terminating at the last layer when the final action is to output the prediction ŷ . We depict this interaction in Figure 1 . For this interaction to be Markov , the computation of the final reward requires the input x . The rewards received during the episode are zero , with a reward given upon termination that is a function of the error between prediction and target , such as err ( ŷ , y ) = ( ŷ − y ) 2 with reward −err ( ŷ , y ) . The probability of y depends on x , and so the reward depends on x on this last step . We can obtain Markov states by simply including x in each observation . This part of the input can be ignored by the policy , as it is for most NN architectures . Formally , finite-horizon undiscounted MDP consists of state space S = X × { 0 } ∪ X × Rn × { 1 , . . . , k − 1 } , actions A = Rn , where each start state s0 = [ x , 0 ] and later states for j ≥ 1 are sj = [ x , aj−1 , j ] . The transitions are deterministic , with a reward of zero on each step until termination at step k giving reward rk+1 = −err ( ŷ , y ) . A typical RL agent learns a separate policy for each horizon . This corresponds to learning one stationary policy , because the step-index is included in the state . Practically , however , these ( stochastic ) policies πj ( aj |oj ) for j ∈ { 0 , . . . , k } are updated separately , without considering the single stationary policy . If the agent is using policy gradients , then it uses parameterized distributions for the policies , such as Gaussians with means parameterized by the activations from the last layer . These policies can be updated using a REINFORCE update , using the gradient of the CoAN objective . The CoAN objective corresponds to a policy gradient objective . An episode trajectory , with index information implicit and assuming actions are activations , is o0 = x , a0 , r1 = 0 , o1 = [ a0 ] , . . . , ok = [ ak−1 ] , ak = ŷ , rk+1 = −err ( ŷ , y ) , Termination . Because π ( a0 , a1 , . . . , ak|x ) = π ( a0|x ) ∏k j=1 πj ( aj |aj−1 ) , the probability of this trajectory is p ( x ) π ( a0 , a1 , . . . , ak|x ) p ( y|x ) . This is because p ( x ) gives the probability of the start state ; π ( a0 , a1 , . . . , ak|x ) gives the probability of the trajectory from s0 because the state outcomes are deterministic given the action ( activation or pre-activation ) ; and p ( y|x ) defines the probability of the reward on termination , since the target is stochastic . This provides a straightforward policy gradient objective , for undiscounted return G def= −err ( ak , y ) max π ∫ p ( x ) π ( a0 , a1 , . . . , ak|x ) p ( y|x ) G dx da0 . . . dak dy = max π Eπ , p ( x , y ) [ G ] ( 1 ) If we assume each policy πj has parameters θj , then the stochastic gradient of this objective separates out into the following stochastic gradients for each policy separately : G∇θj lnπj ( aj |oj ) . We show this formally in Proposition 1 in Appendix A . A similar result has been shown for CoANs used in the RL setting [ 24 , Theorem 1 ] and for stochastic computation graphs [ 60 , Theorem 2 ] ; we include the result specifically for this case because it avoids much of the complications from those other works . The locality of the policy gradient means policies can be updated locally with their own gradients , with a shared global return signal . We can also easily incorporate control variates to reduce the variance of the gradient , called baselines . The update is ( G− V ( oj ) ) ∇θj lnπj ( aj |oj ) , where the learned baseline V ( oj ) estimates expected return given oj : G− V ( oj ) corresponds to the advantage for the actions selected . Several different critics have been proposed for stochastic computation graphs [ 60 ] ; we discuss how to make similar baselines for this finite horizon problem in Appendix B . The CoAN is trained as a stochastic network , but is evaluated under the standard ( deterministic ) backprop setting . The question we ask is : can a CoAN , trained with local updates , produce predictions with similar accuracy to those produced by an NN , trained with backprop ? We therefore test the CoAN using the greedy actions from the co-agents , to provide deterministic predictions . Effectively , we are comparing the parameters θ produced by the CoAN to those produced by backprop.1 Remark : We described the formalism for the simpler feedforward setting , to facilitate understanding . All the ideas , however , extend to more generic acyclic networks , because every acyclic network has a topological ordering on nodes . The state input for each coagent still consists of the input nodes . At each time step , whichever coagents have all their nodes evaluated—namely have their state input available—can produce their action . This propagates forward until a prediction is produced .

Summary: The paper investigates using reinforcement learning to train neural networks. Specifically, it considers formulates the problem in which individual layers are treated as steps in a finite horizon partially observable MDP, and explores using either a global agent or individual co-operating agents at each layer to maximize the reward (the negative final error).  The paper then presents an empirical study of using REINFORCE in this cooperative agent framework to train one or two hidden layer networks on MNIST/Boston Housing while varying different parameters like the policy variance, different partitioning schemes of the network into agents, different baseline approaches to improve stability, and activations. The key takeaways of this experiment are that (1) the CoAN+REINFORCE approach generally underperforms backprop, (2) baselines to improve stability don't seem to help, and (3) the primary reason for the poor performance is the non-deterministic behavior of the other agents (as indicated by having a single agent optimizing all layers performs the best by far).   Motivated by this, the paper presents an off-policy Q learning approach to training cooperative agents, which uses bootstrapping and trains by assuming the best action is taken at the next step. While the paper shows this improves over on-policy SARSA, it still underperforms REINFORCE, which underperforms backprop.   Lastly, the paper shows that in a continual learning setting the CoAN+REINFORCE approach can be more robust than backprop.