Natural gradient enables fast sampling in spiking neural networks
Abstract
For animals to navigate an uncertain world, their brains need to estimate uncertainty at the timescales of sensations and actions.
Sampling-based algorithms afford a theoretically-grounded framework for probabilistic inference in neural circuits, but it remains unknown how one can implement fast sampling algorithms in biologically-plausible spiking networks.
Here, we propose to leverage the population geometry, controlled by the neural code and the neural dynamics, to implement fast samplers in spiking neural networks.
We first show that two classes of spiking samplers—efficient balanced spiking networks that simulate Langevin sampling, and networks with probabilistic spike rules that implement MetropolisHastings sampling—can be unified within a common framework.
We then show that careful choice of population geometry, corresponding to the natural space of parameters, enables rapid inference of parameters drawn from strongly-correlated high-dimensional distributions in both networks.
Our results suggest design principles for algorithms for sampling-based probabilistic inference in spiking neural networks, yielding potential inspiration for neuromorphic computing and testable predictions for neurobiology.
Introduction
Neural circuits perform probabilistic computations at the sensory, motor and cognitive levels [1–4].
From abstract representations of decision confidence [5] to estimates of sensory uncertainty in visual cortex [6, 7], evidence of probabilistic representations can be found at all levels of the cortical processing hierarchy [8].
To be behaviorally useful, these probabilistic computations must occur at the speed of perception [9].
However, how neuronal dynamics allow brain circuits to represent uncertainty in high-dimensional spaces at perceptual timescales remains unknown [3, 10–12].
Several neural architectures for probabilistic computation have been proposed, including: probabilistic population codes [13], which in certain cases allow a direct readout of uncertainty; direct encoding of metacognitive variables, such as decision confidence [4, 5, 8]; doubly distributional codes [14, 15] which distinguish uncertainty from multiplicity; and sampling-based codes [9, 16–24], where the variability in neural dynamics corresponds to a signature of exploration of the posterior probability.
Most experiments quantifying uncertainty representations in single biological neurons have only
[Figure]
varied parameters along one or two dimensions, such as in Bayesian cue combination [2, 10, 25].
In these conditions, many algorithms can perform adequately.
However, probabilistic inference becomes more challenging as the entropy of the posterior distribution—which often scales with dimensionality—increases [26, 27].
Some algorithms that work well in low dimensions, such as probabilistic population codes, may scale poorly to high-dimensional settings [3].
Of the proposed approaches to probabilistic computation in neural networks, sampling-based codes are grounded in the strongest theoretical framework [26, 28–32], and have been used to perform inference at scale [33–35].
Moreover, they predict specific properties of neural responses in visual cortex, including changes in Fano factor and frequency of oscillations with tuning and stimulus intensity [23].
Previous works have proposed several approaches to accelerate sampling in biologically-inspired algorithms.
Hennequin et al.
[9] showed that adding non-reversible dynamics to rate networks can reduce the sample autocorrelation time.
However, they did not study convergence of the sampling distribution, and did not consider the biologically-relevant setting of spiking networks.
Savin and
Denève [20] used a distributed code to parallelize sampling in spiking networks, but only considered two-dimensional distributions.
Therefore, it remains unclear how accurate sampling from highdimensional distributions at behaviorally-relevant timescales can be achieved using spiking networks.
In this paper, we show how the choice of the geometry of neural representations at the population level [36, 37], set by the neural code and the neural dynamics, can accelerate sampling-based inference in spiking neural networks.
Ideas from information geometry allow us to perform inference in the natural space of parameters, which is a manifold with distances measured by the Fisher-Rao metric (Figure 1.c)
Concretely, we leverage recently-proposed methods for accelerating sampling from the machine learning literature [27, 41–43] to design novel efficient samplers in spiking neural networks.
The structure and major contributions of this paper are divided as follows:
• In §2, we construct from first principles a novel spiking neural network model for sampling from multivariate Gaussian distributions.
This model is based on a probabilistic spike rule that implements approximate Metropolis-Hastings sampling.
We show that efficient balanced networks (EBNs) [20, 44–48] emerges as a limit of this model in which spiking becomes deterministic.
• In §3, we show that population geometry enables rapid sampling in spiking networks.
Leveraging the “complete recipe” for stochastic gradient MCMC [42], we establish principles for the design of efficient samplers in spiking neural networks.
Then, we show how neural population geometry enables fast sampling—on the timescale of tens of milliseconds—in two limits of the model introduced in §2: EBNs in which sampling is driven by stochastic Langevin dynamics [20], and networks in which sampling is driven purely by a Metropolis-Hastings probabilistic spiking rule.
• Finally, in §4 we conclude by discussing the implications of our results in the context of prior works, and highlight their limitations as well as remaining open questions.
In particular, we comment on possibilities for future experimental studies of sampling in biological spiking networks, and applications to neuromorphic computing.
Spiking networks for sampling-based probabilistic inference
We begin by proposing a framework for probabilistic inference in spiking neural networks in which the spiking rule implements a
Metropolis
-Hastings step.
We show that EBNs [
can be recovered as a limiting case of this
more
general framework
In this work, we will keep our discussion quite general, and state our results for sampling from a generic Gaussian distribution.
However, the problem we aim to solve can be given a concrete interpretation in a neuroscience context, and could also be extended to non-Gaussian distributions.
The goal of a neural network performing probabilistic inference is to estimate a posterior distribution
P ( θ | x ) over n p latent variables (or parameters) θ given an input x (Figure 1).
The input could correspond to the activity of sensory neurons in early sensory processing (e.g. input onto ganglion cells in the retina or onto mitral cells in the olfactory bulb) or inputs into a cortical column that linearly sense features in the environment through an affinity matrix.
We provide a detailed discussion of this linear Gaussian model in Appendix C.2.
In the rest of the paper, we will usually abbreviate the distribution from which we want to sample as P ( θ ) , rather than writing P ( θ | x ) .
2.1 Deriving an approximate Metropolis-Hasting spiking sampler
We will build a network of n n spiking neurons to approximately sample an n p -dimensional Gaussian distribution of time-varying mean θ ( t ) and fixed covariance Ψ .
We first consider the case in which θ is constant, and then generalize the resulting algorithm to the case in which it is slowly time-varying.
As in prior work on probabilistic inference using spiking networks, we take the samples z to be linearly decoded from the filtered spike trains r of n n neurons [20, 44–48].
Working in discrete time for convenience and clarity (as in [47]), we let
[Equation]
where
[Equation] [Equation]
) is the low-pass filtered history of spikes o t ∈ { 0 , 1 } n n for some decay constant 0 ≤ η ≤ 1 .
Metropolis-Hastings sampling constructs a Markov chain by drawing a proposed next state from some distribution, and then deciding whether to accept or reject that proposal based on a probabilistic rule [26, 49].
The acceptance ratio is given in terms of the relative posterior probability of the proposed and current states.
Here, we will randomly propose which neuron spikes at a given timestep.
In trying to build a Metropolis-Hastings sampler [26, 49] using a probabilistic spiking rule, we are immediately faced with two problems.
First, the spikes are sign-constrained and discrete.
Second, the dynamics of the filtered spike history incorporate a decay term, hence the readout z will change even if no spikes are emitted.
These conditions mean that the proposal density over z will not be symmetric, and that the Markov process will not satisfy the condition of detailed balance [26, 49].
We note that previous work has shown that violations of detailed balance can accelerate sampling [17, 18, 50], but we will not carefully explore this possibility in the present work.
To obtain a symmetric proposal distribution with sign-constrained spikes, we assume that the network is divided into two equally-sized populations with equal and opposite readout weights, i.e. that the readout matrix is of the form Γ =
[+ M , − M ] for some matrix M ∈ R n p × n n / 2 .
This could be accomplished by dividing the total population of neurons into excitatory and inhibitory populations with weights that are fine-tuned to be equal and opposite.
The second problem can be solved by assuming that η = 0 , i.e., that we have access to a perfect integrator of the spike trains.
At the t -th timestep, we choose one neuron j uniformly at random, and let the spike proposal be o ′ =
e j , where e j is the j -th standard Euclidean basis vector (i.e., ( e j ) k
= δ jk ).
This yields a candidate readout
[Equation]
If η = 0 and the balance assumption on Γ is satisfied, then the proposal distribution is exactly symmetric, in the sense that the probabilities of reaching z ′ from z t − 1 and of reaching z t − 1 from z ′ are equal.
Then, if we accept the proposed spike with probability
[Equation]
we obtain a Metropolis-Hastings sampling algorithm for the discretized Gaussian [26, 49].
We now relax the assumption of the perfect integration, and assume only that η ≪ 1 .
Introducing decay smooths the readout, which in the perfect integrator is discretized.
With the same proposal distribution as before, we take the acceptance ratio of the accept-reject step to be
[Equation]
This choice has two important features.
First, the decay means that the proposal distribution will be asymmetric, and the Markov chain will no longer satisfy the condition of detailed balance [26, 49].
However, the resulting error will be small if η ≪ 1 .
Second, by comparing the likelihood of the proposal, P [(1 − η )
z t − 1 + Γe j ] , to the likelihood of the next state without the proposed spike but with the decay P [(1 − η )
z t − 1 ] (instead of the likelihood of the current state P [ z t − 1 ] as in the
Metropolis-Hastings algorithm), this choice implements a sort of look-ahead step that should allow the algorithm to partially compensate for the decay in the rate.
With the choices above, we show in Appendix B.2 that one can write the acceptance ratio (4) as
[Equation]
where V t − 1 = − (1 − η )
Ωr t − 1 + Γ ⊤ Ψ − 1 θ has the interpretation of a membrane potential,
[Equation]
has the interpretation of a spiking threshold, and the the recurrent weight matrix is defined as
[Equation]
Thus far, we have assumed that the mean signal is constant.
The natural generalization of this algorithm to a time-varying mean signal θ t is to take the membrane potential to be V t = − (1 − η )
Ωr t + Γ ⊤ Ψ − 1 θ t .
This leads to the voltage dynamics
[Equation]
which, when combined with the probabilistic spiking rule with uniform proposals and acceptance ratio (5), yields our final algorithm.
This will not be an exact Metropolis-Hastings sampler unless the mean is constant, the decay term vanishes ( η = 0 ), and the readout matrix Γ satisfies an exact balancing condition.
In particular, if these conditions are violated, the Markov chain will not satisfy detailed balance.
However, if they are violated only weakly, this algorithm can still be a reasonable approximation to a true sampler.
We will provide empirical evidence for this intuition in §3.
2.2 The continuous-time limit
We now consider the continuous-time limit of the model introduced above.
This limit corresponds to taking the limit in which spike proposals are made infinitely often, and regarding the dynamics written down previously as a forward Euler discretization of an underlying continuous-time system.
For a timestep ∆ between spike proposals, we let the discrete-time decay rate be η = ∆ /τ m for a time constant τ m , thusly named because it has the interpretation of a membrane time constant.
Then, we show in Appendix B.4 that the ∆ ↓ 0 limit of the discretized rate dynamics (1) yields the familiar continuous-time dynamics
[Equation]
In continuous time, the spike train o ( t ) is now composed of Dirac delta functions, as the discretized spikes are rectangular pulses of width ∆ in time and height 1 / ∆ .
We next consider the voltage dynamics of the leaky integrator for a varying mean signal (8), which have a similar continuum limit:
[Equation]
In this limit, the rate will decay by only a infinitesimal amount between a rejected spike proposal and the next proposal, meaning that the error incurred by neglecting the asymmetry in the proposal distribution due to the decay should be negligible.
We also note that, though the probabilistic spike rule (5) does not explicitly include a reset step, the dynamics (10) prescribe that the j -th neuron’s membrane voltage should be decremented by 2 T j after it spikes.
2.3 Efficient balanced networks as a limit of the spiking Metropolis-Hastings sampler
This spiking network samples the posterior distribution by emitting spikes probabilistically, but we can use the same architecture to re-derive EBNs, which approximate continuous dynamical systems using spiking networks [44–48].
If we take Ψ ∝ I n p (for I n p the n p × n p identity matrix)
the voltage dynamics (10) are identical to those of the EBN.
The greedy spiking rule of the EBN can be recovered in this framework by taking a limit in which the variance of the Gaussian target distribution vanishes.
Concretely, we let Ψ =
ψ
I n p , and define re-scaled variables ˜ V ( t ) ≡ ψ V ( t ) , ˜ Ω ≡ ψ Ω , and ˜ T ≡ ψ T that will remain O (1) even as we take ψ ↓ 0 .
In terms of these new variables, the acceptance ratio is A = min { 1 , exp[( ˜ V j ( t ) − ˜
T
j ) /ψ ] } , which tends to A =
Θ( ˜ V j ( t ) − ˜ T j ) as ψ ↓ 0 .
This explicitly recovers the greedy spike rule used in EBNs.
In Appendix B, we further analyze how the overall scales of Ψ and Γ affect the probabilistic spike rule.
The network with voltage dynamics (10) samples a distribution with mean θ and covariance Ψ .
Instead of sampling using the structured proposal distribution on spikes, this network can implement sampling through slowly varying Langevin dynamics on θ .
In the limit where the spike rule becomes greedy, this recovers the spiking sampler studied by Savin and Denève [20].
We will discuss this model further in §3.
Population geometry for fast sampling 3.1 Leveraging the geometry of inference to accelerate sampling
We first review recent work from the machine learning literature for how population geometry can be chosen to accelerate simple Langevin sampling, which establishes principles for the design of fast samplers.
For probability distributions belonging to the exponential family, including the Gaussian distributions on which we focus in this work, one can write the density P ( z ) in terms of an energy function U ( z ) such that P ( z ) ∝ exp[ − U ( z )] .
The classic algorithm to sample such a distribution is the discretization of the naïve Langevin dynamics
[Equation]
(11) where W ( t ) is a standard Brownian motion [28, 29, 32, 51].
By simply following these dynamics, one can obtain samples from the target distribution and therefore an estimate of the uncertainty at the timescale taken by the network to sufficiently explore the target distribution.
The Langevin dynamics (11) can be directly implemented in a rate network [9, 19] or approximately implemented in a spiking network [20], but their convergence properties for high-dimensional distributions have not been investigated in a neuroscience context.
It is well known in statistics that, as the dimensionality of the target distribution increases, convergence of Langevin sampling to the target distribution slows dramatically.
Furthermore, the discretization step can induce errors that cause the variance estimated from sampling to exceed the target variance [28, 29, 32, 51].
To overcome these issues, prior work in statistics and machine learning has proposed algorithms that can accelerate the sampling [27, 30, 41–43, 52–54] which we leverage here to propose our fast samplers in spiking neural networks.
These ideas were unified into a common framework by Ma et al.
[42] (and see also [43, 55]) who proposed a “complete recipe” for stochastic gradient MCMC:
[Equation]
(12) with BB ⊤ = D .
The matrix fields D ( z ) and S ( z ) modify the dynamics but keep the target distribution unchanged.
D is positive semi-definite and defines the local geometry of the space in which the inference is occurring, while S is skew-symmetric and adds non-reversible dynamics.
When D and S are state dependent, the correction term Φ = div( D + S ) must be included [42].
The “complete recipe” provides a general framework to design samplers based on Langevin dynamics.
Samplers based on Riemannian geometry can be designed by choosing D to be the inverse of the
Fisher information matrix G (or an approximation thereof), yielding the natural gradient ∇ nat U =
G − 1 ∇ U (Figure 1) [38–41, 56, 57].
Samplers incorporating dummy variables can be designed by expanding the parameter space and using the matrices D and S to obtain the desired dynamics [30, 42, 43, 52–54, 58].
Similarly, prior works have proposed methods to accelerate sampling in biologically inspired neural networks by parallelizing the inference [20], using Hamiltonian dynamics [21] or by adding non-reversible dynamics [9, 17, 18].
The “complete recipe” provides a general framework that encompasses all these examples, allowing for the principled design of biologically-plausible samplers.
In the simple setting of linear rate networks sampling from a Gaussian distribution with constant mean µ and covariance Σ , the complete recipe for state-independent D and S takes the form
[Equation]
In Appendix D, we show in the zero-mean case µ = 0 that large eigenvalues of the covariance matrix Σ introduce long timescales in the sample cross-covariance for naïve Langevin dynamics ( D = B
= I n p , S = 0 ), which slow convergence to the target distribution.
Concretely, the 2-Wasserstein distance between the ensemble sampling distribution at time t and the target is
[Equation]
(14) where σ i are the eigenvalues of the covariance matrix.
Networks performing inference in the natural space ( D = Σ , B = Σ 1 / 2 ) are insensitive to these large eigenvalues, with 2-Wasserstein distance
[Equation]
For this analytical result, we consider the distribution of samples at time t over realizations of the noise process, which differs from the distribution of samples over time within a single realization, as used elsewhere in the paper.
This choice is made because the joint distribution of samples within a single realization is challenging to characterize [59].
These results show how natural gradient can accelerate inference in linear rate networks, and complement Hennequin et al.
[9]’s study of how adding irreversible dynamics can reduce the sample autocorrelation time (i.e., D = B
= I n p ,
S ̸ = 0 ).
3.2 Fast sampling through population geometry in efficient balanced networks
We first consider a sampler based on efficient neural networks [20, 44–46] that leverages the geometry of the inference to implement efficient sampling at the level of the population.
In previous work,
Savin and Denève [20] derived a sampler implementing naïve Langevin dynamics (we provide a full derivation using the notation from the present paper in Appendix C).
Although they proposed to accelerate sampling by implementing parallel inference loops, they do not leverage the geometry of the inference nor do they test the convergence in high ( n p > 2 ) dimensions.
Here, we approximate the “complete recipe” dynamics (13) using an EBN, and show that performing inference using natural gradients helps with speed and accuracy in high dimensions.
As in §3, we consider sampling from a Gaussian distribution of possibly time-varying mean µ and covariance Σ .
4 As we principally study the effect of the geometry, we henceforth set S = 0 .
Note that even though the underlying dynamics of (13) are reversible if S = 0 , the non-linearities introduced by spiking lead to a non-reversible sampler.
To embed these dynamics in the EBN framework, we take the target mean in the voltage dynamics (10) to evolve according to the Ornstein-Uhlenbeck process (13), which we generalize to include an intrinsic timescale τ s .
Following [20], we approximate the true state of the sampling dynamics by the instantaneous estimate from the filtered spike history.
Then, as detailed in Appendix C, we obtain an EBN with the voltage dynamics
[Equation]
[Equation]
where the recurrent weight matrix is Ω = Γ ⊤ Γ , and the spiking rule is greedy, with thresholds
T j = ( Γ ⊤ Γ ) jj / 2 .
M ean 0 0 .
V a r i an c e 0 0 .
W 2 d i s t an c e 2 4 8 16 32 64 n p
V a r i an c e 2 4 8 16 32 64 n p
W 2 d i s t an c e 2 4 8 16 32 64 n p
M ean
[Figure]
i ii iii c d i ii iii 1 time (s)
M o v
i ng
[Figure]
e r
[Figure]
f i naivenatural  0.
[Figure] [Figure]
ISI (s
[Figure]
Variance and d.iii .
mean 2-Wasserstein distance.
In panels c,d shaded error patches show 95% confidence intervals computed via bootstrapping over 100 realizations in all panels; In Figure E.1 we show the full time course of the inferred statistics for one set of network parameters and in Figure E.
we show the statistics after the full 1 .
5 s, which are qualitatively similar.
See Appendix E for detailed numerical methods.
To illustrate how correlations between parameters affect sampling in high dimensions, we will focus on equicorrelated multivariate distributions.
The covariance matrix Σ of such a distribution is parameterized by an overall variance σ > 0 and a correlation coefficient − 1 < ρ < 1 such that Σ ij = σ [ δ ij + ρ (1 − δ ij )] .
We will explore the performance of the sampling algorithms across values of ρ for different dimensions of the parameter ( n p ) and neuron ( n n ) spaces.
For a multivariate
Gaussian distribution N ( µ , Σ ) , the Fisher information matrix is G = Σ − 1 and we will therefore use
D = G − 1 = Σ (and B = √ D ) in our geometry aware implementation (see Appendix C.3).
In our simulations, we compare the accuracy of sampling of the naïve implementation ( D = B = I ) with the geometry-aware version over a 50 millisecond window, which is roughly twice the membrane time constant τ m = 20 ms, as well as at steady-state.
In Figures 2, E.1 and E.3, we show that the geometry-aware sampler is more robust to increasing the correlation of the parameters and the dimensionality, allowing inference at behavioral timescales.
As in non-spiking Langevin sampling, discretization leads to an overestimation of variance [26], but this effect, although still present, is strongly reduced in the geometry-aware implementation.
Note here, that the geometry is imposed through the dynamics of the membrane potentials via the recurrent connectivity in the network.
3.3 Fast sampling through population geometry using probabilistic spike rules
In the preceding subsection, we have shown how neural population geometry, implemented through neural dynamics, can accelerate the speed of approximate sampling in an efficient balanced network.
However, this approach suffers from a fundamental conceptual gap: the firing rates of the spiking network are being used to simulate non-spiking Langevin dynamics; the spiking network itself has not been designed to sample.
Specifically, the discretization introduced by the spiking exacerbates the errors introduced by the discrete time implementation of the sampling dynamics, leading to
[Figure]
ii
[Figure]
overestimation of the stimulus variance.
In this subsection, we take a first step towards bridging this gap by considering an alternative limit of the general model introduced in §2: the case in which sampling is performed leveraging only the stochasticity in the spike rule.
Our objective here is not to demonstrate a fully biologically-plausible or practically useful sampling algorithm; rather, it is to illustrate the importance of population geometry in a minimal model for probabilistic spiking.
As in 3.2, we focus on sampling from equicorrelated Gaussian distributions N ( µ , Σ ) .
In this case, we set Ψ = Σ and θ = µ in (10).
Naïve sampling here corresponds to choosing Γ in some sense generically (see Appendix B.5), while geometry-aware sampling corresponds to choosing Γ
such that ΓΓ ⊤ ≃ Σ up to overall constants of proportionality
In Figures 3, 4, and E.4, we show that naïve choices of Γ lead to vanishing spike rates at strong correlations ρ and large parameter-space dimensionalities n p .
This results in dramatic underestimation of the mean and variance of the target distribution, which is resolved by choosing the geometry appropriately, again allowing inference at behavioral time-scales.
Moreover, these networks show Poisson-like variability in spiking statistics (Figure 3), consistent with cortical dynamics [60, 61].
In Appendix B.5, we provide a more careful analysis of the strongly-correlated limit ρ ↑ 1 .
Informally, we show that the probability of spiking should vanish if Γ is chosen sufficiently naïvely and the mean of the target distribution is uniform across parameter dimensions, i.e., µ t ∝ I n p .
This analysis is not specialized to a particular case, and holds generally for the model of §2.
Therefore, careful choice of population geometry, as implemented by the neural code, is required for fast sampling in this model.
Discussion
In this paper, we have shown how careful choice of neural population geometry enables fast sampling in spiking neural networks.
We presented a unified framework in which EBN samplers approximating
Langevin dynamics with greedy spiking and approximate Metropolis-Hastings samplers with deterministic voltage dynamics and probabilistic spiking can be unified.
We then leveraged population geometry to perform rapid sampling at behaviorally-relevant timescales in these two disparate limits of our general model.
We now discuss some of the limitations of our work, and highlight possible directions for future inquiry.
ρ a b a v e r age f i r
i ng r a t e (
H z )    naïvenatural 0.25 0.5 0.
s ho r t - t i m e m ean 1.
s ho r t - t
i m e v a r i an c e 1.
s ho r t - t i m e W101010ρ
[Figure]
[Figure]
i ii iii iv i ii iii iv a v
[Figure]
r age f i r i
ng r a t e ( k
H z )
[Figure]
s ho r t - t
i m
[Figure]
m ean
[Figure]
s ho r t - t
i m e
[Figure]
r i an
[Figure]
e
[Figure]
ho r t - t
i m
[Figure]
n p
[Figure] [Figure]
n p
[Figure]
[Figure] [Figure]
n
[Figure]
648 16n p
[Figure] [Figure]
naïvenatural
[Figure] [Figure]
[Figure] [Figure]
requires careful choice of geometry.
b.i
At moderately strong correlations ρ =
75 and high dimensions spiking is suppressed in networks with naïve geometry, but not when the natural geometry is used.
Here, we use 10 neurons per parameter, i.e estimate of mean signal in the after
See
Like the original EBN model, the probabilistic spiking model introduced in §2 suffers from the limitations that it requires instantaneous propagation of spike information and that only one neuron is allowed to spike at a time [20, 44–48].
Moreover, the discretization timestep enforces a hard cutoff on the maximum spike rate.
Some of these limitations could be partially circumvented by generalizing the spike proposal distribution to allow multiple spikes.
However, such a model would still suffer from the issue that an accept/reject step that accounts for the effect of spikes from multiple neurons will require instantaneous communication across the network.
This limitation could possibly be overcome within the framework of asynchronous Gibbs sampling [62, 63], which ignores the requirement that updates should be coordinated across the network.
The analysis of §2 shows that the models of §3 can be viewed as limiting cases of a single framework.
It is likely that the parallels between these limiting models could be further strengthened by viewing the Gaussian noise in the voltage dynamics of the EBN sampler as an approximation to the effect of the stochastic spiking of other neurons in the Metropolis-Hastings sampler.
Such an approximation could be obtained in the limit of large network size given an appropriate treatment of asynchronous updates, provided that one could neglect the coupling between spikes in different neurons induced by the proposal distribution [61, 64–66].
Careful analysis of the relationship between these two sources of stochasticity will be an interesting subject for future investigation.
Moreover, it will be interesting to investigate how they might be integrated in a single network, which would result in a spiking sampler somewhat reminiscent of the Metropolis-adjusted Langevin algorithms used in machine learning [67, 68].
Here, we have focused entirely on Gaussian target distributions.
As we sketch in Appendices B.8 and
C.4, both the probabilistic spiking sampler and the EBN sampler could in principle be extended to general exponential family targets.
However, the most naïve extensions to non-Gaussian distributions would involve non-linear and non-local interactions in the voltage dynamics.
In particular, one would have to account for state-dependence in the matrix fields D
( z ) and S ( z ) , which would require the inclusion of the field
Φ ( z ) in the complete recipe (12) (see ref.
Mapping such interactions onto biological mechanisms requires more careful consideration.
In recent work, Nardin et al.
[69] have proposed how dynamical systems with polynomial nonlinearities can be approximated by deterministic EBNs with multiplicative synapses.
This approach could be extended to the stochastic setting of networks designed to sample, which will be interesting to test in future work.
In this work, we did not constrain neurons to follow Dale’s law, and single neurons therefore have both excitatory and inhibitory effects on their neighbors.
Many frameworks have been proposed to map the connectivity of unconstrained network algorithms onto distinct excitatory and inhibitory neuron types [9, 70, 71].
These refinements of the biological plausibility will not affect our key argument of accelerating inference through a favorable population geometry.
However, different possible implementations that comply with Dale’s law will make different predictions for experimentallymeasureable biophysical properties.
For example, although less numerous, fast spiking inhibitory neurons have higher firing rates [70, 71], which could allow the approximate symmetry of the readout, as required by the construction of §2, to be maintained.
In biological spiking networks, probabilistic spike emission and probabilistic synaptic release are natural sources of stochasticity [72–75].
These two layers of probabilistic computation provide additional flexibility in processing beyond the simple accept/reject step considered here.
As a result, it is likely that one could construct sampling algorithms that are at once more biophysically detailed and more computationally efficient than the simple network constructed in §2.
Further investigation of how violations of detailed balance through these mechanisms and the matrix S in the “complete recipe" framework could enable faster sampling will be a particularly interesting objective [17, 18, 50].
The algorithms proposed in this work could also enable fast sampling in neuromorphic circuits [76–78].
As they require only the local membrane voltage to compute accept/reject steps, these algorithms could be implemented in a distributed neuromorphic architecture.
They would then potentially be limited only by the timescale of individual computing units (which are much faster than biological neurons) rather than the dimensionality of the inference problem [79].
The sampling processes considered in this work focus on short-timescale perceptual inference.
However, similar probabilistic inference can occur at longer timescales of learning [80–82].
These long-timescale sampling processes would allow networks to flexibly infer their synaptic weights.
Importantly, learning the task—i.e., learning the matrix Σ − 1 —and learning the representational geometry— D , S , and Γ —are distinct processes that can occur in parallel.
Once the task structure is learned, adapting the geometry of the population code to match changing stimulus geometry can be achieved through meta-learning [83].
Recent works have analytically studied the population geometry that results from this sampling procedure in rate-based networks [84], and proposed algorithms for efficient learning in EBNs [46, 85–87] and other classes of spiking networks [88–91].
However, the interactions between fast activity sampling and slow network parameter sampling in neural circuits—particularly spiking networks—remain poorly understood [82].
Characterizing how adaptive population geometry accelerates perceptual inference in dynamic environments will be an important step towards a more complete understanding of probabilistic inference in neural circuits.
Acknowledgments and Disclosure of Funding
PM was supported by the Harvard Mind Brain Behavior Interfaculty Initiative.
JAZ-V and CP were supported by a Google Faculty Research Award and NSF Award DMS-2134157.
This work was also supported by a grant from the National Institute of Health (R01-DC-017311) to VNM.
We thank
Tony Zador for co-organizing the NeurIPS 1996 workshop on "Synaptic transmission: reliability and variability", whose follow-up discussions led to this paper.
A Table of variables and parameters used in the models and simulations
We provide a table stating the dimensionality and values taken by the variables and parameters used throughout the paper.
[Table]
B Approximate Metropolis-Hastings sampling using probabilistic spiking rules
In this Appendix
[Table]
we provide a step-by-step construction of the spiking sampler introduced in §2 of the main text.
As introduced in §2 of the main text, our goal is to use z t =
Γr t (B.1) to sample a Gaussian distribution P ( z ) with mean θ and covariance Ψ , where r t = (1 − η ) r t − 1 + o t (B.2) is the filtered spike history for a decay constant 0 ≤ η ≤ 1 (note that o t ∈ { 0 , 1 } n n ).
In §B.1, we construct a circuit that samples a Gaussian distribution using a discrete-time perfect integrator of the spike train (i.e., η = 0 ).
In §B.2, we relax this assumption, yielding an approximate sampler using leaky integration in discrete-time, and discuss the behavior of this model in the continuum limit.
In general, the proposal distribution could be computed using a stochastic gradient step followed by an accept/reject step, yielding Metropolis-adjusted Langevin dynamics [67, 68].
In this case, the accept/reject step allows the algorithm to compensate for some of the sampling error introduced by discretization at the expense of needing to compute a likelihood ratio, which can be expensive for high-dimensional distributions.
S
B.1 A simple sampler assuming perfect integration and balance
We first construct a sampling circuit under the assumptions that we have access to a perfect integrator of the spike train, i.e. that η = 0 , and that the readout matrix is of the form
[Equation]
for some matrix M ∈ R n p × n n / 2 .
At the t -th timestep, we choose one neuron j uniformly at random, and let the spike proposal be o ′ =
e j , where e j is the j -th standard Euclidean basis vector (i.e., ( e j ) k
= δ jk ).
This yields a candidate readout
[Equation]
Under the symmetry assumption on Γ , the acceptance ratio is given by
[Equation]
as the proposal distribution is exactly symmetric in z ′ and z t .
Then, we accept the proposed spike with probability A .
Concretely, for u ∼ U [0 , 1] , we take
[Equation]
For P ( z ) a Gaussian distribution with mean θ and covariance Ψ , we have
[Equation]
Defining the matrix
[Equation]
we interpret the first term in the log-odds ratio as a membrane potential,
[Equation]
and the second as a threshold
[Equation]
such that the acceptance ratio is
[Equation]
With this definition, the membrane voltage simply integrates the spike train:
[Equation]
We now observe that, if n n > n p , it is possible that a spike may not contribute to the parameter estimate.
Concretely, we say that a spike, or the corresponding neuron, is irrelevant if it is annihilated by Γ , i.e., if Γe j = 0 .
Irrelevant spike proposals are always accepted with probability one, as we have z ′ =
z t − 1 .
Moreover, the membrane voltage is not changed by the emission of irrelevant spikes.
For these reasons, we are free to re-define the population rate r t to exclude such spikes.
Therefore, a timestep with an irrelevant spike is equivalent to not updating the network at all, and we could choose to re-define the network such that only relevant neurons are included.
However, though there will exist some non-trivial set of vectors that are annihilated by Γ if n n > n p , the situation in which this null space is axis-aligned (thus implying the existence of irrelevant neurons) is not generic.
S
We note that appropriate initialization of the membrane potential (for the desired mean) is important, as otherwise some bias will be introduced.
Thus, as written, this model cannot easily accommodate a time-varying mean signal θ t .
This shortcoming could obviously be addressed by taking
[Equation]
which yields voltage dynamics
[Equation]
Alternatively, one could also take
[Equation]
and re-define the threshold for the j -th neuron to be time-varying:
[Equation]
The former of these approaches is reasonable from a biological perspective, as the new term Γ ⊤ Ψ − 1 ( θ t − θ t − 1 ) in the voltage dynamics has the interpretation of a signal θ t − θ t − 1 fed through an input weight matrix Γ ⊤ Ψ − 1 .
For this sampling procedure to work, the mean should be slowlyvarying.
B.2 Relaxing the assumption of perfect integration
We now relax the assumption of the perfect integration, and assume only that η ≪ 1 .
With the same proposal distribution as before, we take the acceptance ratio of the accept-reject step to be
[Equation]
As noted in the main text, this choice implements a sort of look-ahead step.
Moreover, if we took the acceptance ratio to depend on the likelihood of the proposed state with decay, P [(1 − η ) Γr t − 1 + Γe j ] , relative to the likelihood P [ Γr t − 1 ] of the current state (rather than the next state with decay but without the proposed spike), the resulting log-odds ratio would include terms of order η that are quadratic in the rate r t .
With the choices above, the log-odds ratio is
[Equation] [Equation]
which, like in the perfect integrator model, is linear in the rate.
As in the perfect integrator model, we define the recurrent weight matrix
[Equation] [Equation]
and interpret the first term in the log-odds ratio as a membrane potential,
[Equation]
and the second as a threshold
[Equation]
such that the acceptance ratio is
[Equation]
With this definition, the membrane voltage evolves as
[Equation]
Therefore, this model differs from the perfect integrator model of §B.1 only in the voltage dynamics; the perfect integrator is recovered exactly if we set η = 0 .
As in the perfect integrator model, the natural generalization of these leaky dynamics to time-varying mean signal θ t is to take
[Equation]
which leads to the dynamics
[Equation]
This will, of course, not be an exact sampler unless the mean is constant.
If the mean is slowly-varying, however, it should be a good approximate sampler.
We remark that we can re-write the membrane voltage as
[Equation]
S
B.3 Adding an elastic net prior on the rates
We now consider the effect of adding an elastic net prior on the rates, as was considered in the original work of Boerlin et al.
[Equation]
Without loss of generality
[Equation]
we consider the case in which a decay term is included, as we can then recover the perfect integrator by setting η = 0 .
Again defining the acceptance ratio in terms of a comparison against the rate with decay but without the proposed spike, the addition of the prior adds
[Equation]
to the log-odds ratio.
This yields a modified recurrent weight matrix Ω
[Equation]
and a modified threshold
[Equation]
but the expression for the membrane voltage in terms of these parameters is identical in functional form:
[Equation]
Therefore, adding the elastic net prior changes the definitions of the weight matrix that maps rates to voltages and of the threshold, but not the overall form of the result, hence it does not introduce any new conceptual difficulties.
We remark that we could include the constant factor α in the membrane voltage as we do in our implementation of EBNs (see Appendix C) rather than in the threshold, which would give it a somewhat different biological interpretation but would have no algorithmic effect.
B.4 The continuous-time limit
In this subsection, we consider the continuous-time limit of the models introduced above.
This limit corresponds to taking the limit in which spike proposals are made infinitely often, and regarding the dynamics written down previously as a forward Euler discretization of an underlying continuous-time system.
For clarity, we write the discrete timesteps, denoted in previous sections simply by t , as t d here, and reserve the unsubscripted symbol t for the continuum variable.
For a timestep ∆ , we let t =
∆ t d , and let the discrete-time decay rate be η = ∆ /τ m for a ‘membrane’ time constant τ m .
We may then write the discretized rate dynamics (B.2) as
[Equation]
which, taking ∆ ↓ 0 , of course yields the familiar continuous-time dynamics
[Equation]
In continuous time, the spike train o ( t ) is now composed of Dirac delta functions, as the discretized spikes are rectangular pulses of width ∆ in time and height 1 / ∆ .
We next consider the voltage dynamics of the leaky integrator for a varying mean signal (B.26), which may similarly be re-written as
[Equation]
S
In the continuum limit, we retain only the contribution of the first of the two terms involving the discrete-time spike train, as the other yields pulses of width ∆ and height unity, which yield a negligible contribution to the integral.
Thus, we have the dynamics
[Equation]
From these dynamics, one could then recover the continuum limit of the perfect integrator dynamics by taking τ m → ∞ .
In this limit, the rate will decay by only a infinitesimal amount between a rejected spike proposal and the next proposal, meaning that the error incurred by neglecting the asymmetry in the proposal distribution due to the decay should be negligible.
B.5 Analyzing the strongly-correlated limit of sampling from an equicorrelated Gaussian
Here, we consider how these models behave when sampling from an equicorrelated Gaussian distribution, i.e., a distribution with covariance matrix
[Equation]
for correlation coefficient ρ ∈ ( − 1 , +1) .
We are particularly interested in the strongly-correlated limit ρ ↑ 1 .
As we will consider the case in which the marginal variance does not scale with the correlation, our choice of unit marginal variance is made without loss of generality.
For this covariance matrix, the Sherman-Morrison formula yields [92]
[Equation]
We first consider the case in which the mean signal θ is identically equal to zero.
We argue that, for choices of Γ that are in some sense sufficiently naïve, the probability that relevant spikes are emitted should tend to zero as ρ ↑ 1 .
This corresponds to showing that V t,j − T j → −∞ as ρ ↑ 1 for all indices j corresponding to relevant neurons.
For θ = 0 , we have
[Equation]
Consider the first timestep, with r 0 = 0 , for which we have
[Equation]
Near ρ = 1 , we then have the expansion
[Equation]
under the assumption that Γ is an O (1) function of ρ in this region, and thus cannot introduce additional possible divergences.
For relevant spikes, we have the strict inequality ( Γ ⊤ Γ ) jj > 0 .
Moreover, by the Cauchy-Schwarz inequality, we have ( 1 ⊤ n p Γe j ) 2 ≤ n p ( Γ ⊤ Γ ) jj , with equality if and only if Γe j ∝ 1 n p .
If Γe j ∝ 1 n p , then spikes in neuron j affect the readout precisely only along the common mode, and V 0 ,j − T j does not diverge as ρ ↑ 1 .
However, this case is quite fine-tuned.
For generic Γ not satisfying this alignment condition, we have the strict inequality
[Equation]
for all relevant neurons.
Then, if ( Γ ⊤ Γ ) jj vanishes no more rapidly than 1 − ρ as ρ ↑ 1 —i.e., if ( Γ ⊤ Γ )
jj / (1 − ρ ) → ∞ as ρ ↑ 1 —we have that V 0 ,j − T j → −∞ as ρ ↑ 1 .
This holds, for instance, if ( Γ ⊤ Γ )
jj is a constant function of ρ .
Thus, under these conditions, the probability that a relevant spike is emitted at the first timestep vanishes as ρ ↑ 1 .
Heuristically, this in turn implies that the (relevant component of the) rate at the second timestep will be zero with probability one.
Therefore,
S we may iterate this argument forward in time, showing that the probability of emission of relevant spikes should vanish in the limit ρ ↑ 1 .
This argument will not be affected by the addition of an elastic net penalty unless the coefficients α and λ are taken to diverge as ρ is taken to unity, as the coefficients are strictly non-positive.
The situation is somewhat more complicated if the mean signal is not identically zero, in which case we have
[Equation]
Following our previous analysis, at the first timestep we have
[Equation]
B
[Equation]
47) near ρ = 1 .
There are now two possible divergent terms, which can compete to change the sign of
V 0 ,j − T j as ρ ↑ 1 .
One case of interest is when θ t ∝ 1 n p .
Then,
[Equation]
and we have
[Equation]
as in the case when θ was strictly zero, implying that V 0 ,j − T j → −∞ under the abovementioned conditions on Γ .
Another illustrative case is θ t =
Γe
[Equation]
With this fine-tuning,
[Equation]
hence we expect V 0 ,j − T j → + ∞ as ρ ↑ 1 under the abovementioned constraints on Γ .
Thus, in this case, the spike probability should tend to one, showing that complications can arise in the case of a non-uniform mean signal.
B.6 Analyzing the small- and large-variance limits
We now consider how these models behave in the limits in which the variance of the target distribution is small or large, or, nearly equivalently, the limits in which the scale of the readout matrix is large or small, respectively.
In §2.3 of the main text, we argued that the EBN is recovered in the limit in which the variance of the target distribution tends to zero.
Concretely, we let Ψ = ψ ˜ Ψ for some fixed matrix ˜ Ψ , and take the zero-variance limit ψ ↓ 0 .
In terms of the renormalized voltage ˜ V ≡ ψ V , weight matrix ˜ Ω = ψ Ω , and threshold ˜ T = ψ T , the spike acceptance ratio is
A = min { 1 , exp[( ˜ V j ( t ) − ˜ T j ) /ψ ] } .
This tends to A =
Θ
( ˜ V j ( t ) − ˜ T j ) as ψ ↓ 0 assuming that the re-scaled variables remain order one, recovering the greedy spiking rule used in the EBN.
If we instead take the large-variance limit ψ ↑ ∞ , the acceptance ratio tends to 1 under the assumption that ˜ V j ( t ) − ˜ T j remains O (1) .
Then, each spike proposal is accepted with probability one, and the total firing rate of the population is 1 / ∆ spikes per second for a timestep ∆ , meaning that the population-averaged rate is 1 / ( n n ∆) spikes per second.
Re-scaling the readout matrix Γ differs from re-scaling the target covariance matrix Ψ because the target mean and recurrent spikes appear with the same power of Ψ (in particular, Ψ − 1 ) but different powers of Γ in the voltage dynamics, scaling as O ( Γ ) and O ( Γ 2 ) , respectively.
Thus, in the large- Γ limit in discrete time, the contribution of the target mean should be negligible relative to that of the recurrent spiking input, while the opposite should hold in the small- Γ limit.
However, these different scalings should not affect the limiting behavior of the acceptance ratio.
In Figure B.1, we probe how sweeping the scale of Γ over five orders of magnitude affects sampling from a Gaussian distribution of fixed variance.
We show that the population-averaged spike rate tends to 1 / ( n n ∆) spikes per second in the small- Γ limit, while the spike rate tends to zero in the large- Γ limit.
S a a v e r age f i r
i ng r a t e ( k
H z ) 1.
s ho
r t - t i m e m ean 1.
s ho r t - t
i m e v a r i an c e 1.
s ho r t - t
i m e Wi ii iii iv naïvenatural variance of random Γ ij 1010 2 10 0 10 -2 10 -4 10n p
= 10, n n =b a v e r age f
i r i
ng r a t e (
H z )  s ho r t - t
i m e m ean 1.
s ho r t - t
i m e v a r i an c e 1.
s ho r t - t
i m e Wi ii iii iv naïvenatural variance of random Γ ij 1010 2 10 0 10 -2 10 -4 10n p
= 32, n n =10101010101010101/( n n Δ ) 1/( n n Δ )
B.7 Alternative spike proposal distributions and membrane voltage bounds
In this appendix, we have considered the simplest possible spike proposal distribution: at each timestep, we choose one neuron uniformly at random as a candidate.
With this choice, obviously, only one neuron can spike at each timestep, which parallels the spiking rule used in the EBN: only the neuron with the maximum membrane voltage among the entire population is allowed to spike [20, 44].
With a discretization timestep ∆ , the maximum total spike rate is then 1 / ∆ .
In the EBN literature, previous analysis has shown that this constraint is vital to avoid pathological spiking patterns [20, 44–47], and introduced alternatives such as Poisson spiking rules [47] or hand-tuned refractory periods [48].
We remark that, in these models, the membrane voltage is always strictly bounded from above, which does not hold in our setting.
As reset is achieved only through spike emission, and is strictly speaking a decrement of the voltage rather than a true reset, a neuron can exceed the threshold voltage if it is not chosen as a candidate to spike.
B.8 Sampling from non-Gaussian exponential families
Though our main focus is on Gaussian target distributions, in this appendix we briefly discuss the possibility of constructing a probabilistic spiking sampler for other exponential families.
For a distribution with density
[Equation]
B.51) for an energy function U , the acceptance ratio (B.18) becomes
[Equation]
In analogy to the Gaussian case, one could then define the difference of energies to be the difference of the membrane voltage and the threshold.
However, the resulting membrane voltage would be a nonlinear function of the firing rate, and would in general evolve according to non-linear dynamics.
S
Most directly, assuming ∥ Γe j ∥ to be small and U to be not too quickly varying, one could make the second-order approximation
[Equation]
and define the membrane voltage and threshold as
[Equation]
and
[Equation]
respectively.
With this choice, the threshold would be state-dependent, but one could incorporate its state-dependence into a re-defined voltage.
C Sampling in efficient balanced networks C.1 Encoding a dynamical system in efficient balanced networks
In this section, we provide a pedagogical derivation of the efficient balanced network [20, 44] using the notation we use throughout the paper.
The goal of the network is to encode an estimate z ( t ) of a signal θ (a vector of size n p × 1 in a population of n n spiking neurons.
The estimate is obtained by reading out a low pass version of the population spiking activity:
[Equation]
where Γ is the n p × n n readout matrix and r ( t ) is the low-pass filter spike history:
[Equation]
where τ m is the time constant of the readout neurons and o i is the spike train of neuron i , o i
= 1 if the neuron spiked at t
and o i = 0 otherwise.
The goal of the network is to minimize the squared error between the signal and the estimate with an elastic net prior on the firing rate in order to find a good solution while keeping the population spiking activity relatively low:
[Equation]
In the standard implementation of efficient balanced networks, neurons use a greedy spiking rule.
A neuron should fire if emitting a spike will lower the loss function,
i
[Equation]
e., if
[Equation]
Using the loss function and the definition of the estimate ( z ( t ) =
Γr ( t ) )
[Equation]
we can rewrite the spiking rule for neuron i as:
[Equation]
Removing terms that appear on both sides
[Equation]
we get
[Equation]
which we can further simplify to
[Equation]
S
The term on the left hand can be interpreted as the voltage potential of neuron i :
[Equation]
while the term on the right hand size can be interpreted as the firing threshold 5 :
[Equation]
If the voltage potential exceeds this threshold, then the neuron will fire and lower its voltage potential back below threshold.
Using (C.8) we can express the dynamics of the voltage potential as a function of the dynamics of θ and its estimate z :
[Equation]
We can rewrite this equation to obtain an expression for the membrane dynamics as a function of the signal, the firing rates, and the spike trains of the neurons.
By adding and subtracting 1 τ m Γ ⊤ i θ − 1 τ m α and noting that Γ ⊤ i Γr + λr i − Γ ⊤ i θ + α = − V
i , we further simplify the expression to obtain membrane dynamics as a function of the membrane potential, the effect of a new spike on the circuit and the encoded dynamical system:
[Equation]
We therefore obtain the membrane dynamics for the efficient balanced spiking network proposed in [20, 44].
We can rewrite the dynamics in vector form:
[Equation]
with Ω =
( Γ ⊤ Γ
+ λ I n n ) .
(C.14)
Using this scheme one can encode of variety of dynamical systems, including Langevin dynamics [20, 45, 46].
C.2 Sampling in efficient balanced networks using naïve Langevin dynamics
Using the scheme presented in §C.1 above, Savin and Denève [20] proposed to implement a dynamical system corresponding to the naïve Langevin dynamics of a multivariate normal distribution.
We use a linear Gaussian model similarly to several studies of neuroscience inspired sampling-based networks [9, 20].
These networks estimate the posterior probability of hidden sources ( θ ) given sensory inputs x corrupted by Gaussian noise, p ( x | θ ) =
N ( x ; Aθ , σ 2 n I ) , and prior expectations on the values of the hidden sources p ( θ ) =
N ( θ ; 0 , C ) .
The mean µ and covariance Σ of the posterior probability of the features given the input, p ( θ | x ) ∝ p ( x | θ ) p ( θ ) , are µ = 1 σ 2 n Σ A ⊤ x and Σ = (cid:16) C
+ 1 σ 2 n A ⊤ A (cid:17) − 1 respectively.
Up to an irrelevant constant offset, the corresponding energy function is U ( θ )
=
12 ( θ − µ ) ⊤ Σ − 1 ( θ − µ ) , and its gradient is ∇ U
( θ ) =
Σ − 1 ( θ − µ ) .
We define τ s as the timescale of the inference process and we set ϵ t = 1 τ s .
Then, we can write down the naïve
Langevin dynamics from (11):
[Equation]
regularizing term α as a fixed offset in the voltage potential but it can equivalently be included as an offset in the spiking threshold T i , as discussed in Appendix B.3.
S
These dynamics can be approximated by an efficient balanced network by replacing in (C.14) θ by z .
As discussed in [20], this approximation will introduce an acceptable error in most practical situations:
[Equation]
and we therefore obtain the membrane dynamics for the efficient balanced network proposed by
Savin and Denève [20]:
[Equation]
Note here that we have two timescales: the timescale of neuronal representations τ m ∼ 20 ms, controlled by the biophysical properties of the neurons, and the timescale τ s of the Langevin diffusion encoded by the network.
C.3 Sampling in efficient balanced networks using the complete recipe for stochastic gradient MCMC
The model proposed by Savin and Denève [20] implements naïve Langevin dynamics, which are known to be slow in high dimensions.
Instead, we can use the “complete recipe” for stochastic gradient MCMC [42] to write another sampler with the same equilibrium distribution but more favorable convergence properties.
For any positive semi-definite matrix D and skew-symmetric matrix S , the following dynamics will converge to N ( µ , Σ ) as their stationary distribution:
[Equation]
with BB ⊤
= D .
We can use (C.14) to encode this more general formulation into an efficient balanced network, yielding the following membrane dynamics:
[Equation]
Any choice of positive semi-definite matrix D leads to a valid sampler, but extensive work inspired by Amari’s seminal work on natural gradient descent [38, 39] has shown that a principled choice is to take D to be the inverse of the Fisher information matrix G [41, 42].
For the multivariate Gaussian distribution, the Fisher information matrix is given by:
[Equation]
We should therefore choose D =
G − 1 = Σ .
Note that for the multivariate Gaussian distribution, the
Fisher information matrix is identical to the Hessian and is location independent - it does not depend on the value of θ .
For more complex distributions, the Fisher information matrix might be difficult to compute and an approximation can be used as long as it is valid (positive semi-definite) within the complete recipe framework [42, 57, 93] and state dependent matrices can be corrected for using the term Φ from the complete recipe in (12) [42].
In this work, we have considered only hand-tuned or random choices for the matrices controlling the geometry.
Previous work by Hennequin et al.
[9], when framed within the complete recipe, proposes methods to find a skew-symmetric matrix S which accelerates the dynamics.
Non-reversibility is indeed known to accelerate learning, but analysing networks with such dynamics is notoriously difficult [27, 94].
In contrast, approximations for the inverse Fisher information matrix are readily computable, even in non-Gaussian settings [57, 93].
C.4 Sampling from non-Gaussian exponential families using efficient balanced networks
To sample from a non-Gaussian exponential family distribution with energy U ( θ ) using an EBN, we can simulate the “complete recipe” for a general exponential-family as stated in (12) of the main text, which we reproduce below:
[Equation]
where BB ⊤ = D [13].
Here, we have introduced an auxiliary time constant τ s , as in our previous discussion of the Gaussian case.
Then, the EBN voltage dynamics (C.14) become
[Equation]
where we have again made the approximation of the Langevin sampling trajectory by the approximate sampling trajectory z
= .
It is easy to see that these dynamics will in general involve non-linear dependence on the population firing rate r through the readout z = Γr .
D Natural gradient enables fast sampling in linear rate networks
In this section, we analyze how natural gradients enable fast sampling in rate networks designed to sample from zero-mean Gaussian distributions.
Our starting point is the complete recipe for sampling from a zero-mean Gaussian distribution with covariance Σ :
[Equation]
where D is a symmetric positive-semidefinite matrix and S is skew-symmetric.
In previous work,
Hennequin et al.
[9] showed how these dynamics can be interpreted as a linear rate network, and demonstrated that careful choice of S can accelerate sample autocorrelation timescales in networks with D
= I n p .
Here, our objective is to show how the geometry of inference, set by D , can accelerate sampling in these rate networks.
Our analysis is a straightforward application of the classic theory of Ornstein-Uhlenbeck processes [28, 51, 95].
Assuming for simplicity a deterministic initial condition z (0) = 0 , the solution to this stochastic differential equation is given by the Itô integral
[Equation]
which has ensemble mean zero and ensemble covariance
[Equation]
The ensemble distribution of z ( t ) is of course Gaussian, with mean zero and covariance given by the equal-time covariance function C ( t ) ≡ C ( t, t ) .
We assume that ( D + S ) Σ − 1 is a non-defective matrix with all eigenvalues having positive real part, such that the process has a well-defined stationary state.
By construction, the covariance matrix of the stationary state is precisely the target covariance matrix Σ .
In the stationary state, we have a simplified two-point function
[Equation]
where we have observed that
[Equation]
S
One standard case in which the integral defining the non-stationary covariance function can be evaluated explicitly is if A = ( D + S ) Σ − 1 is a normal matrix (i.e., if AA ⊤ =
A ⊤ A ) [28].
Then, there exists a unitary matrix U such that UU † =
I n p and
[Equation]
which, for all t ≥ t ′ , yields
[Equation]
for
[Equation]
We will measure the rate of convergence of the ensemble distribution to the stationary distribution in the Kullback-Leibler (KL) divergence and 2-Wasserstein distance for several choices of D .
We note that this is not the same measure as considered in our numerical simulations, where we examine the distribution of samples over time within a single trajectory, and measure the mean 2-Wasserstein distance between univariate marginals.
As noted in Appendix E, an organism must usually make estimates based on the distribution of samples over a single trajectory.
However, even at equilibrium, it is challenging to analytically characterize the 2-Wasserstein distance between samples from a
Gaussian distribution and the underlying population distribution [59].
Therefore, we will consider the ensemble distribution at a given time.
We denote the KL divergence and 2-Wasserstein distances as a function of time by K ( t ) and W 2 ( t ) , respectively.
As all distributions of interest are Gaussian, we have the relatively simple formulas
[Equation]
and
[Equation]
D.1 Naïve Langevin dynamics
We first consider naïve Langevin sampling with D =
I n p and S = 0 .
In this case, A =
( D +
S )
Σ − 1 = Σ − 1 is symmetric and therefore normal, hence, letting the diagonalization of Σ be
[Equation]
the result above yields
[Equation]
for [
[Equation]
15) for all t ≥ t ′ .
In particular, equal-time covariances are governed by
[Equation]
In matrix form,
[Equation]
) and
[Equation]
This choice yields stationary covariance
[Equation]
hence large eigenvalues of Σ will introduce long autocorrelation timescales.
S
In this case, K ( t ) and W 2 ( t ) are simple to compute thanks to the fact that Σ commutes with C ( t ) , yielding
[Equation]
and
[Equation]
In terms of the eigenvalues σ 1 , . . .
, σ n of Σ , these distances are
[Equation]
and
[Equation]
D.2 Sampling in the space of natural parameters
We now consider sampling in the space of natural parameters, with D = Σ and S = 0 .
In this case,
A = DΣ − 1 =
I n p is trivially normal, hence, for t ≥ t ′ ,
[Equation]
D.24) with
[Equation]
in particular.
In this case, the stationary covariance is
[Equation]
26) hence the stationary autocorrelation timescale will be independent of the spectrum of Σ .
This setup can also easily be generalized to the case in which S ̸ = 0 .
In this case, A =
( D + S ) Σ − 1 =
I n p + SΣ − 1 is not in general a normal matrix.
However, we can evaluate the covariance function by exploiting the particular structure of the problem.
Factoring out the terms involving the identity matrix and using the skew-symmetry of S , we have
[Equation]
We now observe that
[Equation]
hence e −
[Equation]
t
[Equation]
s
[Equation]
1 ( t − t ′ ) Σ , and therefore
[Equation]
The equal-time covariance is then
[Equation]
as in the case S = 0 .
Thus, the rate of convergence of the ensemble sampling distribution to the stationary distribution will be unaffected if we add this skew-symmetric term.
With the addition of the skew-symmetric term, the stationary covariance is
[Equation]
which reflects the fact that this term introduces non-reversible dynamics.
The spectrum of this matrix is, however, not easy to analyze in general.
Once again, K ( t ) and W 2 ( t ) are easy to compute thanks to the fact that C ( t ) commutes with Σ , yielding
[Equation]
and
[Equation]
Comparing the W 2 distance for naïve Langevin sampling (D.23) with (D.36), we can see that sampling in the space of natural parameters eliminates the sensitivity of convergence speed to large eigenvalues of the covariance matrix.
Moreover, comparing the stationary cross-covariance at timelag τ ≡ t − t ′ for naïve Langevin sampling, C s ( τ ) =
e − Σ − 1 | τ | Σ , to that for sampling in the natural space, C s ( τ ) =
e −| τ | Σ , we can see that the same qualitative difference is present.
We note that
Hennequin et al.
[9] focused on the speed of decay in C s ( τ ) .
Therefore, in this simple setting, there is a clear intuitive picture of why natural gradients enable fast sampling.
E Numerical methods and supplementary figures
In this appendix, we describe our numerical methods and include supplementary figures.
All simulations were run in M ATLAB 9.10 (R2021a) or 9.12 (R2022a) (The MathWorks, Natick, MA,
USA) on desktop workstations (CPU: Intel i9-9900K or Xeon W-2145, 64GB RAM).
They were not computationally intensive, and required less than 24 hours of compute time in total.
The code used to generate all figures is available from GitHub: https://github.com/Pehlevan-Group/
FastSpikingSampler .
For the sweeps in ρ we tested 100 values of ρ ∈ [0; 0 . 99] .
For the dimension sweeps we varied n p ∈
[2 , 4 , 8 , 16 , 32 , 64] and k = n n n p ∈ [1 , 5 , 10 , 20] is the number of neurons per parameters.
Here, we showed results for k = 10 .
As expected, sampling failed in the case of k = 1 as the sign constraint introduced by spiking restricts the network to efficiently sample only one half on the values for each parameter but results were qualitatively similar for k = n n /n p ∈
In Figures 2, E.2, E.3, 4, and E.4, convergence statistics are computed based on distributions of samples over time, that is, the values visited by the sampler over the course of a single trial.
Note that this is different than many machine learning studies, which instead consider distributions across realizations at a single timepoint.
However, for an organism, probabilistic inference must be performed within a single trial.
Because estimation of the full 2-Wasserstein distance between high-dimensional distributions is computationally expensive [96], we instead computed the mean across dimensions of the 2-Wasserstein distances between the marginals of the sampling and target distributions.
For Figures 2, E.2 and E.3, we sampled with a discretization timestep of ∆ = 10 − 4 s and a membrane time constant of τ m = 20 ms. Γ by drawing independent and identically distributed Gaussian element.
We uses τ s = 0 .
01 τ m and scaled the elastic net regularization parameters using α = λ = √ n n .
For Figures 3, 4, and E.4, we simulated a Metropolis-Hastings sampling network with a discretization timestep of ∆ = 10 − 5 s and a membrane time constant of τ m = 20 ms.
In the naïve case, we generated the readout matrices as Γ =
[ − M , M ] for random matrices M with independent and identically distributed Gaussian elements.
To perform sampling in approximately the natural space, we chose Γ = Σ 1 / 2 [ − M , M ] for M a random matrix with i.i.d.
Gaussian elements.
For the dimension sweeps in 4c and E.4b, we scale the variance of the elements of M to be 1 /n p .
We probe the sensitivity of the sampler to the variance of the random matrix in Figure B.1, showing
S
[Figure] [Figure]
that its performance is robust within some range of variances.
The distributions in Figure 3b were generated using 1000 realizations of the randomness in the proposal and accept/reject steps for a single realization of the random matrix M .
Statistics in Figures 4 and E.4 were computed acrossrealizations of the random matrix M and of the randomness in the proposal and accept/reject steps.
S a i 0 0 .
Wd i s t an c e ii iii i ii iii b 2 4 8 16 32n p
M ean 0 0 .
M e an 2 4 8 16 32 64 n p
V a r i an c e 0 0 .
V a r
i
[Figure] [Figure]
p
[Figure]
S ρ a a v e r age f
i r i
ng r a t e (
H z )    naïvenatural s t ead y - s t a t e m ean
[Figure]
s t ead y - s t a t e
[Figure]
r i
[Figure] [Figure]
ead y
[Figure]
s t
[Figure]
t e
[Figure]
b i ii iii iv a v
[Figure]
r age f i r i
ng r a t e ( k
H z )  n p
[Figure]
n
[Figure] [Figure]
n
[Figure] [Figure]
n
[Figure] [Figure]
i ii
[Figure]
ead y -
[Figure]
t
[Figure]
t
[Figure]
ean
[Figure]
ead y
[Figure] [Figure]
ead
[Figure]
t
[Figure]
t
[Figure]
W
[Figure]
naïvenatural
[Figure]
S