System-size expansion and Gaussian moment closure for Quiescent-Active-Refractory model

Epilogue: These notes concern the system size expansion and moment closure later published as "Neural field models for latent state inference: Application to large-scale neuronal recordings" [pdf] (more) 

These notes derive the Kramers-Moyal system size expansion and approximate equation for the evolution of means and covariances of a single-compartment neural mass model with Quiescent (Q), Active (A), and Refractory (R) states. The derivations here are identical to the standard ones for the Susceptible (S), Infected (I), and Recovered (R) (SIR) model commonly used in epidemiology. 

[get these notes as PDF]

Normally, in the system-size expansion we take the number of neurons $N$ to be large. In this case we can expand the master equation in terms of the state change for a single-neuron, which is a small parameter $\mathcal O(1/N)$

But, re-scaling the equations by the system size $N$ only complicates the notation. It is equivalent to take the expansion in terms of a single neuron. This is $\mathcal O(1)$, but we understand that as $N$ grows large it becomes increasingly accurate relative to the size of the total population.

The system

$$ \begin{aligned} Q &\xrightarrow{\rho_{\textrm{q}}} A\\ Q+A &\xrightarrow{\rho_{\textrm{e}}} A+A\\ A &\xrightarrow{\rho_{\textrm{a}}} R\\ R &\xrightarrow{\rho_{\textrm{r}}} Q\\ \end{aligned} $$

In retinal waves, the spontaneous reaction $Q \xrightarrow{\rho_{\textrm{q}}} A$ is rare. It is poorly modeled by a mean-field approach. Instead, we treat this reaction separately as an external input (Poisson a.k.a. shot noise).

Kramers-Moyal expansion

Consider the non-spatial case for illustration. The allowed transitions of the neural population can be described in terms of a master equations. For review, the possible reactions include creation and annihilation of cells to and from the Quiescent (Q), Active (A), and Refractory (R) states

$$ \begin{aligned} \partial_t \Pr(Q,A,R) &= \Pr(Q{+}1,A{-}1,R) \rho_{\textrm{e}} (A-1) (Q+1) \\ &+ \Pr(Q,A{+}1,R{-}1) \rho_{\textrm{a}} (A+1) \\ &+ \Pr(Q{-}1,A,R{+}1) \rho_{\textrm{r}} (R+1) \\ &-\Pr(Q,A,R) \left[ \rho_{\textrm{e}}AQ + \rho_{\textrm{a}}A + \rho_{\textrm{r}}R \right] \end{aligned} $$

The terms, in order, reflect (1) $Q+A\rightarrow A+A$ transition into this state, (2) spontaneous $A\rightarrow R$ transition into this state, (3) spontaneous $R\rightarrow Q$ transition into this state, (4) the sum of all possible transitions away from this state.

We approximate $\Pr(Q,A,R)$ as Gaussian and derive terms for the evolution of its mean and covariance by a second-order Taylor expansion of the master equation, parameterized by the system size $N$, the number of interacting agents within a single well-mixed compartment. First, re-write the steps in the master equation in terms of the system size $N$, explicitly. $Q+A+R=N$.

$$ \begin{aligned} \partial_t \Pr(Q,A,R) &= \Pr(S{+}1,A{-}1,R) \rho_{\textrm{e}} (A{-}1)(Q{+}1) \\ &+ \Pr(Q,A{+}1,R{-}1) \rho_{\textrm{a}}(A{+}1) \\ &+ \Pr(Q{-}1,A,R{+}1) \rho_{\textrm{r}}(R{+}1) \\ &- \Pr(Q,A,R) \left[ \rho_{\textrm{e}}AQ + \rho_{\textrm{a}}A + \rho_{\textrm{r}}R \right] \end{aligned} $$

Now, let $\epsilon(\varepsilon_q,\varepsilon_a,\varepsilon_r)$ be a Taylor expansion of the step operator (denoting the identity operator as $1$). This represents the change in probability density $\Pr(Q,A,R)$ for an infinitesimal change in intensity for the $Q,A,R$ fields.

$$ \begin{aligned} \epsilon(\varepsilon_q,\varepsilon_a,\varepsilon_r) &= 1 + \varepsilon_q \partial_Q + \varepsilon_a \partial_A + \varepsilon_r \partial_R \\ &+ \tfrac{1}{2} \left( \varepsilon^2_q \partial^2_Q + \varepsilon^2_a \partial^2_A + \varepsilon^2_r \partial^2_R \right) \\ &+ \left( \varepsilon_q \varepsilon_a \partial_Q \partial_A + \varepsilon_a \varepsilon_r \partial_A \partial_R + \varepsilon_r \varepsilon_q \partial_R \partial_Q \right) \end{aligned} $$

We can approximate the first three terms of the master equation in terms of this expansion (abbreviating $\Pr(Q,A,R)$ as $P$)

