Four ways to do Gaussian process regression

 

Form (a)

In most textbooks or tutorials you'll see the posterior mean written as:

$$\begin{array}{}\begin{array}{cc}{y}_2& \sim N(\mu ,\Sigma )\\ \mu & ={\mu }_2+{\Sigma }_{12}^{\top }[{\Sigma }_{11}+{\Sigma }_{ϵ}{]}^{-1}(y_1-{\mu }_1)\\ \Sigma & ={\Sigma }_{22}-{\Sigma }_{12}^{\top }[{\Sigma }_{11}+{\Sigma }_{ϵ}{]}^{-1}{\Sigma }_{12},\end{array}& \text{(64)}\end{array}$$

where ${\Sigma }_{11}$ and ${\Sigma }_{22}$ are the prior covariance of the observation and output points, respectively; ${\Sigma }_{12}$ is the cross-covariance of the observation and output points; ${\Sigma }_{ϵ}$ is the measurement noise covariance; ${\mu }_2$ and ${\mu }_1$ are the prior mean at the observation and output points, respectively, and $y_1$ are the observations.

When observations are sparse, this form is computationally efficient efficient, because ${\Sigma }_{11}$ and ${\Sigma }_{ϵ}$ are small, and the expression ${\Sigma }_{12}^{\top }[{\Sigma }_{11}+{\Sigma }_{ϵ}{]}^{-1}{\Sigma }_{12}$ is low-rank. This form is especially convenient when updating a GP model with a single observation, in which case it reduces to a rank-1 update

$$\begin{array}{}\begin{array}{cc}\mu & ={\mu }_2+\frac{1}{{\sigma }_{11}^2+{\sigma }_{ϵ}^2}{\Sigma }_{12}^{\top }(y_1-{\mu }_1)\\ \Sigma & ={\Sigma }_{22}-\frac{1}{{\sigma }_{11}^2+{\sigma }_{ϵ}^2}{\Sigma }_{12}^{\top }{\Sigma }_{12}.\end{array}& \text{(65)}\end{array}$$

However, in our application we have an extended time-series where a rat visits each location many times. There are many more observations than output points, and ${\Sigma }_{11}$ and ${\Sigma }_{ϵ}$ are large matrices. Binning together nearby measurements reduces the complexity. However, coverage of the space is quite dense. ${\Sigma }_{11}$ and ${\Sigma }_{ϵ}$ are almost as large as ${\Sigma }_{22}$, even excluding bins with no observations. We therefore assume that the observations and output points are evaluated on a common grid, in which case the equations take the form:

$$\begin{array}{}\begin{array}{cc}\mu & ={\mu }_0+{\Sigma }_0^{\top }[{\Sigma }_0+{\Sigma }_{ϵ}{]}^{-1}(y-{\mu }_0)\\ \Sigma & ={\Sigma }_0-{\Sigma }_0^{\top }[{\Sigma }_0+{\Sigma }_{ϵ}{]}^{-1}{\Sigma }_0,\end{array}& \text{(a)}\end{array}$$

where ${\Sigma }_0$ ${\mu }_0$ are the prior covariance and mean, ${\Sigma }_{ϵ}$ is the measurement error covariance, and $y$ are the measurements

This form has the following useful properties:

• ${\Sigma }_0$ and ${\Sigma }_{ϵ}$ can be singular, provided ${\Sigma }_0+{\Sigma }_{ϵ}$ is not. In GP regression, this allows measurements with zero error, and also allows priors with zero eigenvalues. For example, one might set a prior that has zeros for high-frequency components, to encode a strong assumption that the posterior function should be smooth.

• The update is low-rank and fast when observations are limited

• If ${\Sigma }_0+{\Sigma }_{ϵ}$ is nonsingular, then it is positive definite. This allows one to compute $[{\Sigma }_0+{\Sigma }_{ϵ}{]}^{-1}y$ quickly via Cholesky factorization. (However, it is even faster to use a Krylov subspace solver, if the prior ${\Sigma }_0$ is in a form that supports fast matrix-vector product via the FFT.)

• The final matrix multiplications by ${\Sigma }_0$ can be computed via FFT. (The data need to be zero padded to do this, but this padding can be stripped away after performing the convolution, so there is no added complexity cost. Contrast this to form b, in which the zero-padding must be retained before solving a linear system, which increases the complexity).

