Predictive Coding Approximates Backprop along Arbitrary Computation Graphs

1 INTRODUCTION . Deep learning has seen stunning successes in the last decade in computer vision ( Krizhevsky et al. , 2012 ; Szegedy et al. , 2015 ) , natural language processing and translation ( Vaswani et al. , 2017 ; Radford et al. , 2019 ; Kaplan et al. , 2020 ) , and computer game playing ( Mnih et al. , 2015 ; Silver et al. , 2017 ; Schrittwieser et al. , 2019 ; Vinyals et al. , 2019 ) . While there is a great variety of architectures and models , they are all trained by gradient descent using gradients computed by automatic differentiation ( AD ) . The key insight of AD is that it suffices to define a forward model which maps inputs to predictions according to some parameters . Then , using the chain rule of calculus , it is possible , as long as every operation of the forward model is differentiable , to differentiate back through the computation graph of the model so as to compute the sensitivity of every parameter in the model to the error at the output , and thus adjust every single parameter to best minimize the total loss . Early models were typically simple artificial neural networks where the computation graph is simply a composition of matrix multiplications and elementwise nonlinearities , and for which the implementation of automatic differentation has become known as ‘ backpropagation ’ ( or ’ backprop ’ ) . However , automatic differentiation allows for substantially more complicated graphs to be differentiated through , up to , and including , arbitrary programs ( Griewank et al. , 1989 ; Baydin et al. , 2017 ; Paszke et al. , 2017 ; Revels et al. , 2016 ; Innes et al. , 2019 ; Werbos , 1982 ; Rumelhart and Zipser , 1985 ; Linnainmaa , 1970 ) . In recent years this has enabled the differentiation through differential equation solvers ( Chen et al. , 2018 ; Tzen and Raginsky , 2019 ; Rackauckas et al. , 2019 ) , physics engines ( Degrave et al. , 2019 ; Heiden et al. , 2019 ) , raytracers ( Pal , 2019 ) , and planning algorithms ( Amos and Yarats , 2019 ; Okada et al. , 2017 ) . These advances allow the straightforward training of models which intrinsically embody complex processes and which can encode significantly more prior knowledge and structure about a given problem domain than previously possible . Modern deep learning has also been closely intertwined with neuroscience ( Hassabis et al. , 2017 ; Hawkins and Blakeslee , 2007 ; Richards et al. , 2019 ) . The backpropagation algorithm itself arose as a technique for training multi-layer perceptrons – simple hierarchical models of neurons inspired by the brain ( Werbos , 1982 ) . Despite this origin , and its empirical successes , a consensus has emerged that the brain can not directly implement backprop , since to do so would require biologically implausible connection rules ( Crick , 1989 ) . There are two principal problems . Firstly , backprop in the brain appears to require non-local information ( since the activity of any specific neuron affects all subsequent neurons down to the final output neuron ) . It is difficult to see how this information could be transmitted ’ backwards ’ throughout the brain with the required fidelity without precise connectivity constraints . The second problem – the ‘ weight transport problem ’ is that backprop through MLP style networks requires identical forward and backwards weights . In recent years , however , a succession of models have been introduced which claim to implement backprop in MLP-style models using only biologically plausible connectivity schemes , and Hebbian learning rules ( Liao et al. , 2016 ; Guerguiev et al. , 2017 ; Sacramento et al. , 2018 ; Bengio and Fischer , 2015 ; Bengio et al. , 2017 ; Ororbia et al. , 2020 ; Whittington and Bogacz , 2019 ) . Of particular significance is Whittington and Bogacz ( 2017 ) who show that predictive coding networks – a type of biologically plausible network which learn through a hierarchical process of prediction error minimization – are mathematically equivalent to backprop in MLP models . In this paper we extend this work , showing that predictive coding can not only approximate backprop in MLPs , but can approximate automatic differentiation along arbitrary computation graphs . This means that in theory there exist potentially biologically plausible algorithms for differentiating through arbitrary programs , utilizing only local connectivity . Moreover , in a class of models which we call parameter-linear , which includes many current machine learning models , the required update rules are Hebbian , raising the possibility that a wide range of current machine learning architectures may be faithfully implemented in the brain , or in neuromorphic hardware . In this paper we provide two main contributions . ( i ) We show that predictive coding converges to automatic differentiation across arbitrary computation graphs . ( ii ) We showcase this result by implementing three core machine learning architectures ( CNNs , RNNs , and LSTMs ) in a predictive coding framework which utilises only local learning rules and mostly Hebbian plasticity . 2 PREDICTIVE CODING ON ARBITRARY COMPUTATION GRAPHS . Predictive coding is an influential theory of cortical function in theoretical and computational neuroscience . Central to the theory is the idea that the core function of the brain is to minimize prediction errors between what is expected to happen and what actually happens . Predictive coding views the brain as composed of multiple hierarchical layers which predict the activities of the layers below . Unpredicted activity is registered as prediction error which is then transmitted upwards for a higher layer to process . Over time , synaptic connections are adjusted so that the system improves at minimizing prediction error . Predictive coding possesses a wealth of empirical support ( Friston , 2003 ; 2005 ; Bogacz , 2017 ; Whittington and Bogacz , 2019 ) and offers a single mechanism that accounts for diverse perceptual phenomena such as repetition-suppression ( Auksztulewicz and Friston , 2016 ) , endstopping ( Rao and Ballard , 1999 ) , bistable perception ( Hohwy et al. , 2008 ; Weilnhammer et al. , 2017 ) and illusory motions ( Lotter et al. , 2016 ; Watanabe et al. , 2018 ) , and even attentional modulation of neural activity ( Feldman and Friston , 2010 ; Kanai et al. , 2015 ) . Moreover , the central role of top-down predictions is consistent with the ubiquity , and importance of , top-down diffuse connections between cortical areas . Predictive coding is consistent with many known aspects of neurophysiology , and has been translated into biologically plausible process theories which define candidate cortical microcircuits which can implement the algorithm . ( Spratling , 2008 ; Bastos et al. , 2012 ; Kanai et al. , 2015 ; Shipp , 2016 ) . In previous work , predictive coding has always been conceptualised as operating on hierarchies of layers ( Bogacz , 2017 ; Whittington and Bogacz , 2017 ) . Here we present a generalized form of predictive coding applied to arbitrary computation graphs . A computation graph G = { E , V } is a directed acyclic graph ( DAG ) which can represent the computational flow of essentially any program or computable function as a composition of elementary functions . Each edge ei ∈ E of the graph corresponds to an intermediate step – the application of an elementary function – while each vertex vi ∈ V is an intermediate variable computed by applying the functions of the edges to the values of their originating vertices . In this paper , vi denotes the vector of activations within a layer and we denote the set of all vertices as { vi } . Effectively , computation flows ’ forward ’ from parent nodes to all their children through the edge functions until the leaf nodes give the final output of the program as a whole ( see Figure 1 and 2 for an example ) . Given a target T and a loss function L = g ( T , vout ) , the graph ’ s output can be evaluated and , and if every edge function is differentiable , automatic differentiation can be performed on the computation graph . Predictive coding can be derived elegantly as a variational inference algorithm under a hierarchical Gaussian generative model ( Friston , 2005 ; Buckley et al. , 2017 ) . We extend this approach to arbitrary computation graphs in a supervised setting by defining the inference problem to be solved as that of inferring the vertex value vi of each node in the graph given fixed start nodes v0 ( the data ) , and end nodes vN ( the targets ) . We define a generative model which parametrises the value of each vertex given the feedforward prediction of its parents , p ( { vi } ) = p ( v0 . . . vN ) = ∏N i p ( vi|P ( vi ) ) 1 , and a factorised , variational posterior Q ( { vi } |v0 , vN ) = Q ( v1 . . . vN−1|v0 , vN ) = ∏N i Q ( vi|P ( vi ) , C ( vi ) ) , where P ( vi ) denotes the set of parents and C ( vi ) denotes the set of children of a given node vi . From this , we can define a suitable objective functional , the variational free-energy F ( VFE ) , which acts as an upper bound on the divergence between the true and variational posteriors . F = KL [ ( Q ( v1 . . . vN−1|v0 , vN ) ‖p ( v0 . . . vN ) ] ≥ KL [ ( Q ( v1 . . . vN−1 ) |v0 , vN ) ‖p ( v1 . . . vN−1|v0 , vN ) ] ≈ N∑ i=0 Ti i ( 1 ) Under Gaussian assumptions for the generative model p ( { vi } ) = ∏N i N ( vi ; v̂i , Σi ) , and the vari- ational posterior Q ( { vi } ) = ∏N i N ( vi ) , where the ‘ predictions ’ v̂i = f ( P ( vi ) ; θi ) are defined as the feedforward value of the vertex produced by running the graph forward , and all the precisions , or inverse variances , Σ−1i are fixed at the identity , we can write F as simply a sum of prediction errors ( see Appendix D or ( Friston , 2003 ; Bogacz , 2017 ; Buckley et al. , 2017 ) for full derivations ) , with the prediction errors defined as i = vi − v̂i . These prediction errors play a core role in the framework and , in the biological process theories ( Friston , 2005 ; Bastos et al. , 2012 ) , are generally considered to be represented by a distinct population of ‘ error units ’ . Since F is an upper bound on the divergence between true and approximate posteriors , by minimizing F , we reduce this divergence , thus improving the quality of the variational posterior and approximating exact Bayesian inference . Predictive coding minimizes F by employing the Cauchy method of steepest descent to set the 1This includes the prior p ( v0 ) , which simply has no parents . dynamics of the vertex variables vi as a gradient descent directly on F ( Bogacz , 2017 ) . dvi dt = ∂F ∂vi = i − ∑ j∈C ( vi ) j ∂v̂j ∂vi ( 2 ) The dynamics of the parameters of the edge functions θ such that v̂i = f ( P ( vi ) ; θ ) , can also be derived as a gradient descent on F . Importantly these dynamics require only information ( the current vertex value , prediction error , and prediction errors of child vertices ) locally available at the vertex . dθi dt = ∂F ∂θi = i ∂v̂i ∂θi ( 3 ) To run generalized predictive coding in practice on a given computation graph G = { E , V } , we augment the graph with error units ∈ E to obtain an augumented computation graph G̃ = { E , V , E } . The predictive coding algorithm then operates in two phases – a feedforward sweep and a backwards iteration phase . In the feedforward sweep , the augmented computation graph is run forward to obtain the set of predictions { v̂i } , and prediction errors { i } = { vi − v̂i } for every vertex . Following Whittington and Bogacz ( 2017 ) , to achieve exact equivalence with the backprop gradients computed on the original computation graph , we initialize vi = v̂i in the initial feedforward sweep so that the output error computed by the predictive coding network and the original graph are identical . In the backwards iteration phase , the vertex activities { vi } and prediction errors { i } are updated with Equation 2 for all vertices in parallel until the vertex values converge to a minimum of F . After convergence the parameters are updated according to Equation 3 . Note we also assume , following Whittington and Bogacz ( 2017 ) , that the predictions at each layer are fixed at the values assigned during the feedforward pass throughout the optimisation of the vs. We call this the fixed-prediction assumption . In effect , by removing the coupling between the vertex activities of the parents and the prediction at the child , this assumption separates the global optimisation problem into a local one for each vertex . We implement these dynamics with a simple forward Euler integration scheme so that the update rule for the vertices became vt+1i ← vti−η dFdvti where η is the step-size parameter . Importantly , if the edge function linearly combines the activities and the parameters followed by an elementwise nonlinearity – a condition which we call ‘ parameter-linear ’ – then both the update rule for the vertices ( Equation 2 ) and the parameters ( Equation 3 ) become Hebbian . Specifically , the update rules for the vertices and weights become dvidt = i − ∑ j jf ′ ( θj v̂j ) θ T j and dθi dt = if ′ ( θiv̂i ) v̂i T , respectively .

Summary: The paper extends prior work on equivalence between predictive coding and backprop in layered neural networks to arbitrary computation graphs. This is empirically tested first on a simple nonlinear scalar function, and then on a few commonly used architectures (CNNs, RNNs, LSTMs), confirming the theoretical results. The importance of this advance is highlighted by noting that the demonstrated equivalence shows how in principle modern architectures could be implemented in biological neural systems, and that the highly parallel nature of predictive coding could lead to efficient implementations in neuromorphic hardware.