$$ \begin{aligned} \Pr(Q{+}1,A{-}1,R) \rho_{\textrm{e}} (A{-}1)(Q{+}1) &\approx \epsilon\left(1,-1,0 \right) \rho_{\textrm{e}} AQ P \\ \Pr(Q,A{+}1,R{-}1) \rho_{\textrm{a}}(A{+}1) &\approx \epsilon\left(0,1,-1 \right) \rho_{\textrm{a}}A P \\ \Pr(Q{-}1,A,R{+}1) \rho_{\textrm{r}}(R{+}1) &\approx \epsilon\left(-1,0,1 \right) \rho_{\textrm{r}}R P \end{aligned} $$

Substituting the Taylor approximations into the master equation (for brevity, let $\rho_{\textrm{qa}} = \rho_{\textrm{e}} A$. Note that this depends on $A$ and is not constant).

$$ \partial_t P = \epsilon\left(1,-1,0 \right) \rho_{\textrm{qa}} Q P\\ + \epsilon\left(0,1,-1\right) \rho_{\textrm{a}}A P\\ + \epsilon\left(-1,0,1 \right) \rho_{\textrm{r}}R P \\ -\left[ \rho_{\textrm{qa}}Q + \rho_{\textrm{a}}A + \rho_{\textrm{r}}R \right] P $$

Regrouping terms and dividing by $P$, we get:

$$ \frac {\partial_t P} P = \left[ \epsilon\left(1,-1,0 \right) -1 \right] \rho_{\textrm{qa}} Q + \left[ \epsilon\left(0,1,-1 \right) -1 \right] \rho_{\textrm{a}}A + \left[ \epsilon\left(-1,0,1 \right) -1 \right] \rho_{\textrm{r}}R $$

Substituting the definition of $\epsilon(\varepsilon_s,\varepsilon_i,\varepsilon_r)$ and collecting terms leads to a second-order description of the master equation, which can be interpreted as the drift and diffusion terms for the evolution of the mean and the covariance of a multivariate Gaussian approximation of $\Pr(Q,A,R)$.

$$ \begin{aligned} \frac {\partial_t P} P &= \left[ \left(\partial_Q-\partial_A\right) + \tfrac{1}{2} \left(\partial^2_Q+\partial^2_A - 2\partial_Q \partial_A \right) \right] \rho_{\textrm{qa}} Q \\ &+ \left[ \left(\partial_A-\partial_R\right) + \tfrac{1}{2} \left(\partial^2_A+\partial^2_R - 2\partial_A \partial_R \right) \right] \rho_{\textrm{a}}A \\ &+ \left[ \left(\partial_R-\partial_Q\right) + \tfrac{1}{2} \left(\partial^2_R+\partial^2_Q - 2 \partial_R \partial_Q \right) \right] \rho_{\textrm{r}}R \end{aligned} $$

The first and second order terms form the second-order expansion of the master equation yield the drift and diffusion terms for a Fokker-Plank equation, respectively. Denote state $(Q,A,R)$ as $\mathbf{x} = (Q,A,R)$

Collecting first-order terms yields a drift term $\mu(\mathbf{x})$

$$ -1 \left[ \partial_Q\left[ \rho_{\textrm{r}}R - \rho_{\textrm{qa}}Q \right] + \partial_A\left[ \rho_{\textrm{qa}}Q - \rho_{\textrm{a}}A \right] + \partial_R\left[ \rho_{\textrm{a}}A - \rho_{\textrm{r}}R \right] \right] $$$$ \mu(\mathbf{x}) = \begin{bmatrix} \rho_{r}R {-}\rho_{qa}(A)Q \\ \rho_{qa}(A)Q {-}\rho_{a}A \\ \rho_{a}A {-}\rho_{r}R \\ \end{bmatrix} $$

Collecting second-order terms yields a diffusion term $\Sigma(\mathbf{x})$

$$ \tfrac{1}{2} \left( \partial^2_Q\left[\rho_{\textrm{qa}}Q+\rho_{\textrm{r}}R \right]+ \partial^2_A\left[\rho_{\textrm{a}}A +\rho_{\textrm{qa}}Q \right]+ \partial^2_R\left[\rho_{\textrm{r}}R +\rho_{\textrm{a}}A \right] \\ -2 \left[ \partial_Q \partial_A \rho_{\textrm{qa}} Q + \partial_A \partial_R \rho_{\textrm{a}}A + \partial_R \partial_Q \rho_{\textrm{r}}R \right] \right) $$$$ \Sigma(\mathbf{x}) = \begin{bmatrix} \rho_{qa}Q + \rho_{r}R &-\rho_{qa}Q &-\rho_{r}R \\ -\rho_{qa}Q & \rho_{a}A + \rho_{qa}Q &-\rho_{a}A \\ -\rho_{r} R & -\rho_{a}A & \rho_{r} R + \rho_{a}A \\ \end{bmatrix} $$

The drift $\mu(\mathbf{x})$ and diffusion $\Sigma(\mathbf{x})$ terms define a Fokker-Plank equation that can be used to propagate a state-space model for the $Q,A,R$ system forward in time:

