Paper: Complex priors and flexible inference in recurrent circuits with dendritic nonlinearities (ICLR-2024)

Reviewer comment (Experiments):
Regarding the sentence in page 5: "Overall, these results indicate that the constraints imposed by biology may have a minimal effect on the quality ... compared to DMs". Isn't it too early to state such conclusion. The experiment is on a quite toy-dataset.

Author response:
About the use of toy datasets: see common reply. Briefly, we now numerically demonstrate the scalability of our approach using MNIST.

Reviewer comment (Writing):
In the second line of Section 2, "nosy" should be corrected to "noisy."

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

Reviewer comment (Writing):
In Figure 1D, the caption mentions "multi-compartment neuron." Shouldn't it be "multi-compartment neurons"? I thought each gray triangle is a multi-compartment neuron.

Author response:
Multi-compartment neuron should be plural (typo).

Reviewer comment (Writing):
The paper repeatedly mentions "optimized ReLU nonlinearities." I did not really understand the meaning of "optimized" in them?

Author response:
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.

Reviewer comment (Writing):
Is $\sigma_\lambda$ in Equation 1 the same with $\beta_\lambda$ in Figure 1?

Author response:
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.

Reviewer comment (Theory):
Regarding the following sentence in page 6: "In principle, this might lead to catastrophic accumulation of errors and large sampling biases; however, the attractor dynamics prevent this from happening." Can authors elaborate on this a bit more? It is not clear to me why?

Author response:
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.

Reviewer comment (Theory):
In page 6, it is written "$\mathbf{g}\_t (\mathbf{x}\_{t-1}, s) = \frac{1}{\sigma_s^2} \mathbf{M}_s \mathbf{M}\_s^T (\mathbf{x} - \mathbf{x}\_c)$". Is it $\mathbf{g}\_t (\mathbf{x}\_{t-1}, s) = \frac{1}{\sigma\_s^2} \mathbf{M}\_s \mathbf{M}\_s^T (\mathbf{x}\_{t-1} - \mathbf{x}\_c)$?

Author response:
Yes, this is correct. We have fixed the typo to include the time step in $\mathbf{g}$.

Reviewer comment (Theory):
I have major concerns about the justification and derivation of two central claims (dendritic nonlinear and global oscillating signal) about neural circuit implementation of diffusive models. A possibility is that the text doesn't explain the key math steps sufficiently well. I look forward to seeing some justification in the rebuttal.

Author response:
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.

Reviewer comment (Theory):
The paper directly states the nonlinear dendritic operation \( f(x_t, \beta_t) \) after Eq. 5 without explaining how it is derived. I am confused about how this nonlinearity comes out. Is it directly from \( \mu_\theta (x_{\lambda + t}, \lambda) \) in Eq. 2? In this case, does it imply the dendritic nonlinearity needs to be readjusted if the transition operator in the reverse process is changing? Or is it just used as a way to capture the nonlinearity in biological neurons?

Author response:
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).

Reviewer comment (Theory):
I don't understand how the global oscillating signal is derived. Although the author explained the \( \beta_t \) is analogous to the sequence of noise variance in the diffusive model, it seems that the diffusive model doesn't have such an oscillating signal if I understood correctly. I have no idea how Eq. 5 was derived, what assumptions it relies on, and why the \( \beta_t \) becomes a sinusoid function there.

Author response:
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.

Reviewer comment (Theory):
The author says the dendrite is modeled as an ANN whose parameters were trained via gradient descent. Does this imply some mechanism to adjust the dendritic parameters in real neurons? If so, what are the possible biological mechanisms? I don't see related discussions in the paper.

Author response:
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.

Reviewer comment (Theory):
Overall, I suggest the author explain the key motivations, and assumptions in deriving the nonlinear dendrite and global oscillating signals in modulating dendritic noises around Eq. 5. They seem quite critical in this work. In addition, it is better to write a concrete, commonly used recurrent neural dynamics in a single line, and compare it directly with the reverse process in diffusive models. If my major concerns are solved, I'd like to increase my rating.

Author response:
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.

