Cortico-cerebellar networks as decoupling neural interfaces

1 Introduction . Efficient credit assignment in the brain is a critical part of learning . However , how the brain solves the credit assignment problem remains a mystery . One of the central issues of credit assignment across multiple stages of processing is the need to wait for later stages to finish their computation before learning can take place . In deep artificial neural networks , where processing is divided into a forward ( i.e . prediction ) and feedback ( i.e . error backpropagation ) phase , these constraints are explicit [ 1–5 ] . Given hidden activity h at a certain layer or timestep , network computations ∗Equal contributions 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . at subsequent layers/timesteps must first be computed before feedback gradients δh are available for network parameters to be updated . Not only does this introduce computational inefficiency , in particular for the temporal case where backpropagation through several timesteps is required , but it is also considered biologically unrealistic [ 6 ] . Recently , a framework was introduced to decouple forward and feedback processing in artificial neural networks – decoupled neural interfaces ( DNI ; Jaderberg et al . [ 5 ] ) 2 – in which a connected but distinct neural network provides the main network with predicted forward activity or feedback gradients , thereby alleviating the locking problems . Here , we propose that a specialised brain area , the cerebellum , performs a similar role in the brain . In the classical view the cerebellum is key for fine motor control and learning by constructing internal models of behaviour [ 7–11 ] . More recently , however , the idea that the cerebellum is also involved in cognition has gained significant traction [ 12–14 ] . An increasing body of behavioural , anatomical and imaging studies points to a role of the cerebellum in cognition in both human and non-human primates [ 12 , 14–16 ] . In particular cerebellar impairments have been observed in language [ 15 ] , working memory [ 17 ] , planning [ 18 ] , and others modalities [ 19 ] . Coupled with its notoriously uniform structure , these observations suggests that the cerebellum implements a universal function across the brain [ 7–9 , 20 ] . Moreover , experimental studies looking at cortico-cerebellar interactions have directly demonstrated that cerebellar output is crucial for the development and maintenance of neocortical states [ 17 , 21 , 22 ] . However , to the best of our knowledge , no computational formulation exists which explicitly describes the function of such interactions between the cerebellum and cortical areas . With this in mind and inspired by deep learning DNIs we introduce a systems-level model of corticocerebellar loops . In this model and consistent with the cerebellar universal role , the cerebellum serves to break the spatio-temporal locks inherent to both feedforward and feedback information processing in the brain , akin to DNI . In particular , we posit that the cerebellar forward internal model hypothesis of the cerebellum is equivalent to DNI-mediated unlocking of feedback gradients . Following this view the cerebellum not only provides motor or sensory feedback estimates , but also any other modality encoded by a particular brain region . In this regard we introduce a cortico-cerebellar RNN ( ccRNN ) model which we test on sensorimotor tasks : ( i ) a simple line drawing task [ 23–25 ] and ( ii ) drawing tasks with temporally complex input based on the MNIST dataset , but also ( iii ) on a cognitive task – caption generation [ 15 ] . Our results support the decoupling view of cortico-cerebellar networks by showing that they improve learning while reducing ataxia-like behaviours in a range of tasks , being qualitatively consistent with a wide range of experimental observations [ 15 , 23–25 ] . 2 Cerebellum as a cortical feedback prediction machine . We first describe DNIs following Jaderberg et al . [ 5 ] and then establish the link to cortico-cerebellar networks . We focus on “ backward ” ( feedback ) DNI which directly facilitates learning and is the model considered in this paper , but we also consider “ forward ” DNI to be strongly relevant to cerebellar computation ( see SM A ) . Although in this paper we focus on the DNI in a temporal setting , to better explain DNI we first briefly explain the spatial case . This case not only serves as a general basis for the temporal case ( i.e . one can think of a recurrent neural network as a specific instance of a feedforward network ) , but also highlights the possibility to place cerebellar modules between different pathways . Assume that a feedforward neural network consists of N layers , with the ith layer ( 1 ≤ i ≤ N ) performing a “ computational step ” fi with parameters θi . Given input x at layer 1 , the output of the network at its final layer is therefore given by fN ( fN−1 ( . . . f2 ( f1 ( x ) ) . . . ) ) . Let Fji denote the composition of steps from layer i to layer j ( inclusively ) . Finally , let hi denote the ( hidden ) activity at layer i , so that hi = fi ( hi−1 ) with h0 = x . We now illustrate the learning constraints of standard artificial neural networks used in deep learning . Suppose that a network is in its feedback ( or backpropagation ) phase , with current input-target pair ( x , y ) . To update the layer parameters θi the gradient dLdθi is required , where L = L ( y , ŷ ) is the loss which compares the target value against the model output ŷ = FN1 ( x ) under some loss function L ; we then apply gradient descent on the parameters , θi ← θi−α dLdθi , with learning rate α > 0 . Suppose however that the network has only recently received the input x and is currently only at layer i of 2DNIs are related to earlier work on using network critics to train neural networks [ 2 ] . the forward computation . In order to update the corresponding parameters of that layer , θi , the layer must first wait for all remaining layers to finish fj ( j > i ) for the loss to be computed . Only then are the various gradients of the loss are backpropagated and dLdθi is finally available . The enforced wait makes layer i “ feedback locked ” 3 to FNi+1 and imposes an complete dependence between learning of the layer in question and the computations in the remaining network . The basis of DNI is to break the feedback lock by sending a copy of the hidden layer i activity hi to a separate neural network C , termed “ synthesiser ” in [ 5 ] and which we propose as being implemented by the cerebellum in the brain , which then returns a synthetic gradient δ̂i – an estimate of the real loss gradient with respect to the layer , i.e . Ci ( hi ) = δ̂i ≈ δi where δi : = dLdhi . We can then update our original layer immediately according to gradient descent , effectively breaking the feedback lock with θi ← θi − αC δ̂i dhi dθi ; δ̂i = Ci ( hi ) ( 1 ) How can we trust the synthetic gradient to improve performance ? The parameters of C are themselves learnt so as to best approximate the observed feedback gradient . That is , once the standard forward and feedback computations are performed in the rest of the network ( which may take some time ) and dL dhi is available , we minimise the objective function LC = ||δ̂i − dLdhi || . Assuming C learns a decent approximation , the original layer parameters should then move ( roughly ) along the direction of the true loss gradient . Consistent with the need for this approximation a model with a fixed cerebellum C ( i.e . not learnt ) gives poor or often no learning , indicating that C output is not simply facilitating learning via a stabilising effect ( Fig . S7 ) . Temporal feedback in RNNs Using the same terminology but replacing layer indices i with time t and setting θt = θ ∀t , we can apply the above to recurrent neural networks ( RNNs ; Fig . 1a ) . Now the synthesiser C predicts the effect of the current hidden activity ht to future losses Lτ > t = � τ > t Lτ , effectively providing the RNN with ( approximate ) temporal feedback before the sequence may even be finished . As an optimization algorithm to obtain the true temporal feedback δt we use backpropagation through time ( BPTT ) . In principle , this feedback is only available once all the future feedback is available and C itself would be feedback locked . To circumvent this issue C ( hi ) learns using a mixture of nearby true feedback signals and a bootstrapped ( self-estimation ) term ( Fig . 1b ) , where for each t we now minimise ||δ̂t − δ̄t|| with δ̄t = �T τ > t dLτ dht + C ( hT ) dhTdht [ 5 ] . Here the first term includes the nearby feedback signals backpropogated within some limited time horizon T , and the second term includes the bootstrapped cerebellar prediction which estimates gradient information beyond this horizon ( analogous to temporal difference algorithms used to estimate value functions in reinforcement learning ) . In this case the cerebellum feedback facilitates learning by enabling the RNN to learn using ( predicted ) future feedback signals that would otherwise be inaccessible , enabling a future-aware online learning setup . In particular , this model has already been demonstrated to thrive in the harsh but more biologically reasonable setting of reduced temporal windows for BPTT ( small T ; cf [ 5 ] ) . 2.1 Mapping between cortico-cerebellar networks and decoupled neural interfaces . Building on the temporal feedback DNI we introduce a systems model for cortico-cerebellar computation , which we denote as ccRNN ( cortico-cerebellar recurrent neural network ; Fig . 1 ) . The model includes two principal components : an RNN which models a recurrent cortical network ( or area , e.g . motor cortex or prefrontal cortex ) , and a feedforward neural network - the synthesiser C - which receives RNN activity before returning predicted feedback δ̂t and which we interpret as the cerebellum . The cerebellum receives the current cortical activity ht through its mossy fibres which then projects onto granule cells and Purkinje cells and sends back the predicted cortical feedback δ̂t . At the architectural level our model is in keeping with the observed anatomy of the cerebral-cerebellar circuitry . Recurrent connections are widespread in the neocortex and are considered functionally important for temporal tasks [ 26 ] , whereas the cerebellum exhibits a relatively simple feedforward architecture [ 7–9 ] . Moreover , an important defining feature of the cerebellum is its expansion at the granular layer : with > 50 billion granule cells ( more than the rest of the brain ’ s neurons combined ) 3We use feedback locked networks to incorporate both the “ update ” and “ backward ” lock as described in [ 5 ] . there is a considerable divergence from the ∼ 250 million input mossy fibres . To replicate this , in our experiments we include a high divergence ratio ( 1 : 4-5 ) between the RNN network size and the hidden layer of the cerebellar network ( see SM for network sizes ) . In our experiments this ratio helped the cerebellar module learn more quickly ( not shown ) . At the functional level , our model suggests that the cerebellum should encode a variable dependent on error signals ( in particular , the error gradient with respect to neocortical activity ) , consistent with neural recordings showing cerebellar-encoding of kinematic errors during motor learning [ 27–29 ] , but also reward prediction signals in cognitive tasks [ 30 ] . Finally , we conform to the classical view that the inferior olive would be the site at which cerebellar prediction and feedback are compared to compute the cerebellar cost LC which drives cerebellar learning [ 7–9 , 31 ] . Crucially , we also highlight the cerebellum ’ s ability to learn in the case of delayed feedback - which is necessarily the case under this model - via its specially adapted timing rules for plasticity at the parallel fibre synapse [ 32 ] .

Summary: During a network update step all layers need to wait for the feedforward pass (or its temporal RNN analogue) to reach "the end" and only then the error is calculated, gradient computed, weights updates. So effectively most of the time the lower layers are just waiting for the feedback to trickle down. This is called "feedback locking" and is inefficient. The paper proposes to add a second network to the architecture that learn to predict intermediate layer gradient from their activations, allowing an immediate update. The overall architecture bears hypothesized similarity to the role of cerebellum in primate brain.  I find the idea interesting and the execution faithful and detailed. My main concern is whether the chosen control experiment is doing a good job at eliminating confounding factors potentially responsible for the better performance of the proposed method (see below for detailed explanations and a proposal).  I am reserving my final assessment till after authors' rebuttal.  ---- UPDATE ---- Authors' rebuttal has fully addressed my concerns and resolved them in a positive manner. I am now much more confident in the value of the contribution. Raising the score to "8: Top 50% of accepted NeurIPS papers, clear accept"