Complex priors and flexible inference in recurrent circuits with dendritic nonlinearities

Thank you for the positive overall evaluation of the work. We agree that numerically validating that the autocorrelation of the RNN-generated samples decaying to zero at small lags does not mathematically ‘prove’ that this will be the case in all conditions; the same goes when demonstrating that the additional biological constraints do not harm sampling performance. We can still say that –empirically– it is the case across the distributions that we have tried (including new results on MNIST), both for the prior and the posterior samples.

Regarding the role of attractor dynamics in correcting for the approximations, one way to think about the approximations introduced in the service of biological plausibility is as inducing small noisy deviations from what would have been the ideal trajectory, defined by the solution without the approximations. Since in our model the representation of distributions in the embedding manifolds is driven by the attractor flow fields, adding moment by moment stochasticity on top of such dynamics can be partially corrected for by the push towards the manifold. The dynamics will not end up exactly at the same location on the manifold compared to the exact version, but they will still sample from the right distribution. About the use of toy datasets: see common reply. Briefly, we now numerically demonstrate the scalability of our approach using MNIST.

We have corrected the minor typos and so on (see updated pdf). Questions:

• Multi-compartment neuron should be plural (typo).
• We used the phrase “optimized ReLU nonlinearities” to mean that the dendritic nonlinearities have been optimized on the denoising objective defined by Eq. (4). These nonlinearities are functionally instantiated by linear weight matrices and ReLU activations. We have modified the main text for clarity.
• It is indeed the case that $\sigma_\lambda$ in Eq 1 is the same as $\beta_\lambda$ in Fig. 1. To keep our notation consistent, we changed Fig 1 to use $\sigma_\lambda$ instead.
• This sentence relates to our approximation of the likelihood, which consists of evaluating $g_t$ at $x_{t-1}$ rather than at $\mu_\theta(x_{t-1})$. There is a substantial penalty associated with this approximation, since $x_{t-1}$ is the noisy sample at time $t-1$, while $\mu_\theta(x_{t-1})$ is the network’s (denoised) estimate of the clean sample at time $t$, which lies closer to the manifold of clean images. Since this approximation error is incurred at every timestep, it is possible for these errors to accumulate over time and lead to large sampling biases in the resulting posterior distribution. Instead, we find in our numerical simulations that this is not the case; rather, the discrepancy (measured using the KL divergence) tends to decrease over the course of the inference process.
• Yes, this is correct. We have fixed the typo to include the time step in $\mathbf{g}$.
• The last sentence of the caption of Figure 3 does indeed correspond to Figure 3D. We have fixed the caption.

We thank the reviewer for recognizing the originality of our work and we apologize for the lack of clarity in the description of the model.

Dendritic nonlinearities: The collection of single neuron nonlinearities (which have a shared parametric form, but with neuron-specific parameters), are the direct result of the denoising optimization procedure, and their operations map one to one onto the reverse transition operator $\mu_\theta(x_{\lambda+1}, \lambda)$ in Eq. 2. This indeed implies that the dendritic nonlinearity parameters would change when modeling a different prior distribution (biologically achievable through a combination of synaptic and branch plasticity, although we do not model the mechanics of learning here).

Global oscillations: there is a one-to-one map between the oscillation phase in our model and the noise schedule in the traditional diffusion models at sampling time. In more detail: in diffusion models, the degree of noise at step $\lambda$ is given by a noise schedule $\sigma_\lambda$ that decreases monotonically over the course of inference as a noisy image $x_\lambda$ is iteratively denoised. In other words, images with the largest degree of noise are found at $\lambda=\Lambda$ and clean images are produced at $\lambda=0$. While the exact form of the schedule can matter in the details, there is a range of possible choices available. In the interest of biology we used a sinusoidal interpolation between $\lambda=\Lambda$ and $\lambda=0$. This allows us to establish a direct map between $\Lambda$ and the ascending phase of a circuit oscillation (see Eq.5), such that images with the largest and smallest degrees of noise correspond to the trough and peak of the oscillation, respectively (this is a somewhat arbitrary choice, motivated by some of the past literature). There is no direct correspondence to the descending phase of the $\beta_t$ oscillation in diffusion models, which is a unique element of our solution. The traditional forward process would provably achieve the correct outcome computationally (since its entire purpose is to map the actual input distribution into a gaussian), something which is not formally guaranteed for the oscillation. However, that solution would not match the biological constraint that the neural dynamics driving the evolution of network activity cannot dramatically change every few ms or tens of ms, whereas we can use the same RNN dynamics with increasing levels of noise and keep things consistent.

Training an ANN-based model for dendrites: the biological implementation of this process is beyond the scope of this paper, but there are some technical results that suggest that one can define a local denoising objective that does not require knowledge of the clean image, but only of the the general noise statistics used to corrupt it (Raphan and Simoncelli, 2014), so it is possible in principle to construct a denoising optimization procedure that involves relatively local computations. To which extent those operations map to experimental observations on synaptic and branch plasticity remains to be determined.

Question answers: -in Fig.1E The only difference between the ascending and descending phases is whether the degree of noise that is present in the neural transition operator is decreasing or increasing. The increasing level of noise in the dynamics effectively drives the activity pushing it increasingly off manifold. The statistics of exactly where one ends up are not provably identical to those imposed by the traditional feedforward process, but when interacting with the attractor dynamics in the ascending oscillation phase it seems to be close enough to effectively decorrelate samples across cycles.

