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 $\mathit{\lambda }=\rho (\mathit{\theta })$, where $\rho (\cdot )$ is a firing-rate nonlinearity and $\mathit{\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 $\mathit{\theta }$ can be written as $$\begin{array}{}\text{(1)}& \begin{array}{rl}\ln Pr(\mathbf{y}|\mathbf{z})& ={\mathbf{y}}^{\mathrm{\top }}\mathit{\theta }-{\mathbf{1}}^{\mathrm{\top }}A(\mathit{\theta })+\text{constant},\end{array}\end{array}$$

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

Assume that the synaptic activations $\mathit{\theta }$ are driven by shared latent variables $\mathbf{z}$ with a Gaussian prior $\mathbf{z}\sim \mathcal{N}({\mathit{\mu }}_z,{\mathbf{\Sigma }}_z)$. Let $\mathit{\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 $\mathit{z}$ is approximated as Gaussian, i.e. $Pr(\mathbf{z}|\mathbf{y})\approx Q(\mathbf{z})$, where $Q(\mathbf{z})=\mathcal{N}({\mathit{\mu }}_q,{\mathbf{\Sigma }}_q)$. We then optimize ${\mathit{\mu }}_q$ and ${\mathbf{\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}}[Q(\mathbf{z})\Vert Pr(\mathbf{z})]$ from the prior to the posterior, while maximizing the expected log-likelihood $⟨Pr(\mathbf{y}|\mathbf{z})⟩$: $$\begin{array}{}\text{(2)}& \begin{array}{rl}{D}_{\text{KL}}[Q(\mathbf{z})\Vert Pr(\mathbf{z}|\mathbf{y})]& =D_{\text{KL}}[Q(\mathbf{z})\Vert Pr(\mathbf{z})]-⟨\ln Pr(\mathbf{y}|\mathbf{z})⟩+\text{constant}.\end{array}\end{array}$$

(In these notes, all expectations $⟨\cdot ⟩$ 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}}[Q(\mathbf{z})\Vert Pr(\mathbf{z})]$ has the closed form : $$\begin{array}{}\text{(3)}& \begin{array}{rl}{D}_{\text{KL}}[Q(\mathbf{z})\Vert Pr(\mathbf{z})]& =\frac{1}{2}\{({\mathit{\mu }}_z-{\mathit{\mu }}_q{)}^{\mathrm{\top }}{\mathbf{\Sigma }}_z^{-1}({\mathit{\mu }}_z-{\mathit{\mu }}_q)+\mathrm{tr}\left({\mathbf{\Sigma }}_z^{-1}{\mathbf{\Sigma }}_q\right)+\ln |{\mathbf{\Sigma }}_z^{-1}{\mathbf{\Sigma }}_q|\}+\text{constant}.\end{array}\end{array}$$

For our choice of the canonically-parameterized natural exponential family, the expected negative log-likelihood can be written as: $$\begin{array}{}\text{(4)}& \begin{array}{rl}-⟨\ln Pr(\mathbf{y}|\mathbf{z})⟩& ={\mathbf{1}}^{\mathrm{\top }}⟨A(\mathit{\theta })⟩-{\mathbf{y}}^{\mathrm{\top }}\mathbf{B}{\mathit{\mu }}_q+\text{constant}.\end{array}\end{array}$$

Neglecting constants and terms that do not depend on $({\mathit{\mu }}_q,{\mathbf{\Sigma }}_q)$, the overall loss function to be minimized is: $$\begin{array}{}\text{(5)}& \begin{array}{rl}\mathcal{L}({\mathit{\mu }}_q,{\mathbf{\Sigma }}_q)& =\frac{1}{2}\{({\mathit{\mu }}_z-{\mathit{\mu }}_q{)}^{\mathrm{\top }}{\mathbf{\Sigma }}_z^{-1}({\mathit{\mu }}_z-{\mathit{\mu }}_q)+\mathrm{tr}\left({\mathbf{\Sigma }}_z^{-1}{\mathbf{\Sigma }}_q\right)+\ln |{\mathbf{\Sigma }}_z^{-1}{\mathbf{\Sigma }}_q|\}+{\mathbf{1}}^{\mathrm{\top }}⟨A(\mathit{\theta })⟩-{\mathbf{y}}^{\mathrm{\top }}\mathbf{B}{\mathit{\mu }}_q\end{array}.\end{array}$$

Closed-form expectations

To optimize $\text{(5)}$, we need to differentiate it in ${\mathit{\mu }}_q$ and ${\mathbf{\Sigma }}_q$. These derivatives are mostly straightforward, but the expectation $⟨A(\mathit{\theta })⟩$ poses difficulties when $A(\cdot )$ is nonlinear. We'll consider some choices of firing-rate nonlinearity for which the derivatives of $⟨A(\mathit{\theta })⟩$ have closed-form expressions when $\mathit{\theta }$ is Gaussian.

Because we've assumed a Gaussian posterior on our latent state $\mathbf{z}$, and since $\mathit{\theta }=\mathbf{B}\mathbf{z}$, the synaptic activations $\mathit{\theta }$ are also Gaussian. The vectors ${\mathit{\mu }}_{\theta }$ and ${\mathit{\sigma }}_{\theta }^2$ for the mean and variance of $\mathit{\theta }$, respectively, are: $$\begin{array}{}\text{(6)}& \begin{array}{rl}{\mathit{\mu }}_{\theta }& =\mathbf{B}{\mathit{\mu }}_q\\ {\mathit{\sigma }}_{\theta }^2& =\mathrm{diag}\left[\mathbf{B}{\mathbf{\Sigma }}_q{\mathbf{B}}^{\mathrm{\top }}\right]\end{array}\end{array}$$

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 }⟨A(\theta )⟩$ and ${\partial }_{{\sigma }^2}⟨A(\theta )⟩$, with respect to $\mu$ and ${\sigma }^2$ respectively, are: $$\begin{array}{}\text{(7)}& \begin{array}{rl}{\partial }_{\mu }⟨A(\theta )⟩& =⟨A^{\prime }(\theta )⟩=⟨\rho (\theta )⟩\\ {\partial }_{{\sigma }^2}⟨A(\theta )⟩& =\frac{1}{2{\sigma }^2}⟨(\theta -{\mu }_{\theta })A^{\prime }(\theta )⟩=\frac{1}{2{\sigma }^2}⟨(\theta -\mu )\rho (\theta )⟩.\end{array}\end{array}$$