Problem: This form is unsuitable if ${\Sigma }_{ϵ}^{-1}$ is singular. In order to leverage the FFT for matrix calculations, we need to perform the GP regression over a regular and periodic grid. To avoid opposite ends of the space from interacting, we need to add as many zeros to the edge of our data as our kernel is wide. This means that, inherently, some bins will be missing data. Missing observations can be represented as zero precision (inverse variance), so ${\Sigma }_{ϵ}^{-1}$ can be constructed—but it will not be invertable, so ${\Sigma }_{ϵ}$ does not exist.

Form (b)

It's also common to encounter the following form, when deriving the posterior mean for the product of two multivariate Gaussian distributions:

$$\begin{array}{}\begin{array}{cc}\mu & ={\mu }_0+{[{\Sigma }_0^{-1}+{\Sigma }_{ϵ}^{-1}]}^{-1}{\Sigma }_{ϵ}^{-1}(y-{\mu }_0).\end{array}& \text{(b)}\end{array}$$

This form has the following useful properties

• ${\Sigma }_{ϵ}^{-1}$ can be singular, so we can include "null" observations when calculating the regression over a regular grid
• ${\Sigma }_{ϵ}^{-1}=\mathrm{diag}\left[\tau \right]$ is diagonal, so ${\Sigma }_{ϵ}^{-1}y=\mathrm{diag}[\tau \circ y]$, which is trivial to compute
• If ${\Sigma }_0$ is nonsingular, then ${\Sigma }_0^{-1}$ exists and can be computed quickly via the FFT
• ${\Sigma }_0^{-1}+{\Sigma }_{ϵ}^{-1}$ is positive definite. ${[{\Sigma }_0^{-1}+{\Sigma }_{ϵ}^{-1}]}^{-1}\mathrm{diag}[\tau \circ y]$ can therefore be solve efficiently using Cholesky factorization.
• Priors that assume correlations arise from nearest-neighbor interactions can be represented as a prior precision ${\Sigma }_0^{-1}$ that is nonzero only for entries for pairs of adjacent regions. This sparsity, when combined with Krylov subspace algorithms for solving the linear system, this allows for fast solutions on arbitrary topologies, for which spectral methods might not be possible.

Problem: This form is unsuitable if ${\Sigma }_0$ is singular, and is numerically unstable if ${\Sigma }_0$ is ill-conditioned. So, it requires regularization to ensure that no eigenvalue of ${\Sigma }_0$ is too small. For larger problems, it may require so much regularization that the prior ${\Sigma }_0$ is altered substantially, affecting the accuracy of the inferred posterior mean. Additionally, computing ${\Sigma }_0^{-1}$ via FFT requires adding zero padding to handle the circular boundary conditions correctly. This padding cannot be removed before solving the subsequent linear system without introducing boundary artifacts, so this can lead to slightly larger linear system to solve. Since the time complexity of this is $O\left(L^6\right)$ in terms of the linear dimension of a $2\times 2$ grid, this can lead to a significant slowdown.

Form (c)

We can also pull out ${\Sigma }_0$ from form (b), and solve for the posterior mean as :

$$\begin{array}{}\begin{array}{cc}\mu & ={\mu }_0+{[{\Sigma }_0{\Sigma }_{ϵ}^{-1}+I]}^{-1}\left[{\Sigma }_0{\Sigma }_{ϵ}^{-1}\right(y-{\mu }_0\left)\right].\end{array}& \text{(c)}\end{array}$$

This form has the following useful properties:

• Both ${\Sigma }_0$ and ${\Sigma }_{ϵ}^{-1}$ can be singular
• ${\Sigma }_{ϵ}^{-1}y=\mathrm{diag}[\tau \circ y]$ is trivial, and ${\Sigma }_0\mathrm{diag}[\tau \circ y]$ can be calculated via FFT
• The product ${\Sigma }_0{\Sigma }_{ϵ}^{-1}={\Sigma }_0\circ \left(\tau 1^{\top }\right)$ is trivial.
• ${\Sigma }_0{\Sigma }_{ϵ}^{-1}+I$ is well conditioned, so this calculation is fairly numerically stable, requiring less regularization.
• If ${\Sigma }_0$ contains many zero eigenvalues, then it is low rank . In the special case that ${\Sigma }_0$ is circulant, the FFT offers a fast conversion to/from the eigenspace of ${\Sigma }_0$, and it is possible to use the nonzero Fourier components of ${\Sigma }_0$ as a low-rank basis for calculations.