• sorry for the misdirection, we have now corrected the caption to top/bottom
• we made the choice of indexing the network by time but the nature of the operations really depends by the phase of the oscillation, which maps one to one to the time step within a cycle in regular diffusion. We were hoping that it makes the correspondence between what happens algorithmically and what might happen in the brain clearer, but it does add extra notation, perhaps unnecessarily.

We thank the reviewer for the positive comments.

It may be worth noting that our sampling dynamics are qualitatively quite different from Langevin, even though it is indeed the case that our work builds upon a body of work on MCMC approaches for neural sampling.

Providing direct evidence for individual elements of our solution is indeed difficult, but not outside the scope of future experiments. There is ample indirect evidence that similar kinds of computations are possible in the brain: certain interneurons types have both oscillatory responses and target principal cells’ dendritic arbors, which means that they can modulate the effective neuron nonlinearities in a cyclic fashion. Similarly, the phase of the local oscillation can affect somatic subthreshold variability and change the noisiness of the generated neuron outputs. Finally, there is ample evidence (at least from within the hippocampus and for hippocampal-cortical interactions) that across area communication happens preferentially at certain phases of a local oscillation, and Communication through coherence theories have made similar arguments for the cortex. So, we would argue that all of the individual elements of the model have at least loose experimental support (as in something of that flavor is known to happen, even if it cannot be matched in the details). That said, we do identify one very concrete analysis that would prove or disprove the model in the paper, in the form of activity pattern similarities across different phases of an oscillatory cycle. There is experimental proof of feasibility of such an analysis (used for a slightly different purpose) in the work of Berkes et al.

About the use of toy example models: see common reply. Briefly, the low dimensionality of the distributions considered is a convenient choice rather than a limitation of the model; we now include a MNIST experiment representing actual images to illustrate that point.

Thank you for the feedback. About the general biological plausibility of our scheme (see also reply to reviewer XKTM):

Oscillation: It is fundamentally in the nature of modeling to abstract away some of the inessential details. Our oscillatory signal is such an abstraction. In new results we now show numerically that the core mechanism remains largely unaffected when using an imperfect oscillation, which varies in amplitude and period from cycle to cycle by a significant margin (Suppl. Section B5).

Learning: Learning probabilistic representations is a huge open question in the field. Virtually none of the published circuit level models provide a learning procedure that can construct them. This is true about models of priors (Ganguli and Simoncelli, 2014, although something similar can be learned with a REINFORCE type of rule, Bredenberg et al, 2020) and -to our knowledge- all neural sampling literature. That said, denoising as a learning objective is more local and in some ways mathematically simpler, which makes us hopeful about our ability to map it into plasticity-based biological learning. Indeed, there are some technical results that suggest that one can learn a Bayesian estimator without requiring knowledge of the distribution of clean data, but only of the the general noise statistics used to corrupt it (Raphan and Simoncelli, 2007), which suggests a mechanism by which biological systems may learn to denoise without having ever seen clean data.

About task complexity: see common reply. Briefly, the distributions we consider are actually quite complex in light of the previous neural sampling literature (complexity here is defined as nonlinear low-d manifolds embedded in high-d ambient spaces and strong multimodality). None of the previously published neural sampling models would be able to represent or sample from our ‘simple’ toy examples, so this is already a qualitative improvement over the comp. neuro. SOTA. Further, we now include new experiments with MNIST to show that the model can be scaled up without problems (Suppl.Section B6).

A more detailed note on the limits of distributional complexity that can be in principle represented: there is a natural question about “how complex can it possibly get?” Based on the machine learning literature on diffusion models the answer is at least pixel level natural images, although the number of parameters and complexity of dendritic nonlinearities involved in a naive implementation of that is unrealistically high. In subsequent work we are exploring how hierarchical versions of the same idea would scale up prior complexity while keeping the neural resources required for its representation relatively small.

Detailed question answers:

Recurrence: indeed this is a network of all-to-all connected units, each of which involves a neuron specific nonlinearity, with parameters trained by denoising. The number of neurons is the same as the ambient space (10). The somatic compartment has a somewhat privileged role in that it introduces stochasticity in the neuron responses (although we have tried a version where noise was injected along the dendritic tree, with no particular ill effect). These details are now clearly explained in the updated manuscript.

Autocorrelation function of posterior: Yes, the samples still show zero autocorrelation for posterior inference. The results are shown in supplement, in the section titled B.3 Posterior Autocorrelation.

Flexible inference in other schemes: In principle, a linear PPC circuit that separately encodes the log prior can combine its representation additively with incoming sensory/contextual log likelihoods to perform inference; flexible inference would require a different kind of gating between subpopulations, but it can in principle be done, although to our knowledge it hasn’t been demonstrated explicitly. Nonetheless, this only works for one-dimensional or fully factorized joint distributions. The type of priors the brain needs to represent are arguably more complex which has motivated probabilistic circuits based on sampling. We now provide a way to achieve flexible inference for such complex distributions.

Should the top-down contribution be called ‘prior’ or ‘likelihood’: We agree that our choice of wording is (although formally correct) a little inconsistent with some of the neural literature. We are thinking about these qualifiers from the perspective of an intermediate stage of a hierarchical inference process, and so the definition of the role of top-down information is with regards to the locally represented latents.