Reviewer comment (Theory):
Fig. 1E caption: why does neural dynamics push the network off manifold? Does the author mean the neural dynamics of sensory transmission correspond to injecting noise in a similar fashion with the forward process in a diffusive model? At least this neural dynamics is not the same recurrent neural dynamics which the author claims to implement a reverse process to sample from posteriors.

Author response:
-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.

Reviewer comment (Presentation):
Fig. 4B caption: it should be (top) and (bottom) because there is no (left) and (right) here.

Author response:
- sorry for the misdirection, we have now corrected the caption to top/bottom

Reviewer comment (Theory):
While the model is elegant and combines a SOTA class of generative model in ML (DDPMs) with neural circuitry, the end product feels too artificial in construction and biologically implausible, with its many strong restrictions/assumptions. For one, the diffusion model in this implementation can be trained as per usual, but how would real dendrites in a neuron go through this learning? The authors state the learning aspect is for future work, but that’s a huge “if”—is it really realistic to suppose that such a complex mechanism can emerge in vivo, without any hint for the readers of how it could?

Author response:
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.

Reviewer comment (Theory):
Sampling critically depends on a stationary and permanent oscillation, which is rarely found in the brain. Even hippocampal theta, which the authors draw inspiration from, are irregular in time and frequency, nevermind oscillations in other cortical areas.

Author response:
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).

Reviewer comment (Experiments):
The performance is only demonstrated on a relatively simple task of sampling from a very smooth and low-dimensional manifold. There’s no indication that this would work with mildly more complex high-d distributions, as the authors had originally motivated (and pointed out as a weakness in previous literatures, i.e., mostly Gaussian prior/posteriors).

Author response:
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.

Reviewer comment (Experiments):
It is unclear whether the proposed approach can scale to perform a mildly more complicated task, such as conditionally sampling MNIST, or whether it would be completely infeasible given the architectural constraints.

Author response:
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.

Reviewer comment (Reproducibility):
Is there recurrent interaction between neurons? It’s also a bit unclear how many neurons there are, or are all the results from a single dendritic tree? Also, what exactly are the inputs/outputs of the dendritic networks, and what exactly is the somatic “compartment” doing, or is that functionally just the last layer of the dendrites?

Author response:
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.

Reviewer comment (Theory):
The authors motivate their proposed architecture as flexible, since it can be reused for various inference scenarios by swapping out or combining likelihoods. Wouldn’t this also be true for a model where a different population encodes the prior as well? Why is it necessary that it’s in the dendrites? Or is that just a “semantic” difference, since the branching networks can equally be implemented as different neurons?

Author response:
Detailed question answers:

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.

Reviewer comment (Theory):
Typically, bottom-up sensory info and contextual info are thought to be likelihood and prior, respectively, whereas here they are represented as two different streams of likelihoods. Can the authors comment on this discrepancy?

Author response:
Detailed question answers:

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.

Reviewer comment (Theory):
There is recent experimental evidence implicating oscillations of different frequencies coordinating bottom-up (gamma, ~40Hz) vs. top-down (beta, ~15Hz) signaling (see A. Bastos, EK Miller, etc.), while here it’s crucial that there is a global oscillation of a single frequency. Otherwise, the prior and likelihood sampling are temporally misaligned.

Author response:
Multiple interacting oscillatory bands: thank you for the references. At this stage we don’t have a computational interpretation for the role of multiple oscillatory bands in interarea information exchange. Worth thinking about.

Reviewer comment (Theory):
Most of the time, most areas of the cortex are not experiencing oscillations. Furthermore, oscillations tend to disappear during task engagement (such as 10Hz alpha in visual areas and 20Hz beta in motor areas), which would be a time that is critical for sampling. How reliable would the proposed implementation be in such scenarios?

Author response:
Modulation of slow oscillations by behavioral context: it is not clear how those empirical observations relate to what we are talking about here, but it is likely a separate process. More generally, the global coordination of information across the full visual hierarchy is very interesting, and something we are actively working on right now but it goes beyond the scope of this particular paper.

Reviewer comment (Presentation):
Caption for Figure 1F is missing.

Author response:
Thanks for pointing out the missing caption. Now fixed.