Problem: The matrix ${\Sigma }_0{\Sigma }_{ϵ}^{-1}+I$ will not be symmetric in general, so we cannot use Cholesky factorization to solve the linear system here. This can lead to significant slow-downs for larger systems. However! Once ${\Sigma }_0{\Sigma }_{ϵ}^{-1}y$ is computed via FFT, one can remove the zero-padding before solving the linear system, which reduces the problem size. Since computing circulant matrix operations using FFT requires as much zero-padding as our prior kernel is wide, this reduction in problem size can be significant. Since solving the linear system scales as $O\left(L^6\right)$ for a $L\times L$ spatial grid, a small reduction in $L$ is much more important than the constant-factor improvement provided by Cholesky decomposition.

Additionally, for translationally-invariant kernels on a regular grid, the matrix-vector product $[{\Sigma }_0{\Sigma }_{ϵ}^{-1}+I]u$ can be calculated quickly via FFT. This used with a Krylov-subspace-based solver like minres, and can be orders of magnitude faster than other approaches.

Form (d)

In the special case that the measurement error covariance is constant, ${\Sigma }_{ϵ}={\sigma }_{ϵ}^{2}I$, then GP regression reduces to a convolution, and can be evaluated using the Fourier transform thanks to the convolution theorem:

$$\begin{array}{}\begin{array}{cc}\mu & ={\mu }_0+{\text{F}}^{-1}[\left(\frac{\tilde{\kappa }}{\tilde{\kappa }+{\sigma }_{ϵ}^2}\right)\circ \text{F}(y-{\mu }_0\left)\right],\end{array}& \text{(d)}\end{array}$$

where $\text{F}$ is the Fourier transform.

Problem: The regression must be on a periodic domain in order to compute $\text{F}$ quickly using the FFT. This can cause measurements at opposite sides of the spatial grid to influence each-other. This can be solved by zero-padding. However, the above assumption assumes ${\sigma }_{ϵ}^2$ for all observations, so this zero padding will also be treated as a $y=0$ observation with error ${\sigma }_{ϵ}^2$. This will lead to some tapering toward zero at the boundaries. Alternatively, one may use reflected boundary conditions, copying a mirrored version of $y$ into the padded space. This reflected boundary condition reduces artifacts, but is not identical to solving the original, un-padded GP regression problem.

Note: if ${\Sigma }_{ϵ}$ is not constant, but changes slowly relative to the scale of the prior kernel, then one may also decompose a larger problem into smaller ones that tile the space.

When to use which?

Use form (a) when

• Measurement variance ${\Sigma }_{ϵ}$ and prior covariance ${\Sigma }_{ϵ}$ are well defined.
• Observations are sparse compared to the number of output points.
• This is as fast or faster than (b,c) for any size problem, but the extra matrix multiplications can be expensive for large systems, or when observations are dense.

Use form (b) with when

• The measurements and outputs are evaluated at the same set of points.
• ${\Sigma }_{ϵ}^{-1}$ is well defined, but ${\Sigma }_{ϵ}$ is not.
• ${\Sigma }_0^{-1}$ exists and well-conditioned.
• For medium-sized systems,
  • ${\Sigma }_0^{-1}+{\Sigma }_{ϵ}$ can be solved via Cholesky factorization.
• For large systems, use a Krylov subspace solver and leverage special properties of ${\Sigma }_0^{-1}$:
  • ${\Sigma }_0^{-1}$ is circulant, so we can use the FFT, or
  • ${\Sigma }_0^{-1}$ is sparse, arising from a nearest-neighbor model.

Use form (c) when

• The measurements and outputs are taken at the same set of points
• ${\Sigma }_{ϵ}^{-1}$ is well defined, but ${\Sigma }_{ϵ}$ is not.
• ${\Sigma }_0$ is low-rank so ${\Sigma }_0^{-1}$ does not exist.
• For large systems, use a Krylov subspace solver and leverage special properties of ${\Sigma }_0$:
  • If ${\Sigma }_0$ is circulant, use the FFT
  • Or, use a low-rank approximation ${\Sigma }_0=Q{Q}^{\top }$

Use form (d) when

• The measurements and outputs are evaluated on a regular grid.
• The problem is too large to calculate using ordinary matrix operations
• Measurement error can be approximated as constant ${\Sigma }_{ϵ}\approx {\sigma }_{ϵ}^{2}I$
• Artifacts from zero or mirrored boundary conditions are acceptable, or the kernel is local and it is acceptable to discard the boundary regions.

Gradient and Hessian for variational inference in Poisson and probit Generalized Linear Models

