SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
d \mathbf {x} = (- \alpha \mathbf {x} + \alpha W _ {\text {o u t}} \phi (\widetilde {W} _ {\text {r e c}} \mathbf {x} + I)) d t + W _ {\text {o u t}} \sigma d B _ {t} := \widetilde {F} _ {\theta} ^ {\mathrm {F R}} (\mathbf {x}) d t + \tilde {\sigma} d B _ {t} \tag {11}
$$
where  $\widetilde{F}_{\theta}^{\mathrm{FR}}(\mathbf{x}) = (-\alpha \mathbf{x} + \alpha W_{\mathrm{out}}\phi (\widetilde{W}_{\mathrm{rec}}\mathbf{x} + I))$  and  $\tilde{\sigma} = W_{\mathrm{out}}\sigma$ . Therefore in order for the output neurons to sample from a stationary distribution  $p$ , we need  $s_{\beta}(\mathbf{x}) = (\frac{1}{2}\tilde{\sigma}\tilde{\sigma}^{T})^{-1}\widetilde{F}_{\theta}^{\mathrm{FR}}(\mathbf{x}) := \widetilde{\Sigma}^{-1}\widetilde{F}_{\theta}^{\mathrm{FR}}(\mathbf{x})$  to match the score  $\nabla \log p(\mathbf{x})$ . Here  $\beta = \left\{\widetilde{W}_{\mathrm{rec}},W_{\mathrm{out}},I,\sigma \right\}$  are tunable parameters. Additionally, we assume that  $\widetilde{F}_{\theta}^{\mathrm{FR}}$  is 0 outside a reasonable range of  $\mathbf{x}$ . This assumption is used to prevent  $s_{\beta}(\mathbf{x})$  from behaving wildly outside the bounded region on which  $s_{\beta}(\mathbf{x})$  has the expressivity to match the score. The following theorem proves that with a large enough number of reservoir neurons, the score-matching loss can be arbitrarily small. The proof is given in Appendix C.
Theorem 3. Suppose that we are given a probability distribution with continuously differentiable density function  $p(\mathbf{x}) : \mathbb{R}^m \to \mathbb{R}^+$  and score function  $\nabla \log p(\mathbf{x})$  for which there exist constants  $M_1, M_2, a, k > 0$  such that
$$
p (\mathbf {x}) <   M _ {1} e ^ {- a \| \mathbf {x} \|} \tag {12}
$$
$$
\left\| \nabla \log p (\mathbf {x}) \right\| ^ {2} <   M _ {2} \| \mathbf {x} \| ^ {k} \tag {13}
$$
when  $\| \mathbf{x}\| >L$  for large enough  $L$ . Then for any  $\varepsilon >0$ , there exists a recurrent neural network whose firing-rate dynamics are given by (11), whose recurrent weights, output weights and the diffusion coefficient are given by  $W_{\mathrm{rec}}\in \mathbb{R}^{n\times n}$  of rank  $m$ ,  $W_{\mathrm{out}}\in \mathbb{R}^{m\times n}$ , and  $\sigma \in \mathbb{R}^{n\times m}$  respectively, such that, for a large enough  $n$ , the score of the stationary distribution of the output units  $s_{\theta}(\mathbf{x})$  satisfies  $\mathbb{E}_{\mathbf{x}\sim p(\mathbf{x})}[\| \nabla \log p(\mathbf{x}) - s_{\theta}(\mathbf{x})\| ^2 ] < \varepsilon$ .
This theorem says that for any realistic data distribution with a smooth positive density function, there always exists a reservoir of recurrently-connected neurons whose output units give samples from a
distribution whose score function approximates that of the data distribution to arbitrary precision. Given the bound on the score matching error, Block et al. [10] (cf. Theorem 13) gives bounds in Wasserstein 2-distance between the stationary distribution of the trained recurrent dynamics and the true data distribution. It is also worth noting that the recurrent weight matrix of the neural circuit in the theorem is of low-rank, so regardless of how many neurons there are in the reservoir, we can always find a low-rank recurrent weight matrix such that the output neurons sample from a correspondingly low-dimensional distribution.
3.4 An efficient RNN weight learning algorithm
The training procedure is derived from the proof of Theorem 3. The main idea is to first train an auxiliary neural network with one hidden layer, then transfer the weights of the auxiliary neural network to the weights of the recurrent network that we are considering. More specifically, we first optimize an auxiliary feedforward neural network with one hidden layer  $W_{\mathrm{out}} \phi (\widetilde{W}_{\mathrm{rec}} \mathbf{x} + I)$  using backpropagation with the denoising score matching loss (4) such that
$$
2 \alpha \left(W _ {\text {o u t}} \phi \left(\widetilde {W} _ {\text {r e c}} \mathbf {x} + I\right) - \mathbf {x}\right) \approx \nabla \log p (\mathbf {x}). \tag {14}
$$
Then we can calculate the diffusion coefficient  $\sigma = W_{\mathrm{out}}^{T}(W_{\mathrm{out}}W_{\mathrm{out}}^{T})^{-1}$  and the real recurrent weights  $W_{\mathrm{rec}} = \widetilde{W}_{\mathrm{rec}}W_{\mathrm{out}}$  accordingly. The noise magnitude added to the data samples is decreased exponentially over the entire training period. Figure 2 illustrates how the network can gradually learn the score function during training. We refer the readers to Appendix D for more details.
Note that this method of using an auxiliary neural network is much more computationally efficient than directly matching the score of the stationary distribution of the dynamics (11), for which the score function, which involves the matrix inverse  $(\tilde{\sigma}\tilde{\sigma}^T)^{-1}$ , needs to be recomputed at each optimization step. Furthermore, since the entire training procedure is equivalent to training a feedforward network with one hidden layer, it sidesteps the often challenging temporal computations associated with the Backpropagation Through Time (BPTT) algorithm used to train deterministic RNNs.
Although we assumed that the divergence free field  $G = 0$  for the purpose of theoretical analysis, in practice, fast sampling is a major concern for implementing sampling-based inference models of the brain [25, 1, 19, 38], and reversible stochastic dynamics, i.e. if  $G = 0$ , are known for their slow sampling speed. Fortunately, there is a straightforward way to extend our framework and train RSNs to implement irreversible dynamics with a non-zero divergence-free field  $G$ . This results in improvements in sampling speed (see Appendix E for details).
In this section, we present the results from two tasks. First, we try to let the recurrent neural network learn and generate samples from a 1-D double-mode Gaussian mixture distribution and a 2-D mixture of heavy-tailed Laplace distributions. For the second task, we explore whether a reservoir-sample network with firing rate dynamics is able to sample from the distribution of internal representations computed from PCA filtering of image inputs. All dynamics are simulated with the Euler-Maruyama method. See Appendix H for hyperparameters used and other training details.
4.1 Learning mixture distributions
We consider a 1-D Gaussian mixture distribution whose density function is the average of two Gaussian distributions centered at  $\pm 1$ , i.e.  $p_{\mathrm{data}}(x) = \frac{1}{2} (\mathcal{N}(-1,0.25) + \mathcal{N}(1,0.25))$ . We artificially generate 10000 data points from this distribution and minimize the denoising score-matching loss:
$$
\mathcal {L} \left(W _ {\text {o u t}}, \widetilde {W} _ {\text {r e c}}, I\right) = \mathbb {E} _ {p _ {\text {d a t a}} (x)} \mathbb {E} _ {\tilde {x} \sim \mathcal {N} \left(x, \sigma^ {2}\right)} \left[ \left\| 2 \alpha \left(W _ {\text {o u t}} \phi \left(\widetilde {W} _ {\text {r e c}} x + I\right) - x\right) + \frac {\tilde {x} - x}{\sigma^ {2}} \right\| ^ {2} \right]. \tag {15}
$$
We use  $\phi (\cdot) = RLU(\cdot)$  and set  $\alpha = 1 / 2$ . See Figure 2 for the numerical results. It is worth noting that if we use tanh as the transfer function, the sampler-only networks are able to learn the score function perfectly, as the score function of the Gaussian mixture distribution we considered is exactly spanned by  $f_{1}(x) = x$  and  $f_{2}(x) = \tanh (2x)$ . See Appendix F.1 for an example where sampler-only networks fail to learn the score function even if hyperbolic tangent nonlinearity is used.
Figure 2: Bimodal distribution sampling results. The 3 tractable cases shown are Sampler-only (SO) networks with both synaptic current (SC) dynamics and firing rate (FR) dynamics and Reservoir-sampler (RS) networks with FR dynamics, which are named SO-SC, SO-FR and RS-FR respectively. a-c) The score function learned compared to the true score function (orange curve) as we gradually decrease the noise level (the darker the line, the lower the noise level). We see that RS-FR is capable of perfectly fitting the score function, while SO-SC and SO-FR are only able to fit the score function with piecewise linear functions when using the ReLU transfer function. d-f) Histogram of sampled points, and the (scaled) density function of the target distribution. Again the reservoir-sampler network is able to generate samples whose distribution matches the target distribution, while the sampler-only network is not able to do so due to the incorrectly matched score function.
Next, we show that the model can learn mixtures of heavy-tailed distributions that are evident in natural image statistics and the neural representations in the primary visual cortex [47, 39]. We trained the Reservoir-Sampler network with FR dynamics (RS-FR) model on 20000 sampled data points from a 2-D Laplace mixture distribution, whose density is given by  $p_{\mathrm{data}}(\mathbf{x}) = \frac{1}{2}\left(\mathrm{Lap}\left(\mathbf{0},\left[ \begin{array}{cc}1 & 0.9\\ 0.9 & 1 \end{array} \right]\right) + \mathrm{Lap}\left(\mathbf{0},\left[ \begin{array}{cc}1 & -0.9\\ -0.9 & 1 \end{array} \right]\right)\right)$ , where Lap denotes the multivariate Laplace distribution. The model successfully learned the probability density of the mixture distribution (Figure 3 left vs. middle), and captured the heavy tails of the distribution as measured by the kurtosis (Figure 3 right).
4.2 MNIST generation task
We also tested the sampling ability of our model using the MNIST dataset [32] which contains 60,000 handwritten digits from 0 to 9. We projected MNIST images to a 300-D latent space spanned by the first 300 principal components, and trained the weights of the recurrent neural network as described in Section 3.4 so that the RNN can sample from the latent distribution. To test the model, we generated images by applying inverse PCA projection to samples generated by the model. The schematics and generated images are shown in Figure 4. Note that since we are essentially using a shallow network to match the score, we should not expect comparable performance to generative models that use deep ANNs. Our main goal is to illustrate that the reservoir-sampler network using firing rate dynamics is qualitatively more expressive than other traditional neural sampling models (Appendix G). Finally we also note that it is highly nontrivial for recurrent neural dynamics to complete such a generative task, and to the best of our knowledge, no previous work has achieved such results.
Figure 3: RS-FR model learning a mixture of 2-D Laplace distribution. Left to right: sample density from the true distribution (brighter color denotes higher density); sample density from the learned distribution; marginalized kurtosis of each dimension from the true and learned distribution.
Figure 4: Learning to sample the MNIST image distribution. a) An MNIST image is projected to a 300-D latent space (orange circle) spanned by the first 300 principal components using PCA. The sampler learns to sample from the distribution of this latent space and generates images by applying inverse PCA projection to these samples. The diagram illustrates the RS-FR model. b) The loss curves for 3 different models during training. Every 100 epochs the noise level added to the training samples is reduced (Appendix H), and the noise increases to a higher value because the score matching loss magnitude depends on the noise level. As shown in the inset, the loss of the RS-FR model decreases throughout the training process when using the lowest fixed noise level. Meanwhile, the losses of the other two models remain unchanged. c-f) The images generated for the 3 models compared to the digit images generated from latent training samples.
From the perspective of functional analysis and SDE theory, we prove that under the strict orthogonality constraint, it is essential for neural circuits to have a drift term that has the expressivity to approximate complex score functions, despite the fact that the dynamics do not have to exactly implement the Langevin dynamics. We investigated whether a population of neurons can sample from an arbitrary distribution directly and proved that the drift term of the synaptic current dynamics can only sample from a finite-dimensional function space. Although the drift term of the firing
rate dynamics can approximate functions spanned by different basis functions, the number of basis functions is limited. To address this problem, we proposed the reservoir-sampler network for firing rate dynamics. We found that with learnable diffusion coefficients and a sufficiently large reservoir of hidden neurons, the output neurons described using the firing rate dynamics are able to sample from arbitrary data distributions. Our results partly answer the question of what architecture recurrent neural circuits need so that they are able to sample from complex data distributions.
Our analysis and empirical experiments affirm the universality of stochastic RNNs. However, this universality comes with limitations. First, we have only analytically shown the existence of weights that enable sampling from complex data distributions; there is no guarantee that one will find such weights through backpropagation. Additionally, in order to obtain the tunable diffusion coefficient during training, a matrix inverse is needed (likewise in the FORCE algorithm [50]). Further, the question of how biological circuits compute the specific gradient and implement the denoising score-matching algorithm remains an open question. Moreover, in our current formulation, we are only able to approximate the score function with a shallow network with one hidden layer. 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 [39]). It is unclear if this is because of insufficient number of reservoir neurons. Due to the limitation of the GPU memory, we did not try higher number of reservoir neurons.
Our model differs from recent diffusion models [26, 49], which can be seen as time-inhomogeneous SDEs, and has the advantage of being able to run indefinitely in time, making it a suitable candidate for modeling spontaneous activity in the brain. Moreover, while Song and Ermon [48] optimizes the denoising score matching loss at different noise levels jointly, we adopt a sequential learning procedure by gradually decreasing the noise level of the training samples. This procedure is more aligned with the developmental processes involved in forming visual representations in the infant brain, where the distribution of visual representations are thought to be noisier (less linearly separable) initially [4]. Our study therefore serves as a starting point for building a mechanistic model for probabilistic computation in the brain that has similar generative power to current AI generative models.
Biologically, there are multiple ways to interpret the reservoir neurons and sampler neurons in an RSN. First, reservoir and sampler neurons could be seen as different types of neurons in a single brain area, where the dynamics of sampler neurons converge quickly to the equilibrium point. Second, even more straightforwardly, the sampler neurons could be seen as a more separate set of neurons located downstream of the reservoir. We also wish to suggest an alternative interpretation. Biological neural networks are known to have non-synaptic signaling networks (e.g. pervasive neuropeptidergic signaling [3], extensive aminergic signaling [6] or potential extrasynaptic signaling [53]) in addition to the synaptic connectivity that is typically modeled (i.e., via connection weights). We suggest that it is possible that the computations of the reservoir may be implemented by non-synaptic networks, and then "read out" by certain neurons' spikes. This possibility is supported by the recent finding of a low correlation between functional activity and the synaptic ("structural") connectome in C.elegans [53]. Moreover, if we only take the structural connectome into consideration, then the resulting model of C. elegans would correspond to the sampler-only network, which, as our theory predicts, will have limited sampling capability.
In this paper, we explore how a recurrent neural circuit can sample from complex probability distributions, an important functional motif in probabilistic brain models. We start from a basic assumption that the recurrent neural circuit could be described as an SDE. We show that a recurrently-connected neural population by itself has a limited capability to implement stochastic dynamics that can sample from complex data distributions. In contrast, we prove that firing rate dynamics of the output units of a recurrent neural circuit (a reservoir-sampler network) can sample from a richer range of probability distributions. These theoretical results, together with our preliminary experimental results, provide a sufficient condition for neural sampling-based models to exhibit universal sampling capability. Our results therefore provide a foundation for the next generation of sampling-based probabilistic brain models that can explain a wider range of cognitive behaviors.
We are thankful to Profs. Hong Qian, Bamdad Hosseini and Edgar Walker for their guidance and insight with this project. We gratefully acknowledge the support of the grant NIH BRAIN R01 1RF1DA055669.
[52] Eszter Vértes and Maneesh Sahani. Flexible and accurate inference and learning for deep generative models. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https://proceedings.neurips.cc/paper/2018/bit/955cb567b6e38f4c6b3f28cc857fc38c-Abstract.html.
[53] Eviatar Yemini, Albert Lin, Amin Nejatbakhsh, Erdem Varol, Ruoxi Sun, Gonzalo E Mena, Aravinthan DT Samuel, Liam Paninski, Vivek Venkatachalam, and Oliver Hbert. Neuropal: a multicolor atlas for whole-brain neuronal identification in c. elegans. Cell, 184(1):272-288, 2021.
A Helmholtz-Hodge decomposition and proof of Proposition 1
First we give a proof for Proposition 1:
Proposition. Let  $p$  be the stationary distribution of the neural dynamics, and the diffusion matrix be the identity matrix. If the  $DF$  field  $G$  is strictly orthogonal to the gradient field  $\nabla p$ , meaning that  $G(\mathbf{v}) \cdot \nabla p(\mathbf{v}) = 0$  for all  $\mathbf{v}$ , then the drift term  $F_{\theta}(\mathbf{v})$  can be written as the sum of a divergence-free field  $p^{-1}G$  and a gradient field  $\nabla \log p$ .
Proof. Since  $p$  is the stationary distribution of the neural dynamics, from the Fokker-Planck equation we have
$$
F _ {\theta} (\mathbf {v}) = p ^ {- 1} G + \nabla \log p \tag {16}
$$
where  $G$  is divergence free. Now
$$
\begin{array}{l} \nabla \cdot (p ^ {- 1} G) = \nabla p ^ {- 1} \cdot G + p ^ {- 1} \nabla \cdot G \\ = \frac {\nabla p \cdot G}{p ^ {2}} \quad \text {G i s d i v e r g e n c e f r e e} \\ = 0. \quad \mathrm {G} \text {a n d} \nabla p \text {i s p o i n t w i s e o r t h o g o n a l} \\ \end{array}
$$
Therefore,  $p^{-1}G$  is also divergence free.
Now we give a statement of the Helmholtz-Hodge decomposition, which is an application of the Hodge decomposition theorem to vector calculus on  $\mathbb{R}^n$  [8, 37].
Theorem 4 (Helmoholtz-Hodge Decomposition). A smooth vector field  $\mathbf{V}$  on a bounded domain of  $\mathbb{R}^n$  can be uniquely decomposed into the sum of a gradient field  $\nabla u$  that is normal to the boundary, a divergence-free field  $\mathbf{v}$  that is parallel to the boundary and a harmonic field  $\mathbf{h}$ , i.e.  $\mathbf{V} = \nabla u + \mathbf{v} + \mathbf{h}$ . Moreover, the three components are orthogonal to each other in the  $L_2$  sense.
We note that the form of the decomposition in Proposition 1 coincides with the Helmholtz-Hodge decomposition with a zero harmonic field. Therefore, we can define the orthogonal projection operator  $\mathbb{P}$  that projects  $F_{\theta}$  onto its gradient component  $\nabla \log p$ . Under the strict orthogonality constraint that requires  $\mathbf{G}$  to be orthogonal to  $\nabla p$  pointwise, to be able to sample from a broad range of distributions with  $\nabla \log p$  orthogonal to the boundary, we need that for all  $p$ , there exists  $G_{p}$  and  $\theta$  such that  $F_{\theta}(\mathbf{v}) = p^{-1}G_{p} + \nabla \log p$ . Applying the projection operator on both sides, we have
$$
\mathbb {P} F _ {\theta} (\mathbf {v}) = \mathbb {P} \left(p ^ {- 1} G _ {p} + \nabla \log p\right) = \nabla \log p. \tag {17}
$$
Moreover, if the function space  $\{F_{\theta}\}_{\theta}$  admits a set of basis functions  $\{e_i\}_i$ , i.e.  $F_{\theta} = \sum_{i}c_{i}e_{i}$ , then
$$
\nabla \log p = \mathbb {P} F _ {\theta} = \sum_ {i} c _ {i} (\mathbb {P} e _ {i}). \tag {18}
$$
This means that the space of  $\nabla \log p$  needs to be spanned by  $\{\mathbb{P}e_i\}_i$ , so  $\{F_{\theta}\}_{\theta}$  should also be spanned by an infinite number of basis functions  $\{e_i\}_i$  if we want to sample from a complex class of  $\nabla \log p$ .
B Neural dynamics not following gradient fields
For the meaning of notations, see Section 3.2.
Proposition 5. When  $m \geq 2$ ,  $F_{\theta}^{\mathrm{SC}}(\mathbf{v}) := D(-\mathbf{v} + W\phi(\mathbf{v}) + I)$  is a gradient field if and only if  $W$  is a zero matrix.
Proof.  $(\Rightarrow)$  If  $F_{\theta}^{\mathrm{SC}}(\mathbf{v}) = \nabla G(\mathbf{v})$  is a gradient field, then
$$
\frac {\partial}{\partial v _ {j}} F _ {\theta} ^ {\mathrm {S C}} (\mathbf {v}) _ {i} = \frac {\partial}{\partial v _ {i}} F _ {\theta} ^ {\mathrm {S C}} (\mathbf {v}) _ {j} = \frac {\partial}{\partial v _ {i} \partial v _ {j}} G (\mathbf {v}) \tag {19}
$$
because  $F_{\theta}^{\mathrm{SC}}(\mathbf{v})$  is smooth, and we can switch the order of partial differentiation. Written in the form of (7), we have for all  $v_{i}$  and  $v_{j}$  (where  $i \neq j$ ),
$$
\begin{array}{l} \frac {\partial}{\partial v _ {j}} \left(- d _ {i} v _ {i} + d _ {i} \sum_ {k = 1} ^ {n} W _ {i k} \phi (v _ {k}) + d _ {i} I _ {i}\right) = \frac {\partial}{\partial v _ {i}} \left(- d _ {j} v _ {j} + d _ {j} \sum_ {k = 1} ^ {n} W _ {j k} \phi (v _ {k}) + d _ {j} I _ {j}\right) \\ \frac {\partial}{\partial v _ {j}} d _ {i} W _ {i j} \phi (v _ {j}) = \frac {\partial}{\partial v _ {i}} d _ {j} W _ {j i} \phi (v _ {i}) \\ d _ {i} W _ {i j} \phi_ {v} (v _ {j}) = d _ {j} W _ {j i} \phi_ {v} (v _ {i}), \\ \end{array}
$$
since  $v_{i}$  and  $v_{j}$  are arbitrary and  $\phi$  is nonlinear (so the derivative  $\phi_v$  is not the same for all  $v$ ), this implies  $d_i W_{ij} = d_j W_{ji} = 0$ . But the decay constants cannot be 0. Therefore,  $W_{ij} = W_{ji} = 0$
(  $\Longleftarrow$  ) If  $W$  is a zero matrix, then  $F_{\theta}^{\mathrm{SC}}(\mathbf{v}) = \nabla D\left(-\frac{\mathbf{v}\cdot\mathbf{v}}{2} +I\cdot \mathbf{v}\right)$  is a gradient field.
B.1 Hilbert projection theorem and proof of Proposition 2
Here we directly state the Hilbert Projection Theorem. For detailed proof, we refer readers to standard functional analysis textbooks [44].
Theorem 6. (Hilbert Projection Theorem) Given a Hilbert space  $H$ , and a nonempty closed convex set  $E \subset H$ , for any  $x \in H$ , there exists a unique  $\Pi x \in E$ , such that  $\| \Pi x - x \| = \inf_{c \in E} \| x - c \|$ .
Here  $\Pi$  is called the orthogonal projection operator. Next, we give the proof of Proposition 2.
Proof. Here we consider a special case of a Hilbert space  $L_{2}(\mathbb{R}^{m},\mathbb{R}^{m};p)$ , where the inner product is defined by  $\langle f,g\rangle = \int_{\mathbb{R}^m}|f\cdot g|p(\mathbf{x})d\mathbf{x}$  and we let  $E = \{A\mathbf{v} + B\phi (\mathbf{v}) + I|A,B\in \mathbb{R}^{m\times m};I\in \mathbb{R}^{n\times 1}\}$ . The set of functions  $E$  is a closed convex vector subspace of  $L_{2}(\mathbb{R}^{m},\mathbb{R}^{m};p)$  because  $E$  is closed under addition and finite-dimensional. Hence by Hilbert projection theorem, there exists a unique  $\Pi H\in E$  such that  $\| (1 - \Pi)H\| = \inf_{F\in E}\| H - F\| >0$ . Because  $\left\{F_{\theta}^{\mathrm{SC}}\right\}_{\theta}\subset E$  if the decay matrix  $D$  is diagonal, we have