Expressive probabilistic sampling in recurrent neural networks

In sampling-based Bayesian models of brain function, neural activities are assumed to be samples from probability distributions that the brain uses for probabilistic computation. However, a comprehensive understanding of how mechanistic models of neural dynamics can sample from arbitrary distributions is still lacking. We use tools from functional analysis and stochastic differential equations to explore the minimum architectural requirements for $\textit{recurrent}$ neural circuits to sample from complex distributions. We first consider the traditional sampling model consisting of a network of neurons whose outputs directly represent the samples (sampler-only network). We argue that synaptic current and firing-rate dynamics in the traditional model have limited capacity to sample from a complex probability distribution. We show that the firing rate dynamics of a recurrent neural circuit with a separate set of output units can sample from an arbitrary probability distribution. We call such circuits reservoir-sampler networks (RSNs). We propose an efficient training procedure based on denoising score matching that finds recurrent and output weights such that the RSN implements Langevin sampling. We empirically demonstrate our model's ability to sample from several complex data distributions using the proposed neural dynamics and discuss its applicability to developing the next generation of sampling-based brain models.

Paper

Similar papers

Peer review

Reviewer pX3M5/10 · confidence 5/52023-06-16

Summary

This paper studies circuit algorithms for sampling-based Bayesian inference in continuous-time rate-based recurrent neural networks. Specifically, it argues that using a linear readout from a noisy reservoir provides a substantial expressivity benefit relative to representing the sampled distribution directly in the recurrent population.

Strengths

1. The question of the minimal recurrent circuit architectures that allow sampling from complex distributions is clearly relevant to the neuroscience community at NeurIPS, and I think the main result is interesting (albeit unsurprising). 2. The experiments show to a reasonable degree that the proposed method can sample from non-trivial non-Gaussian distributions (MNIST). This is an improvement on most comp neuro papers on sampling, which largely focus on Gaussian distributions. My enthusiasm for this point is, however, dampened a bit by the fact that the networks here are rate-based and the authors aren't trying to optimize functionals of the dynamics (e.g. as in work from the Lengyel group), so it's not very surprising that this can be made to work.

Weaknesses

1. How does Proposition 1 differ from the results of Ma, Chen, and Fox, "A Complete Recipe for Stochastic Gradient MCMC" (NeurIPS 2015)? Perhaps I've missed something, but I do not see a novel contribution in Section 3.1. 2. The present manuscript does not address the speed of relaxation to the stationary distribution. The ability to obtain accurate samples quickly is a necessary feature of any model of sampling-based inference in the brain, and indeed the question of how to achieve fast sampling in biologically-plausible implementations has been a major focus of past work (e.g., Hennequin et al. 2014, Aitchison and Lengyel 2016, Echeveste et al. 2020, Masset et al. 2022). In my view, this is a key shortcoming of the learning algorithm proposed in Section 3.4: no attention is paid to whether the resulting network enjoys favorable convergence properties. Of course, one could modify the resulting drift and diffusion terms (as in Hennequin et al. 2014 or Masset et al. 2022), but this would require an additional optimization step. The authors should at least discuss this point; adding more experiments to probe this issue would make for a stronger paper. 3. In Lines 293-295 of the Discussion, the authors state that "Our preliminary experiments show that one-hidden-layer RSNs cannot readily approximate high-dimensional heavy-tailed distributions (e.g., those of overcomplete sparse coding representations [36])." If the authors have such results, they should show them. In particular, related to point (2) above, they should show how the convergence speed and number of reservoir neurons required to obtain accurate samples scales with the dimension of the target distribution. As the authors themselves discuss, their approach is currently limited to approximating the score with a single-hidden-layer MLP, and such Cybenko-Hornik-type approximation results may require immensely wide networks. This is also a limitation of the MNIST experiment: MNIST is intrinsically not that high-dimensional, and already here the authors used a network with 20000 reservoir neurons (Line 638)! 4. I think it's important to mention that the striking result of Figure 2 depends critically on the choice of ReLU activation function, particularly as much of the paper is concerned with existence results. If instead one used $\tanh$, then the performance gap between the trained sampler-only firing rate and reservoir firing rate models should be small. In particular, a sampler-only firing rate network can implement exact Langevin sampling in this case. The score of $p(x) = \frac{1}{2} \mathcal{N}(\mu , \sigma^{2}) + \frac{1}{2} \mathcal{N}(-\mu,\sigma^2)$ is $\frac{d}{dx} \log p(x) = \frac{1}{\sigma^{2}} \left[ - x + \mu \tanh\left(\frac{\mu}{\sigma^{2}} x\right) \right]$. This gives the Langevin dynamics $dx(t) = \frac{d \log p}{dx} dt + \sqrt{2} dB(t) = \frac{1}{\sigma^{2}} \left[ - x + \mu \tanh\left(\frac{\mu}{\sigma^{2}} x\right) \right] dt + \sqrt{2} dB(t)$. With $\mu = 1$ as is used in Figure 2, this is exactly of the form of a single-neuron sampler-only firing rate network in eq (8) of the submitted manuscript, with $D = 1/\sigma^{2}$, $W_{rec} = 1/\sigma^{2}$, $I=0$, and $\phi(x) = \tanh(x)$. One can easily verify numerically that this network produces accurate samples from the desired bimodal distribution.