These notes contain some derivations for variational inference in Poisson and probit Generalized Linear Models (GLMs) with a Gaussian prior and approximated Gaussian posterior. ( see also here. )

Problem statement

Consider a population of neurons with firing-intensities $\boldsymbol\lambda =\rho(\boldsymbol\theta)$, where $\rho(\cdot)$ is a firing-rate nonlinearity and $\boldsymbol\theta$ is a vector of synaptic activations (amount of input drive to each neuron). For stochastic models of spiking $\Pr(y|\theta)$ in the canonical exponential family, the probability of observing spikes $\mathbf y$ given $\boldsymbol\theta$ can be written as \begin{equation}\begin{aligned}\ln\Pr(\mathbf y|\mathbf z) &=\mathbf y^\top\boldsymbol\theta-\mathbf 1^\top A(\boldsymbol\theta)+\text{constant},\end{aligned}\end{equation}

where $A(x)$ is a known function whose derivative equals the firing-rate nonlinarity, i.e. $A'(\cdot) =\rho(\cdot)$.

Assume that the synaptic activations $\boldsymbol\theta$ are driven by shared latent variables $\mathbf z$ with a Gaussian prior $\mathbf z\sim\mathcal N(\boldsymbol\mu_z,\boldsymbol\Sigma_z)$. Let $\boldsymbol\theta=\mathbf B\mathbf z$, where "$\mathbf B$" is a matrix of coupling coefficients which determine how the latent factors $\mathbf z$ drive each neuron.

We want to infer the distribution of $\mathbf z$ from observed spikes $\mathbf y$. The posterior is given by Bayes rule, $\Pr(\mathbf z |\mathbf y) =\Pr(\mathbf y |\mathbf z)\Pr(\mathbf z)/\Pr(\mathbf y)$. However, this posterior does not admit a closed form if $A(\cdot)$ is nonlinear. Instead, one can use a variational Bayesian approach to obtain an approximate posterior.

Variational Bayes

In variational Bayes, the posterior on $\boldsymbol z$ is approximated as Gaussian, i.e. $\Pr(\mathbf z |\mathbf y)\approx Q(\mathbf z)$, where $Q(\mathbf z) =\mathcal N(\boldsymbol\mu_q,\boldsymbol\Sigma_q)$. We then optimize $\boldsymbol\mu_q$ and $\boldsymbol\Sigma_q$ to minimize the Kullback-Leibler (KL) divergence from the true posterior $\Pr(\mathbf z |\mathbf y)$ to $Q(\mathbf z)$. This is equivalent to minimizing the KL divergenece $D_{\text{KL}}\left[ Q(\mathbf z)\|\Pr(\mathbf z)\right]$ from the prior to the posterior, while maximizing the expected log-likelihood $\left<\Pr(\mathbf y |\mathbf z)\right>$: \begin{equation}\begin{aligned}D_{\text{KL}}\left[ Q(\mathbf z)\|\Pr(\mathbf z |\mathbf y)\right]&=D_{\text{KL}}\left[ Q(\mathbf z)\|\Pr(\mathbf z)\right]-\left<\ln\Pr(\mathbf y |\mathbf z)\right>+\text{constant}.\end{aligned}\end{equation}

(In these notes, all expectations $\langle\cdot\rangle$ are taken with respect to the approximating posterior distribution.)

Since both $Q(\mathbf z)$ and $\Pr(\mathbf z)$ are multivariate Gaussian, the KL divergence $D_{\text{KL}}\left[ Q(\mathbf z)\|\Pr(\mathbf z)\right]$ has the closed form : \begin{equation}\begin{aligned} D_{\text{KL}}\left[ Q(\mathbf z)\|\Pr(\mathbf z)\right] &= \tfrac 1 2\left\{ (\boldsymbol\mu_z-\boldsymbol\mu_q)^\top \boldsymbol\Sigma_z^{-1} (\boldsymbol\mu_z-\boldsymbol\mu_q) + \operatorname{tr}\left( \boldsymbol\Sigma_z^{-1} \boldsymbol\Sigma_q \right) + \ln| \boldsymbol\Sigma_z^{-1} \boldsymbol\Sigma_q | \right\} +\text{constant}. \end{aligned}\end{equation}

For our choice of the canonically-parameterized natural exponential family, the expected negative log-likelihood can be written as: \begin{equation}\begin{aligned}-\langle\ln\Pr(\mathbf y|\mathbf z)\rangle&=\mathbf 1^\top\langle A(\boldsymbol\theta)\rangle-\mathbf y^\top\mathbf B\boldsymbol\mu_q+\text{constant}.\end{aligned}\end{equation}