For more compact notation, denote the expected firing rate as $\overline{\lambda }=⟨\rho (\theta )⟩$, and denote the expected derivative of the firing-rate in $\theta$ as ${\overline{\lambda }}^{\prime }=⟨{\rho }^{\prime }(\theta )⟩$. Note that $\overline{\lambda }={\partial }_{\mu }⟨A(\theta )⟩$ and $\frac{1}{2}{\overline{\lambda }}^{\prime }={\partial }_{{\sigma }^2}⟨A(\theta )⟩$.

Closed-form expressions for $\overline{\lambda }$ and ${\overline{\lambda }}^{\prime }$ 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, $\overline{\lambda }={\overline{\lambda }}^{\prime }=\exp (\mu +{\sigma }^2/2)$. -Let $\varphi (\cdot )$ and $\mathrm{\Phi }(\cdot )$ denote the probability density and cumulative distribution function, respectively, for a standard normal distribution. Choosing $\rho =\mathrm{\Phi }$ corresponds to a probit GLM. In this case, $\overline{\lambda }=\mathrm{\Phi }(\gamma \mu )$ and ${\overline{\lambda }}^{\prime }=\gamma \varphi (\gamma \mu )$, where $\gamma =(1+{\sigma }^2{)}^{-1}$.

For the probit firing-rate nonlinearity, we will also need to know ${\partial }_{{\sigma }^2}⟨{\rho }^{\prime }(\mathit{\theta })⟩$ to calculate the Hessian-vector product. In this case, ${\rho }^{\prime }=\varphi$. We have from $\text{(7)}$ that ${\partial }_{{\sigma }^2}⟨\varphi (x)⟩=\frac{1}{2{\sigma }^2}⟨\theta (\mu -\theta )\varphi (\theta )⟩$. This can be solved by writing the expectation as an integral and completing the square in the resulting Gaussian integral, yielding: $$\begin{array}{}\text{(8)}& \begin{array}{rl}{\partial }_{{\sigma }^2}⟨\varphi (x)⟩& =\frac{u-1}{\sqrt{8\pi e^u(1+{\sigma }^2{)}^3}},\text{ where }u=\frac{{\mu }^2}{{\sigma }^2+1}.\end{array}\end{array}$$

Derivatives of the loss function