Questions

- The figure axis labels are too small, making them hard to read. Figure 4b is also too small. Please increase the font and panel sizes to improve legibility. - Lines 32-33: Along with the other references on Gaussian sampling, Hennequin, Aitchison, and Lengyel, "Fast Sampling-Based Inference in Balanced Neuronal Networks" (NeurIPS 2014), should be cited. - Line 225: I do not see how the proof of Theorem 3 is "constructive," since it is based on an existence-type universal approximation theorem. - Lines 303-304: "This procedure is more aligned with the developmental processes involved in forming visual representations in the infant brain, where visual representations are thought to be noisier (less linearly separable) initially [3]" Please elaborate; I don't quite follow. - Line 317-318: This is quite vague. Having a clearer proposal for what this model might help explain in neuroscience would make for a much stronger paper. - It would be helpful to repeat the theorem statements in the Supplement. - Lemmas C.2 and C.3 can be replaced by using standard results on the asymptotic expansion of the upper incomplete Gamma function for large argument, see https://dlmf.nist.gov/8.11. - In Line 562, in the proof of Theorem 3, should $m$ be replaced by $n$?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The authors provide a generally adequate discussion of the limitations of their work (modulo the weaknesses discussed above), but fail to show the data for some points (see point 3 under Weaknesses).

Reviewer kc8J7/10 · confidence 2/52023-06-21

Summary

In this paper, the authors address the question of how the dynamics of recurrent neural networks (RNNs) can generate samples from probability distributions of interest. This is of interest to the neuroscience community, since it has been proposed that biological networks can use such sampling-based inference to represent and compute with probability distributions in order to account for their noisy surroundings. The authors prove that standard RNNs have limited expressivity when used directly for sampling-based inference. They proceed to show that RNNs with separate (linear) readouts in contrast have the necessary expressivity to approximate any probability distribution through Langevin sampling. Finally, it is demonstrated empirically that such a 'reservoir' formulation allows RNNs to generate samples from a variety of complex distributions in practice, including multimodal, heavy-tailed, and high-dimensional distributions.

Strengths

The paper contains several interesting theoretical results related to the topic of sampling-based inference using recurrent neural networks. This is an important open problem in computational neuroscience, where it remains unknown how the brain performs (approximate) probabilistic computations and represents the necessary probability distributions for such computations. By providing a firmer theoretical grounding for our understanding of sampling-based inference, this work has the potential to inspire both new computational and experiment studies of probabilistic computations in neural circuits. The paper is also generally well written and demonstrates impressive empirical results, although only on relatively 'toy' problems.

Weaknesses

Much of the motivation and discussion of the paper is written with reference to the neuroscience literature on sampling-based inference, especially in sensory cortices. However, the empirical results all use artificial toy examples with no obvious relation to questions of interest to the neuroscience community. It would have been interesting to e.g. look at natural images or similar and delve a bit deeper into the dynamics learned by the network to make it more interesting to a neuroscience audience. Many of the figure axis labels and legends are too small to be easily readable, and they would benefit from an increased font size.

Questions