Neglecting constants and terms that do not depend on $(\boldsymbol\mu_q,\boldsymbol\Sigma_q)$, the overall loss function to be minimized is: \begin{equation}\begin{aligned} \mathcal L(\boldsymbol\mu_q,\boldsymbol\Sigma_q) &= \tfrac 1 2\left\{ (\boldsymbol\mu_z-\boldsymbol\mu_q)^\top \boldsymbol\Sigma_z^{-1} (\boldsymbol\mu_z-\boldsymbol\mu_q) + \operatorname{tr}\left( \boldsymbol\Sigma_z^{-1} \boldsymbol\Sigma_q \right) + \ln|\boldsymbol\Sigma_z^{-1}\boldsymbol\Sigma_q|\right\} +\mathbf 1^\top\langle A(\boldsymbol\theta)\rangle -\mathbf y^\top\mathbf B\boldsymbol\mu_q \end{aligned} \label{loss}. \end{equation}

Closed-form expectations

To optimize $\eqref{loss}$, we need to differentiate it in $\boldsymbol\mu_q$ and $\boldsymbol\Sigma_q$. These derivatives are mostly straightforward, but the expectation $\langle A(\boldsymbol\theta)\rangle$ poses difficulties when $A(\cdot)$ is nonlinear. We'll consider some choices of firing-rate nonlinearity for which the derivatives of $\langle A(\boldsymbol\theta)\rangle$ have closed-form expressions when $\boldsymbol\theta$ is Gaussian.

Because we've assumed a Gaussian posterior on our latent state $\mathbf z$, and since $\boldsymbol\theta =\mathbf B\mathbf z$, the synaptic activations $\boldsymbol\theta$ are also Gaussian. The vectors $\boldsymbol\mu_\theta$ and $\boldsymbol\sigma^2_\theta$ for the mean and variance of $\boldsymbol\theta$, respectively, are: \begin{equation}\begin{aligned}\boldsymbol\mu_\theta &=\mathbf B\boldsymbol\mu_q\\\boldsymbol\sigma^2_\theta &=\operatorname{diag}\left[\mathbf B\boldsymbol\Sigma_q\mathbf B^\top\right]\end{aligned}\end{equation}

Consider a single, scalar $\theta\sim\mathcal N(\mu,\sigma^2)$. Using the chain rule and linearity of expectation, one can show that the partial derivatives $\partial_{\mu}\langle A(\theta)\rangle$ and $\partial_{\sigma^2}\langle A(\theta)\rangle$, with respect to $\mu$ and $\sigma^2$ respectively, are: \begin{equation}\begin{aligned}\partial_{\mu}\langle A(\theta)\rangle&=\langle A'(\theta) \rangle= \langle \rho(\theta) \rangle \\\partial_{\sigma^2}\langle A(\theta)\rangle&= \tfrac1{2\sigma^2}\left<(\theta-\mu_\theta) A'(\theta)\right>=\tfrac1{2\sigma^2} \left<(\theta-\mu)\rho(\theta)\right>.\end{aligned}\label{dexpect}\end{equation}

For more compact notation, denote the expected firing rate as $\bar\lambda=\langle\rho(\theta)\rangle$, and denote the expected derivative of the firing-rate in $\theta$ as $\bar\lambda'=\langle\rho'(\theta)\rangle$. Note that $\bar\lambda=\partial_{\mu}\langle A(\theta)\rangle$ and $\tfrac 1 2\bar\lambda'=\partial_{\sigma^2}\langle A(\theta)\rangle$.

Closed-form expressions for $\bar\lambda$ and $\bar\lambda'$ exist only in some special cases, for example if the firing-rate function $\rho(\cdot)$ is a (rectified) polynomial. We consider two choices of firing-rate nonlinearity which admit closed-form expressions, "exponential" and "probit".

-Choosing $\rho =\exp$ corresponds to a Poisson GLM. In this case, $\bar\lambda =\bar\lambda' =\exp(\mu +\sigma^2/2)$. -Let $\phi(\cdot)$ and $\Phi(\cdot)$ denote the probability density and cumulative distribution function, respectively, for a standard normal distribution. Choosing $\rho =\Phi$ corresponds to a probit GLM. In this case, $\bar\lambda =\Phi(\gamma\mu)$ and $\bar\lambda' =\gamma\phi(\gamma\mu)$, where $\gamma = (1+\sigma^2)^{-1}$.