With these prelimenaries out of the way, we can now consider the derivatives of $\text{(5)}$ in terms of ${\mathit{\mu }}_q$ and ${\mathbf{\Sigma }}_q$.

Derivatives in ${\mathit{\mu }}_q$

The gradient and Hessian of $\mathcal{L}$ with respect to ${\mathit{\mu }}_q$ are: $$\begin{array}{}\text{(9)}& \begin{array}{rl}{\partial }_{{\mathit{\mu }}_q}\mathcal{L}& ={\mathbf{\Sigma }}_z^{-1}({\mathit{\mu }}_q-{\mathit{\mu }}_z)+{\mathbf{B}}^{\mathrm{\top }}(\overline{\mathit{\lambda }}-\mathbf{y})\\ {\mathrm{H}}_{{\mathit{\mu }}_q}\mathcal{L}& ={\mathbf{\Sigma }}_z^{-1}+{\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[{\overline{\mathit{\lambda }}}^{\prime }]\mathbf{B}\end{array}\end{array}$$

Gradient in ${\mathbf{\Sigma }}_q$

The gradient of $\text{(5)}$ in ${\mathbf{\Sigma }}_q$ is more involved. The derivative of the term $\frac{1}{2}\{\mathrm{tr}({\mathbf{\Sigma }}_z^{-1}{\mathbf{\Sigma }}_q)+\ln |{\mathbf{\Sigma }}_z^{-1}{\mathbf{\Sigma }}_q|\}$ can be obtained using identities provided in The Matrix Cookbook . The derivative of ${\mathbf{1}}^{\mathrm{\top }}⟨A(\mathit{\theta })⟩$ can be obtained by considering derivatives with respect to individual elements of ${\mathbf{\Sigma }}_q$, and is $\frac{1}{2}{\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[{\overline{\lambda }}^{\prime }]\mathbf{B}$. Overall, we find that: $$\begin{array}{}\text{(10)}& \begin{array}{r}{\partial }_{{\mathbf{\Sigma }}_q}\mathcal{L}=\frac{1}{2}\{{\mathbf{\Sigma }}_z^{-1}+{\mathbf{\Sigma }}_q^{-\mathrm{\top }}+{\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[{\overline{\lambda }}^{\prime }]\mathbf{B}\}.\end{array}\end{array}$$

Hessian-vector product in ${\mathbf{\Sigma }}_q$

Since ${\mathbf{\Sigma }}_q$ is a matrix, the Hessian of $\text{(5)}$ in ${\mathbf{\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{array}{}\text{(11)}& \begin{array}{rl}⟨{\mathbf{H}}_{{\mathbf{\Sigma }}_q},\mathbf{M}⟩& ={\partial }_{{\mathbf{\Sigma }}_q}⟨{\mathbf{J}}_{{\mathbf{\Sigma }}_q},\mathbf{M}⟩={\partial }_{{\mathbf{\Sigma }}_q}\mathrm{tr}\left[{\mathbf{J}}_{{\mathbf{\Sigma }}_q}^{\mathrm{\top }}\mathbf{M}\right]\end{array}\end{array}$$

where $⟨\cdot ,\cdot ⟩$ denotes the scalar (Frobenius) product. The Hessian-vector product for the terms ${\mathbf{\Sigma }}_z^{-1}+{\mathbf{\Sigma }}_q^{-\mathrm{\top }}$ in $\text{(10)}$ can be obtained using identities provided in The Matrix Cookbook : $$\begin{array}{}\text{(12)}& \begin{array}{r}{\partial }_{{\mathbf{\Sigma }}_q}\mathrm{tr}\left[{\{{\mathbf{\Sigma }}_z^{-1}+{\mathbf{\Sigma }}_q^{-\mathrm{\top }}\}}^{\mathrm{\top }}\mathbf{M}\right]=-{\mathbf{\Sigma }}_q^{-1}{\mathbf{M}}^{\mathrm{\top }}{\mathbf{\Sigma }}_q^{-1}.\end{array}\end{array}$$

The Hessian-vector product for the term ${\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[{\overline{\lambda }}^{\prime }]\mathbf{B}$ in $\text{(10)}$ is more complicated. We can write $$\begin{array}{}\text{(13)}& \begin{array}{rl}{\partial }_{{\mathbf{\Sigma }}_q}\mathrm{tr}\left[{\{{\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[{\overline{\lambda }}^{\prime }]\mathbf{B}\}}^{\mathrm{\top }}\mathbf{M}\right]& ={\partial }_{{\mathbf{\Sigma }}_q}\mathrm{tr}[\mathbf{B}\mathbf{M}{\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[{\overline{\lambda }}^{\prime }]]\\ & ={\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[\mathbf{B}\mathbf{M}{\mathbf{B}}^{\mathrm{\top }}]\mathrm{diag}[{\partial }_{{\sigma }_{\mathit{\theta }}^2}⟨{\rho }^{\prime }(\mathit{\theta })⟩]\mathbf{B}.\end{array}\end{array}$$

The first step in $\text{(13)}$ 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}}^{\mathrm{\top }}$ and using the fact that ${\overline{\lambda }}^{\prime }=⟨{\rho }^{\prime }(\theta )⟩$. In general, the Hessian-vector product in ${\mathbf{\Sigma }}_q$ is $$\begin{array}{}\text{(14)}& \begin{array}{rl}⟨{\mathbf{H}}_{{\mathbf{\Sigma }}_q},\mathbf{M}⟩& =\frac{1}{2}\{-{\mathbf{\Sigma }}_q^{-1}{\mathbf{M}}^{\mathrm{\top }}{\mathbf{\Sigma }}_q^{-1}+{\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[\mathbf{B}\mathbf{M}{\mathbf{B}}^{\mathrm{\top }}]\mathrm{diag}[{\partial }_{{\sigma }_{\mathit{\theta }}^2}⟨{\rho }^{\prime }(\mathit{\theta })⟩]\mathbf{B}\}\end{array}\end{array}$$

For the exponential firing-rate nonlinearity, ${\partial }_{{\sigma }_{\mathit{\theta }}^2}⟨{\rho }^{\prime }(\mathit{\theta })⟩=\frac{1}{2}\overline{\lambda }$. The solution for the probit firing-rate nonlinearity is given in $\text{(8)}$.

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{array}{}\text{(15)}& \begin{array}{rl}{\partial }_{{\mathbf{\Sigma }}_{q,ij}}\mathrm{tr}\left[\mathbf{C}\mathrm{diag}[⟨f(\mathit{\theta })⟩]\right]& ={\partial }_{{\mathbf{\Sigma }}_{q,ij}}{\left[\mathbf{C}\mathrm{diag}[⟨f(\mathit{\theta })⟩]\right]}_{kk}\\ & ={\partial }_{{\mathbf{\Sigma }}_{q,ij}}{\left[{\mathbf{C}}_{lm}\mathrm{diag}{[⟨f(\mathit{\theta })⟩]}_{mn}\right]}_{kk}\\ & ={\partial }_{{\mathbf{\Sigma }}_{q,ij}}\left[{\mathbf{C}}_{kk}\mathrm{diag}{[⟨f(\mathit{\theta })⟩]}_k\right]\\ & ={\mathbf{C}}_{kk}⟨{\partial }_{{\mathbf{\Sigma }}_{q,ij}}f({\theta }_k)⟩\\ & ={\mathbf{C}}_{kk}{\mathbf{B}}_{ik}^{\mathrm{\top }}{\partial }_{{\sigma }_{\theta }^2}⟨f({\theta }_k)⟩{\mathbf{B}}_{kj}\\ & ={\mathbf{B}}_{ik}^{\mathrm{\top }}{\mathbf{C}}_{kk}{\partial }_{{\sigma }_{\theta }^2}⟨f({\theta }_k)⟩{\mathbf{B}}_{kj}\\ & ={\{{\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[\mathbf{C}]\mathrm{diag}[{\partial }_{{\sigma }_{\mathit{\theta }}^2}⟨f(\mathit{\theta })⟩]\mathbf{B}\}}_{ij}\\ \\ {\partial }_{{\mathbf{\Sigma }}_q}\mathrm{tr}\left[\mathbf{C}\mathrm{diag}[⟨f(\mathit{\theta })⟩]\right]& ={\mathbf{B}}^{\mathrm{\top }}\mathrm{diag}[\mathbf{C}]\mathrm{diag}[{\partial }_{{\sigma }_{\mathit{\theta }}^2}⟨f(\mathit{\theta })⟩]\mathbf{B}\end{array}\end{array}$$

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}}[Q(\mathbf{z})\Vert P(\mathbf{z})]$ between two multivariate Gaussian distributions $Q(\mathbf{z})$ and $P(\mathbf{z})$ with respect to a few parameterizations $\theta$ of the covariance matrix $\mathbf{\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 .