L60: The authors state that they derive a 'biologically plausible' learning rule, but the training procedure described in Eq 14 does not seem particularly biologically plausible, requiring both backprop and matrix inverses (as the authors also highlight themselves in L288-291). Perhaps the introduction could be reworded to make this initial statement more in tune with the subsequent methods/results? L73: it would be useful to briefly state what these regularity conditions entail. L114: Helmholtz typo Figure 2 legend: "SO-SC and SO-FR is (_sic_) only able to fit the score function with piecewise linear function (_sic_) when using ReLU transfer function (_sic_)". Why the contrast between 'piecewise linear' in the legend and what looks like globally linear approximations to the score function in the figure? Ref 5 & 37: formatting errors in author name ("GergHo Orban"). I am slightly confused by the authors' network dynamic equations (Eq 7 & Eq 8). When modeling neural firing rates, a positive (e.g. ReLU) nonlinearity is commonly used since firing rates are strictly positive. However, would the Brownian noise term added to the rates in Eq 8 not allow for negative neural firing rates? I was also surprised by the synaptic current equations, since RNN dynamics are usually formulated in terms of either firing rates or membrane potentials. Potentials and currents are of course equivalent for ohmic units, but I was wondering if there was a reason that the authors decided to describe Eq 7 in terms of currents rather than membrane potentials? The authors have discussed the case of a non-linear RNN with a linear readout and shown that it can approximate arbitrary probability distributions. Is the same true for non-linear RNNs with non-linear readouts? and if so, could the reservoir sampler network be considered a special case of a non-linear RNN where only a subset of units are used to represent samples from the target distribution, and the remainder are used for additional computational power/expressivity?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

The authors have adequately addressed the limitations of their study.

Reviewer wwgi6/10 · confidence 3/52023-07-03

Summary

- This study solves an outstanding problem involving arbitrary density sampling using stochastic neural networks. - The authors propose a a Reservoir Sampling architecture, whereby an auxiliary recurrently-connected population facilitates the sampling from a non-trivial distribution. - The work seems solid and elegant (although SDEs are not my domain of expertise), and the numerical experiments are convincing. UPDATE: Sep 1, 2023. I have read the rebuttal, it addressed my Qs and I maintain my score.

Strengths

- The math looks solid, although I did not have time to go through the supplement. - There are numerical experiments on complex distributions (MNIST PCs), demonstrating that their procedure can operate on non-trivial densities.

Weaknesses

- This work seems to address an _existence_ problem with arbitrary density sampling with RNNs, but how would one validate this kind of model in physiological experiments? - In fact, I would appreciate more discussion on the state of affairs in the neurophysiology of density representation.

Questions