$$ \partial_t \Pr(\mathbf{x}) = - \textstyle\sum_i \partial_{x_i} \mu_i(\mathbf{x}) \Pr(\mathbf{x}) + \tfrac{1}{2} \textstyle\sum_{ij} \partial^2_{x_i,x_j} \Sigma_{ij}(\mathbf{x}) \Pr(\mathbf{x}) $$

Reformulating as a Langevin equation

$$ \mathbf{x} = (Q,A,R) $$$$ d\mathbf{x} = \mu(\mathbf{x}) dt + Q(\mathbf{x}) dW $$

Where $Q$ is a matrix square root of $\Sigma(\mathbf{x})$

$$ Q^\top Q = \Sigma(\mathbf{x}), $$

for example

$$ Q(\mathbf{x}) = \begin{bmatrix} -\sqrt{\rho_{\textrm{qa}}Q} & 0 & \sqrt{\rho_{\textrm{r}}R} \\ \sqrt{\rho_{\textrm{qa}} Q} & -\sqrt{\rho_{\textrm{a}}A}& 0 \\ 0 & \sqrt{\rho_{\textrm{a}}A}& -\sqrt{\rho_{\textrm{r}}R} \\ \end{bmatrix}. $$

To normalize by the number of neurons, divide $\mu(\mathbf{x})$ and $Q(\mathbf{x})$ by $\tfrac{1}{N}$.

Moment closure

So far, we've used a system-size expansion to derive a diffusive approximation of the dynamics. This can be used to simulate $\Pr(\mathbf x)$ in time, as a partial differential equation. It can also be used to sample trajectories from the model as a Langevin equation. We next discuss various approximations to the time evolution of the moments of the state distribution.

Using the drift and diffusion terms above directly gives the so-called linear noise approximation (LNA) to the state-space dynamics. This uses the deterministic rate equations for the evolution of the mean.

In state-space inference models, fluctuations cause the densities of quiescent and excitable neurons to be anti-correlated. This has the effect of stabilizing the system, but is not captured in the LNA. These can be captured at second order by treating the state distribution as Gaussian, and neglecting higher cumulants. This is the so-called "cumulant neglect" or Gaussian moment closure.

The time-evolution of the first moment is derived by taking the time derivative of the expected values of the states over $\Pr(Q,A,R)$: 

$$ \begin{aligned}\partial_t \left\langle Q \right\rangle\left\langle \partial_t Q \right\rangle &= \left\langle \rho_r R - \rho_q Q - \rho_e Q A \right\rangle \ &= \rho_r \left\langle R \right\rangle -\rho_q \left\langle Q \right\rangle -\rho_e \left\langle Q A \right\rangle \\partial_t \left\langle A \right\rangle\left\langle \partial_t A \right\rangle &= \rho_q \left\langle Q \right\rangle +\rho_e \left\langle Q A \right\rangle -\rho_a \left\langle A \right\rangle \\partial_t \left\langle R \right\rangle\left\langle \partial_t R \right\rangle &= \rho_a \left\langle A \right\rangle -\rho_r \left\langle R \right\rangle \end{aligned} $$

Note that the first moment couples to the second moment through pairwise interaction terms $\left\langle Q A \right\rangle$.

The time evolution of the second moment, again, can be derived by taking the time derivative of expectations $\left<x_1 x_2\right>$ where $x_i \in \{Q,A,R\}$, for example $\left\langle Q A \right\rangle$ or $\left\langle AR \right\rangle $:

$$ \begin{aligned} \partial_t \left\langle Q A \right\rangle = \left\langle \partial_t Q A \right\rangle &= \left\langle Q \partial_t A \right\rangle + \left\langle A \partial_t Q \right\rangle \\ &= \left\langle Q (QA\rho_e + Q\rho_q - A \rho_a ) \right\rangle + \left\langle A (-R \rho_r - QA\rho_e - Q\rho_q) \right\rangle \\ &= \rho_e \left\langle QQA \right\rangle + \rho_q \left\langle QQ \right\rangle - \rho_a \left\langle QA \right\rangle - \rho_r \left\langle AR \right\rangle - \rho_e \left\langle AQA \right\rangle - \rho_q \left\langle AQ \right\rangle \end{aligned} $$$$ \begin{aligned} \partial_t \left\langle AR \right\rangle = \left\langle \partial_t AR \right\rangle &= \left\langle R \partial_t A \right\rangle + \left\langle A \partial_t R \right\rangle \\ &= \left\langle R (QA\rho_e + Q\rho_q - A \rho_a) \right\rangle + \left\langle A (A \rho_a - R \rho_r) \right\rangle \\ &= \rho_e \left\langle QQA \right\rangle + \rho_q \left\langle QQ \right\rangle - \rho_a \left\langle QA \right\rangle - \rho_r \left\langle AR \right\rangle - \rho_e \left\langle AQA \right\rangle - \rho_q \left\langle AQ \right\rangle \end{aligned} $$

Gaussian moment closure makes the simplifying assumption that we can replace third-order terms like $\left\langle x_1x_1 x_2 \right\rangle$ with the values that one might expect if the state distribution is a Gaussian with parameters given by the first two moments.