For the probit firing-rate nonlinearity, we will also need to know $\partial_{\sigma^2}\langle\rho'(\boldsymbol\theta)\rangle$ to calculate the Hessian-vector product. In this case, $\rho'=\phi$. We have from $\eqref{dexpect}$ that $\partial_{\sigma^2}\langle\phi(x)\rangle=\tfrac1{2\sigma^2}\left<\theta (\mu-\theta)\phi(\theta)\right>$. This can be solved by writing the expectation as an integral and completing the square in the resulting Gaussian integral, yielding: \begin{equation}\begin{aligned} \partial_{\sigma^2}\langle\phi(x)\rangle &= \frac {u-1} {\sqrt{8\pi e^u(1+\sigma^2)^3}} ,\text{ where }u=\frac{\mu^2}{\sigma^2+1}. \end{aligned} \label{probitrhoprimeexpectgradient} \end{equation}

Derivatives of the loss function

With these prelimenaries out of the way, we can now consider the derivatives of $\eqref{loss}$ in terms of $\boldsymbol\mu_q$ and $\boldsymbol\Sigma_q$.

Derivatives in $\boldsymbol\mu_q$

The gradient and Hessian of $\mathcal L$ with respect to $\boldsymbol\mu_q$ are: \begin{equation}\begin{aligned} \partial_{\boldsymbol\mu_q} \mathcal L &= \boldsymbol\Sigma_z^{-1} (\boldsymbol\mu_q-\boldsymbol\mu_z) + \mathbf B^\top\left( \bar{\boldsymbol\lambda}-\mathbf y\right) \\ \operatorname H_{\boldsymbol\mu_q} \mathcal L &= \boldsymbol\Sigma_z^{-1} + \mathbf B^\top \operatorname{diag}[\bar{\boldsymbol\lambda}'] \mathbf B \end{aligned}\end{equation}

Gradient in $\boldsymbol\Sigma_q$

The gradient of $\eqref{loss}$ in $\boldsymbol\Sigma_q$ is more involved. The derivative of the term $\tfrac 1 2\{ \operatorname{tr}(\boldsymbol\Sigma_z^{-1}\boldsymbol\Sigma_q)+\ln|\boldsymbol\Sigma_z^{-1}\boldsymbol\Sigma_q|\}$ can be obtained using identities provided in The Matrix Cookbook . The derivative of $\mathbf 1^\top\langle A (\boldsymbol\theta)\rangle$ can be obtained by considering derivatives with respect to individual elements of $\boldsymbol\Sigma_q$, and is $\tfrac12\mathbf B^\top\operatorname{diag}[\bar\lambda']\mathbf B$. Overall, we find that: \begin{equation}\begin{aligned} \partial_{\boldsymbol\Sigma_q}{\mathcal L}= \tfrac 1 2 \left\{ \boldsymbol\Sigma_z^{-1}+ \boldsymbol\Sigma_q^{-\top}+ \mathbf B^\top\operatorname{diag}[\bar\lambda']\mathbf B \right\}. \end{aligned} \label{sigmagrad} \end{equation}

Hessian-vector product in $\boldsymbol\Sigma_q$

Since $\boldsymbol\Sigma_q$ is a matrix, the Hessian of $\eqref{loss}$ in $\boldsymbol\Sigma_q$ is a fourth-order tensor. It is simpler to work with the Hessian-vector product. Here, the "vector" is a covariance matrix $\mathbf M$ to be optimized. The Hessian-vector product is given by the following identity: \begin{equation}\begin{aligned} \langle\mathbf H_{\boldsymbol\Sigma_q},\mathbf M\rangle&= \partial_{\boldsymbol\Sigma_q}\langle \mathbf J_{\boldsymbol\Sigma_q},\mathbf M\rangle= \partial_{\boldsymbol\Sigma_q} \operatorname{tr}\left[ \mathbf J_{\boldsymbol\Sigma_q}^\top\mathbf M \right] \end{aligned}\end{equation}

where $\langle\cdot,\cdot\rangle$ denotes the scalar (Frobenius) product. The Hessian-vector product for the terms $\boldsymbol\Sigma_z^{-1}+\boldsymbol\Sigma_q^{-\top}$ in $\eqref{sigmagrad}$ can be obtained using identities provided in The Matrix Cookbook : \begin{equation}\begin{aligned} \partial_{\boldsymbol\Sigma_q} \operatorname{tr}\left[\left\{\boldsymbol\Sigma_z^{-1}+\boldsymbol\Sigma_q^{-\top}\right\}^\top\mathbf M\right]=-\boldsymbol\Sigma_q^{-1} \mathbf M^\top \boldsymbol\Sigma_q^{-1}. \end{aligned}\end{equation}

The Hessian-vector product for the term $\mathbf B^\top\operatorname{diag}[\bar\lambda']\mathbf B$ in $\eqref{sigmagrad}$ is more complicated. We can write \begin{equation}\begin{aligned} \partial_{\boldsymbol\Sigma_q} \operatorname{tr}\left[ \left\{ \mathbf B^\top\operatorname{diag}[\bar\lambda']\mathbf B \right\}^\top\mathbf M \right]&= \partial_{\boldsymbol\Sigma_q} \operatorname{tr}\left[ \mathbf B\mathbf M\mathbf B^\top \operatorname{diag}[\bar\lambda'] \right]\\&= \mathbf B^\top \operatorname{diag}[\mathbf B\mathbf M\mathbf B^\top] \operatorname{diag} \left[ \partial_{\sigma^2_{\boldsymbol\theta}}\langle\rho'(\boldsymbol\theta)\rangle \right] \mathbf B. \end{aligned} \label{hvp2} \end{equation}

The first step in $\eqref{hvp2}$ uses the fact that the trace is invariant under cyclic permutations. The second step follows from Lemma 1 (Appendix, below), with $\mathbf C=\mathbf B\mathbf M\mathbf B^\top$ and using the fact that $\bar\lambda'=\langle\rho'(\theta)\rangle$. In general, the Hessian-vector product in $\boldsymbol\Sigma_q$ is \begin{equation}\begin{aligned}\langle\mathbf H_{\boldsymbol\Sigma_q},\mathbf M\rangle&= \tfrac12\left\{-\boldsymbol\Sigma_q^{-1}\mathbf M^\top \boldsymbol\Sigma_q^{-1}+\mathbf B^\top\operatorname{diag}[\mathbf B\mathbf M\mathbf B^\top] \operatorname{diag}\left[\partial_{\sigma^2_{\boldsymbol\theta}}\langle\rho'(\boldsymbol\theta)\rangle \right]\mathbf B\right\} \end{aligned}\end{equation}

For the exponential firing-rate nonlinearity, $\partial_{\sigma^2_{\boldsymbol\theta}}\langle\rho'(\boldsymbol\theta)\rangle=\tfrac 1 2\bar\lambda$. The solution for the probit firing-rate nonlinearity is given in $\eqref{probitrhoprimeexpectgradient}$.

Conclude

That's all for now! I'll need to integrate these with the various other derivations ( e.g. see also here. ).

Appendix

Lemma 1

(We use Einstein summation to simplify the notation) \begin{equation}\begin{aligned} \partial_{\boldsymbol\Sigma_{q,ij}} \operatorname{tr} \left[ \mathbf C \operatorname{diag}\left[ \langle f(\boldsymbol\theta)\rangle \right] \right]&= \partial_{\boldsymbol\Sigma_{q,ij}} \left[ \mathbf C \operatorname{diag}\left[ \langle f(\boldsymbol\theta)\rangle \right] \right]_{kk}\\&= \partial_{\boldsymbol\Sigma_{q,ij}} \left[ \mathbf C_{lm} \operatorname{diag}\left[ \langle f(\boldsymbol\theta)\rangle \right]_{mn} \right]_{kk}\\&= \partial_{\boldsymbol\Sigma_{q,ij}} \left[ {\mathbf C}_{kk} \operatorname{diag}\left[ \langle f(\boldsymbol\theta)\rangle \right]_{k} \right]\\&= {\mathbf C}_{kk} \langle \partial_{\boldsymbol\Sigma_{q,ij}} f(\theta_k)\rangle\\&= {\mathbf C}_{kk} \mathbf B_{ik}^\top \,\partial_{\sigma^2_{\theta}}\langle f(\theta_k)\rangle \mathbf B_{kj}\\&= \mathbf B_{ik}^\top {\mathbf C}_{kk} \,\partial_{\sigma^2_{\theta}}\langle f(\theta_k)\rangle \mathbf B_{kj}\\&= \left\{ \mathbf B^\top \operatorname{diag}[\mathbf C] \operatorname{diag} \left[ \partial_{\sigma^2_{\boldsymbol\theta}}\langle f(\boldsymbol\theta)\rangle \right] \mathbf B \right\}_{ij} \\ \\ \partial_{\boldsymbol\Sigma_q} \operatorname{tr} \left[ \mathbf C \operatorname{diag}\left[ \langle f(\boldsymbol\theta)\rangle \right] \right]&= \mathbf B^\top \operatorname{diag}[\mathbf C] \operatorname{diag} \left[ \partial_{\sigma^2_{\boldsymbol\theta}}\langle f(\boldsymbol\theta)\rangle \right] \mathbf B \end{aligned}\end{equation}

Note: Derivatives of Gaussian KL-Divergence for some parameterizations of the posterior covariance for variational Gaussian-process inference

These notes provide the derivatives of the KL-divergence $D_{\text{KL}}\left[ Q(\mathbf z) \| P(\mathbf z)\right]$ between two multivariate Gaussian distributions $Q(\mathbf z)$ and $P(\mathbf z)$ with respect to a few parameterizations $\theta$ of the covariance matrix $\boldsymbol\Sigma(\theta)$ of $Q$. This is useful for variational Gaussian process inference, where clever parameterizations of the posterior covariance are required to make the problem tractable. Tables for differentiating matrix-valued functions can be found in The Matrix Cookbook .

Notes: Variability and oscillations

These notes outline a stochastic binary neural network as a toy model of the effects of noise on neural computation. We explore how noise modifies the effective gain in linear-nonlinear networks, and how oscillations might modulate some of the effects of noise.

Causes and consequences of representational drift

Drs. Harvey, O'Leary, and myself, have just published a review in Current Opinion in Neurobiology titled the "Causes and consequences of representational drift" [PDF].

We explore recent results from the work of Alon Rubin and Liron Sheintuch in the Ziv lab and Laura Driscoll in the Harvey lab. Their work has shown that neural representations reconfigure even for fixed, learned tasks. 

We discuss ways that the brain might support this reconfiguration without forgetting what it has already learned. There might be a subset of stable neurons that maintain memories. Alternatively, neural representations can be highly redundant, and support a stable readout even as single cells seem to change. 

Finally, we conjecture that redundancy allows different brain areas to "error-correct" each other, allowing the brain to keep track of the shifting meaning of single neurons in plastic representations.

For a brief preview, here are the pre-publication versions of figures 2 and 3:

Figure 2:

Figure 2: Internal representations have unconstrained degrees of freedom that allow drift. (a) Nonlinear dimensionality reduction of population activity recovers the low-dimensional structure of the T-maze in Driscoll et al. (2017). Each point represents a single time-point of population activity, and is colored according to location in the maze. (b) Point clouds illustrate low-dimensional projections of neural activity as in (a). Although unsupervised dimensionality-reduction methods can recover the task structure on each day, the way in which this structure is encoded in the population can change over days to weeks. (c) Left: Neural populations can encode information in relative firing rates and correlations, illustrated here as a sensory variable encoded in the sum of two neural signals ($y_1 + y_2$). Points represent neural activity during a repeated presentation of the same stimulus. Variability orthogonal to this coding axis does not disrupt coding, but could appear as drift in experiments if it occurred on slow timescales. Right: Such distributed codes may be hard to read-out from recorded subpopulations (e.g. $y_1$ or $y_2$ alone; black), especially if they entail correlations between brain areas. (d) Left: External covariates may exhibit context-dependent relationships. Each point here reflects a neural population state at a given time-point. The relationship between directions $x_1$ and $x_2$ changes depending on context (cyan versus red). Middle: Internally, this can be represented a mixture model, in which different subspaces are allocated to encode each context, and the representations are linearly-separable (gray plane). Right: The expanded representation contains two orthogonal subspaces that each encode a separate, context- dependent relationship. This dimensionality expansions increases the degrees of freedom in internal representations, thereby increasing opportunities for drift.

 

Note: Training stochastic neural networks

Feed-forward neural networks consist of a series of layers. In each layer, outputs from past layers are combined linearly, then passed through some nonlinear transformation. As long as all computations are differentiable, the entire network is differentiable as well. This allows artificial neural networks to be trained using gradient-based optimization techniques (backpropagation).

Methods for training stochastic networks via backpropagation are less well developed, but solutions exist and are the subject of ongoing research (c.f. Rezende et al. 2014 and the numerous papers that cite it). In the context of models of neural computation, Echeveste et al. (2019) trained stochastic neural networks with rectified-polynomial nonlinearities.