- This auxiliary RNN population approach reminds me of a recent poster/paper at cosyne by Duong et al. ("Adaptive whitening in neural populations with gain-modulating interneurons"; ICML 2023 https://openreview.net/forum?id=cEWB5hABV5), where they show that if the auxiliary population is greater than a specific size, it can represent and whiten a multivariate gaussian densities exactly. In your framework, is there a relationship between the dimensionality of reservoir and the expressivity of the network in terms of what kinds of densities it can represent? Could this provide testable predictions for an experimentalist to search for? - Suggestion: I did enjoy reading this manuscript, but there was more emphasis on technical parts than necessary to get your point across. If the text were reframed with less emphasis on the math, I think it would flow much nicer and would be much more approachable.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

There was adequate discussion of limitations.

Reviewer Y31x7/10 · confidence 4/52023-07-06

Summary

The authors proposed a revervoir-sampler network whose firing rate dynamics can sample from an arbitrary probability distribution. They first established the relationship between the sampling power of the neural dynamics and the ability of the dynamics to approximate the score function. Then they showed that the synaptic current dynamics of the traditional sampler-only networks is only able to approximate score functions that are in a finite-dimensional function space. Importantly they proved that the revervoir-sampler network can sample from an arbitrary probability distribution to arbitrary precision. Finally, they developped a biologically-plausible learning rule for the proposed model.

Strengths

This paper is very clearly presented and the scientific question is quite interesting.

Weaknesses

Comparison with previous models should be addressed in the experiment part.

Questions

Q1: Why did the authors choose to use the reservoir network as the building block? Q2: Following the first question, does the RSN model work because of the linear readout layer, which rescales the distributions? Can the authors explaine more about this? Q3: Can the current model be extended to the multiplicative noise scenario which could be the feature of the poisson spike neurons, or (at the network level) be the feature of the E-I balanced spiking net. Q4: How does the RSN model compare to the coupled attractor model which also implement Langevin sampling [1]? [1]: Zhang, W. H., Lee, T. S., Doiron, B., & Wu, S. (2020). Distributed sampling-based Bayesian inference in coupled neural circuits. bioRxiv, 2020-07.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The authors addressed some limitations.

Reviewer sqZn6/10 · confidence 4/52023-07-07

Summary

The paper investigates architectural requirements for recurrent neural circuits to sample from complex distributions using diffusion models. It presents a model where traditional sampler-only networks are enhanced with additional firing-rate dynamics and a set of separate output units, called reservoir-sampler networks (RSNs). An efficient training procedure based on denoising score matching is proposed. Empirical experiments are presented to demonstrate the model's ability to sample from complex data distributions.

Strengths

- The research addresses a significant issue in computational neuroscience and Bayesian learning – how neural dynamics can sample a complex distribution. - Biological plausibility and mathematical computation efficiency are particularly commendable aspects of this work. - The theoretical analysis is well constructed, providing strong support for the model design, including the choice of firing-rate dynamics and the reservoir-sampler network. - The empirical experiments presented provide evidence of the effectiveness of the proposed method.

Weaknesses

- The experiments conducted need improvements. Despite using score matching for model training, the quality of the generated images leaves room for improvement. The incorporation of Unet-like or transformer backbones might have enhanced the results. - The paper lacks a clear, diagrammatic representation to explain the utilization of diffusion models in the design of recurrent neural circuits architecture. - The results and conclusions seems to have limited impact to the general AI/ML areas, giving limited experimental results and comparison (e.g. precision and efficiency ) with existing diffusion probabilistic models is somewhat lacking.

Questions

See weaknesses.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

I am satisfied with the discussion of limitations.

Reviewer kc8J2023-08-10

I appreciate the thorough response from the authors to myself and the other reviewers. The additional experiments with tanh nonlinearities and investigations of relaxation time will be particularly useful additions to the paper, and I look forward to seeing the results on the CIFAR 10 dataset when they are finished. I would also consider adding a sentence or two to the paper about the distinction between an RNN with a linear readout vs. an RNN where only a subset of neurons represent the target distribution, since the latter seems more intuitive from a neuroscience point of view.

Authorsrebuttal2023-08-11

Thank you

Thank you for your review and feedback. We agree with your suggestion for additional experiments on tanh nonlinearities, relaxation time, the CIFAR10 dataset, and more explanations of the distinction between RSN and sampling from part of a larger RNN. We will make sure to incorporate these changes in our revision.

Reviewer pX3M2023-08-10

I thank the authors for their thoughtful reply to my questions and those of the other reviewers. I appreciate their addition of tests relating to convergence time; I think this will enhance the final version. Given these additions, I will raise my score. A few small replies to individual points: - Of course the choice of activation function will affect the choice of distributions that can be learned; I do not think the authors need to replace the existing ReLU figure. I think it is sufficient to note that for certain distributions the score can be realized exactly by sampler-only networks. - I still do not follow the relationship with infant learning. This is, however, a minor point, so I do not think further elaboration is necessarily required.

Authorsrebuttal2023-08-11

Thank you

Thank you for your review and feedback. We agree with your suggestion to note the fact that sampler-only networks can sample exactly from some distributions. We will make sure to incorporate these changes in our revision.

Reviewer wwgi2023-08-10

Thanks for your responses to my review. Because SDE's are not my expertise, there is a possibility that I'm missing technical details, so I will keep my score the same. But I will raise my confidence (2 -> 3) to reaffirm that I think the community would find this interesting. I do find the relation to neurophys and testable predictions a little weak, but that seems to just be where the field is at right now. I will reiterate that there is a larger audience that could appreciate this work, but some of the simple concepts are opaquely hidden behind the over-emphasis on technical presentation its present state. So, I hope the reviewers follow through on reformulating the presentation, including thinking about how to reach a broader neuroscience community and not just the neural probabilistic sampling community.

Authorsrebuttal2023-08-11

Thank you

Thank you for your review and feedback. We are in agreement with your suggestion to reformulate the manuscript for a broader neuroscience audience. We will make sure to incorporate these changes in our revision.

Reviewer sqZn2023-08-15

Thanks

The response addresses most of my concerns, therefore I decide to keep my original ratings.

Authorsrebuttal2023-08-15

Thank you

We thank the reviewer for the review and feedback again.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC