Using A Kalman Filter for Neural Decoding

The task is to decode the movement (x and y position + x and y velocity) of a monkey moving a cursor based on spiking activities of some neurons.

The principle we follow is that precision adds. When we have two independent and noisy measurements of a truth, we can add them with a certain weight and end up with a measurement that is more precise than both raw measurements. This is called a Kalman filter, and it looks mathematically identical to the EMA that I used to do velocity smoothing for a noisy velocity measurement in the inverted pendulum bot.

$$\hat{z}=Kz_1+(1-K)z_2$$

Let’s try to prove that $\hat{z}$ has a higher precision than both $z_1$ and $z_2$ - i.e. the variance of $\hat{z}$ is lower than both of them.

Let $\mu$ be $E[z_1]=E[z_2]$ (Since both are noisy measurements around the same ground truth, no bias).

$$\begin{aligned} \mathrm{Var}\big(Kz_1+(1-K)z_2\big) &= \mathrm{Var}(Kz_1)+\mathrm{Var}\big((1-K)z_2\big) &&\text{independent, so variances add}\\ &= K^2\,\mathrm{Var}(z_1)+(1-K)^2\,\mathrm{Var}(z_2) &&\text{a constant comes out squared}\\ &= K^2\sigma_1^2+(1-K)^2\sigma_2^2 \end{aligned}$$

Differentiate with respect to $K$ and set to zero to minimize the variance:

$$\begin{aligned} \frac{d}{dK}\Big[K^2\sigma_1^2+(1-K)^2\sigma_2^2\Big] &= 2K\sigma_1^2-2(1-K)\sigma_2^2 \\ 2K\sigma_1^2-2(1-K)\sigma_2^2 &= 0 \\ K\big(\sigma_1^2+\sigma_2^2\big) &= \sigma_2^2 \\ K &= \frac{\sigma_2^2}{\sigma_1^2+\sigma_2^2} \end{aligned}$$

What does this remind you of? E&M was only a semester ago for me so this jumps out as a parallel circuit’s resistance.

$$R_{combined}=\frac{R_1R_2}{R_1+R_2}$$

Then, naturally, as the combined resistance is lower for parallel circuits than each resistor’s, so does the variance of the blended measurement. There is a shared mechanism behind both. Circuits also minimize something - power.

$$P=I^2R=I_1^2R_1+I_2^2R_2=I_1^2R_1+(I-I_1)^2R_2 \\ I_1 = I\frac{R_2}{R_1+R_2}$$

And analogous to conductance, the inverses of their variance, as we mentioned, adds up to the combined precision (of $\hat{z}$):

$$\frac{1}{\sigma_{combined}^2}=\frac{1}{\sigma_1^2}+\frac{1}{\sigma_2^2}$$

We have proven that precision adds when we blend two independent and noisy raw measurements with the optimal weight $\frac{\sigma_2^2}{\sigma_1^2+\sigma_2^2}$, their precisions add. We end up with a better measurement than both. Yet, in the case of neural decoding, it looks like we only have one raw measurement - the spikes. So here comes the cool trick:

We get to use the prediction of the previous timestep as an independent input.

$$\hat{x}=Kz+(1-K)x$$

This is similar to a Wiener filter, where we directly use the previous $n$ bins to inform the current one. The reason we can do this is because motion is continuous, and momentum determines a large portion of the following movement.

First Half: Momentum

Let’s assume a linear relationship between the previous state and the current state (why?).

$$\hat{x}_t^- = a\hat{x}_{t-1}$$

But it is not always linear. There is noise, and the monkey changes its movement all the time. Let’s call that error that we cannot account for $w$ with variance $q$.

$$x_{t} = ax_{t-1}+w \\$$ $$\mathrm{Var}(w) = q$$

$w$ and $q$ are properties of the dataset (with regard to $a$) and are fixed. They describe how un-linear (if you get what I mean) the true data is. Furthermore, since when we decode, we don’t have the ground truth $x$ but only the previous prediction $\hat{x}$. That premise could be wrong too. At each timestep, the prediction we make has a varying error, let’s call its variance $P$. It evolves over time by:

$$P^-=a^2P+q$$

$q$ comes from the un-linear noise. and $a^2$ is how the previous error $P$ scale the current one.

Second Half: Spikes

This is not really a half, more like an eighth. The spikes is more straightforward than momentum, since there is no error carrying forward. Let’s just assume the spike give us a similarly noisy and independent measurement of the state $z$ with variance $P_z$.

The Magic Blender

Now, as we proved earlier, we can now find the optimal $K$ to blend the two halves.

$$\begin{aligned} K &= \frac{P^-}{P^-+P_z} \\ \hat{x} &= Kz + (1-K)\hat{x}^- \\ P &= \frac{1}{\frac{1}{P^-}+\frac{1}{P_z}} \end{aligned}$$

At this point, no one can keep track of the symbols anymore, so here is how the process goes:

1. We start with nothing but the first spikes. Translate them to states. Gets $z$, and using it as our first estimated states $\hat{x}$. $P=P_z$.
2. Moving on to the next timestep. We now have the two halves. Before we receive the second spikes, evolve the previous prediction first by our momentum guess: $\hat{x}^- = a\hat{x}, P^-=a^2P+q$.
3. Receiving the spikes $z$, we now do the magical blend: $\hat{x} = Kz + (1-K)\hat{x}^-,P = \frac{1}{\frac{1}{P^-}+\frac{1}{P_z}}$.
4. Go back to step 2 unless no more timestep left.

The superscript $(\quad)^-$ denotes step 2, the intermediate estimate and error solely from momentum, before receiving the next spikes.

The Matrix

Welcome to the real world, Neo.

In reality, we have not 1, but 4 numbers representing the state of a cursor. its x-position, y-position, x-velocity, and y-velocity. We also don’t automatically get the spike’s prediction for the states. All we get is a string of processed data that tells us in each timestep, whether a neuron spiked or not. In my dataset, there are 223 such neurons.

Let us try to map the scalar versions of our equations over.

$$\hat{x}^- = A\hat{x}$$

Here, $\hat{x}^-$ and $\hat{x}$ are vectors of height 4. and $A$ is a $4 \times 4$ matrix. Let $w$ be the error that $A$ cannot account for, and $Q$ be its covariance - a 4 $\times 4$ matrix.

Before we look at what covariance is. Let’s appreciate the fact that we have introduced complexity on our motion model. We used to have a single state in the scalar version - velocity. Now we have four. So each state becomes a linear combination of all the previous states: $a_1p_x+a_2p_y+a_3v_x+a_4v_y$. E.g., the estimated x-position may be $\hat{p}^-=p_x+\Delta t v_x$.

If we try to again find how the variance of their error behave, we are adding the error of the four states. If the states were independent, we can simply add their variances like we did before. However, the four states are clearly not independent. That is why we need to look at covariance.

Covariance

Recall how variance is the expected value of the squared difference between a random variable and its mean. The reason we squared it was to keep it positive, while keeping the algebra easy (unlike absolute value).

How do we compute the variance of a sum of two random variables? If they are independent, we know to just add their repective variances. What if they are not independent? Let’s expand and see.

Let $u$ be $X-E[X]$ and $v$ be $Y-E[Y]$.

$$\begin{aligned} \mathrm{Var}(X+Y)&=E[(u+v)^2] \\ &=E[u^2]+E[v^2]+2E[uv] \\ &=E[(X-E[X])^2]+E[(Y-E[Y])^2]+2E[(X-E[X])(Y-E[Y])] \end{aligned}$$

When $X$ and $Y$ are independent, $E[(X-E[X])(Y-E[Y])]$ is $0$. Since $E[X-E[X]]$ and $E[Y-E[Y]]$ are both $0$. When they are not independent, we get this term $2E[(X-E[X])(Y-E[Y])]$ to adjust.

$$\mathrm{Cov}(X,Y)=E[(X-E[X])(Y-E[Y])]$$

Covariance is about how two random variables tend to vary together. Consider two random variables, $X$ and $Y$. Let them only take values of either $1$ or $-1$. If we sum them, there are 4 possible scenarios:

When the covariance is high (and positive), the variables tend to take similar values (in our case, the same value), the sum of the two variables have a large variance. Otherwise, they tend to take different values, the positive and negative cancels out and we tend to have a small variance (or in this case, 0).

Heading back into our quest of finding the variance of $\hat{x}^-$. With the added covariance to the equation, can we still elegantly represent the uncertainty of the prediction the way $P^-=a^2P+q$ did in the scalar version?

Note how the covariance of a single variable is the variance of the variable. We can absolutely use a matrix-vector multiplication to express how the variance of $\hat{x}$ itself propagates through $A$. We may then add a term $Q$, corresponding to the scalar $q$, for the variance of the error. $A$ also gets squared through covariance operator.

$$\mathrm{Cov}(A\hat{x})=A\mathrm{Cov}(\hat{x})A^\top$$ $$P^-=APA^\top+Q$$

Looking back at the scalar version of the Kalman filter blend, we try to do something similar for the matrices:

$$\hat{x}=Kz+(1-K)\hat{x}^-$$

Yet, remember in the beginning we mentioned how we don’t automatically get the spike’s prediction for the states. We only get the spikes. The intuition would be to decode the spikes directly based on their directional preferences. While that fits our scalar version better, it introduces some complications. If we fitted the decoder in the training block of the dataset, we end up with a decoder that rely partly on the spikes, and partly with knowledge about momentum - that breaks our rules that the two measurements have to be independent. We need to instead translate the states back into spikes - encoding it.

$$z=Hx+e$$

Note that $z$ is now the spikes we observe. $x$ is the true states, and $H$ is a matrix we can fit. For each neuron, we can fit its cosine tuning curve against the training block’s data. We then can calculate the expected spikes for each direction for each neuron.

Now, last step: $K$.

In the scalar version we had:

$$K= \frac{P^-}{P^-+P_z}$$

We already have $P^-$, but what is $P_z$ now? It was the uncertainty in the spikes’ prediction of the state earlier. We assumed we knew it, but we don’t. We only have the spikes now. Of course, neurons have noise too. Their spikes may scatter around the expected value (Poisson distribution). Let us call the covariance matrix of each neuron’s scattering $R$.

Let’s look at one neuron and one state for now, say y-velocity. We model a neuron as direction-selective. If a neuron is sensitive to upward movements, say it fires $h$ more spikes per unit of upward velocity, when we translate 3 unit of upward velocity to spikes of this neuron, we would expect $3h$ spikes above its baseline. $z=hx+\text{noise}$. Where $z$, $h$ and $x$ are all scalars.

$$\mathrm{Var}(\text{noise})=r$$ $$\hat{x}=\frac{z}{h}$$ $$\mathrm{Var}(\hat{x})=\left(\frac{1}{h}\right)^2\mathrm{Var}(z)$$

$z$ only varies because of the noise, since $hx$ is fixed.

$$\mathrm{Var}(\hat{x})=\frac{r}{h^2}$$

That all looks innocent and simple, but we have 4 states and 223 neurons, not 1 and 1. For us to add $P^-$ and the spikes’ uncertainty, they need to mean uncertainty in the same space. $P^-$ is already in the state space with 4 dimensions. We need the other uncertainty to also be in the state space. Let’s go from the ground up again. We have the observed spikes $z$, and we want to translate it to state space and find the variance of the error. With noise in each neuron’s spikes, each neuron disagrees with each other in what the state is. Let’s try to use the good old derivative and see if we can minimize the squared error to get a good estimate instead.

The loss $L$ has to be calculated in spike space, since we do not get the true state during decoding. They also have to be weighted by the precision of each neuron $\frac{1}{r_i}$.

$$L(x)=\sum_{i=1}^{223}\frac{(z_i-h_ix)^2}{r_i}$$

The $x$ inside $L$ is a candidate state I am searching over - not the true state, and not the answer. It is the dial I turn. The one value of it that makes $L$ smallest is the estimate, and I write that one $\hat{x}_z$: the state the spikes vote for.

$$\begin{aligned} \frac{dL}{dx}=\sum_{i=1}^{223}\frac{-2h_i(z_i-h_ix)}{r_i}&=0 \\ x\sum_{i=1}^{223}\frac{h_i^2}{r_i} &= \sum_{i=1}^{223}\frac{h_iz_i}{r_i} \\ x &= \frac{\sum_{i=1}^{223}\frac{h_iz_i}{r_i}}{\sum_{i=1}^{223}\frac{h_i^2}{r_i}} \end{aligned}$$

How absolutely ugly!

Note that everything here are scalars, this is for a single state. Let’s make it 4.

$$L(x)=\sum_{i=1}^{223}\frac{(z_i-\sum_{j=1}^4h_{ij}x_j)^2}{r_i}$$ $$\begin{aligned} \frac{\partial L}{\partial x_k}=\sum_{i=1}^{223}\frac{-2h_{ik}(z_i-\sum_{j=1}^4h_{ij}x_j)}{r_i}&=0 \\ \sum_{i=1}^{223}\frac{h_{ik}\sum_{j=1}^4h_{ij}x_j}{r_i}&=\sum_{i=1}^{223}\frac{h_{ik}z_i}{r_i} \end{aligned}$$

This holds for each $k=1, 2, 3, 4$.

This still looks complicated, but something looks dangerously close to how Axler defines a matrix. Let’s try the right hand side first.

$$\begin{aligned} \sum_{i=1}^{223}\frac{h_{ik}z_i}{r_i}&=\sum_{i=1}^{223}h_{ik}\frac{z_i}{r_i} \\ &= (H^{\top}R^{-1}z)_k \end{aligned}$$

$\frac{z_i}{r_i}=(R^{-1}z)_i$ requires us to assume $R$ is diagonal (which is probably not true, but a trade-off I am taking).

$$\begin{aligned} \sum_{i=1}^{223}\frac{h_{ik}\sum_{j=1}^4h_{ij}x_j}{r_i}&=\sum_{i=1}^{223}\frac{h_{ik}(Hx)_i}{r_i} \\ &= (H^\top R^{-1}Hx)_k \end{aligned}$$

So nice.

$$\begin{aligned} (H^{\top}R^{-1}z)_k &= (H^\top R^{-1}Hx)_k \\ H^{\top}R^{-1}z &= H^\top R^{-1}Hx \\ \hat{x}_z &= (H^\top R^{-1}H)^{-1}H^{\top}R^{-1}z \end{aligned}$$

$\hat{x}_z$ is the state that would minimize the loss function $L$. It is parallel to $\hat{x}^-$, which is the momentum half’s guess of the state.

We are one last step from getting $K$! We still need the covariance of $\hat{x}_z$ as well.

$$\begin{aligned} \mathrm{Cov}\left[(H^\top R^{-1}H)^{-1}H^{\top}R^{-1}z \right]&=(H^\top R^{-1}H)^{-1}H^{\top}R^{-1}\mathrm{Cov}(z)[(H^\top R^{-1}H)^{-1}H^{\top}R^{-1}]^\top \\ &= (H^\top R^{-1}H)^{-1}H^{\top}R^{-1}R[(H^\top R^{-1}H)^{-1}H^{\top}R^{-1}]^\top \end{aligned}$$

This is just $\mathrm{Cov}(Ax)=A \mathrm{Cov}(x)A^\top$.

Notice $R$ is the covariance and therefore symmetric (covariance between $a$ and $b$ is equivalent to covariance between $b$ and $a$). Its transpose equals itself.

$$\begin{aligned} \mathrm{Cov}(\hat{x}_z) &= (H^\top R^{-1}H)^{-1}H^{\top}R^{-1}R[(H^\top R^{-1}H)^{-1}H^{\top}R^{-1}]^\top \\ & = (H^\top R^{-1}H)^{-1}H^{\top}R^{-1}R(R ^{-1}H(H^\top R^{-1}H) ^{-1}) \\ &= (H^\top R^{-1}H)^{-1}H^{\top}R^{-1}H(H^\top R^{-1}H) ^{-1} \\ &= (H^\top R^{-1}H) ^{-1} \end{aligned}$$

So incredibly nice. This matches our 1-neuron-1-state version, where we claimed $\mathrm{Var}(\hat{x})=\frac{r}{h^2}$. And this is the $P_z$ that the scalar filter simply assumed we had - we have finally built it, out of $H$ and $R$.

We are now ready for $K$. Let’s look at the scalar version and guess an analogy:

$$K = P^-(P_z+P^-)^{-1}$$ $$\begin{aligned} \hat{x} &= \hat{x}^- + K(\hat{x}_z - \hat{x} ^-) \\ &= K\hat{x}_z+(I-K)\hat{x}^- \end{aligned}$$

Can we really do this though? I want to look at the total variance after the blend and see if it follows our theme of how precision adds.

$$\begin{aligned} \mathrm{Cov}(\hat{x}) &= KP_zK^\top + (I-K)P^-(I-K)^\top \\ \end{aligned}$$

Let $T=(P_z+P^-)$ for clarity.

$$\begin{aligned} I-K &= I-P^-T^{-1} \\ &= TT^{-1}-P^-T^{-1} \\ &= (T-P^-)T^{-1} \\ &= P_zT^{-1} \end{aligned}$$

Again, $P^-$ and $P_z$ are covariances and therefore are symmetric. $T$ is too, as the sum of two symmetric matrices.

$$\begin{aligned} (P^-T^{-1})P_zT^{-1}P^- + P_zT^{-1}P^-T^{-1}P_z &= [(T-P_z)T^{-1}]P_zT^{-1}(T-P_z) + P_zT^{-1}(T-P_z)T^{-1}P_z \\ &= (I-P_zT^{-1})P_z[I-(T^{-1}P_z)] + P_zT^{-1}(I-P_zT^{-1})P_z \\ &= (P_z-P_zT^{-1}P_z)[I-(T^{-1}P_z)] + P_zT^{-1}(P_z-P_zT^{-1}P_z) \\ &= P_z-P_z(T^{-1}P_z)-P_zT^{-1}P_z+P_zT^{-1}P_zT^{-1}P_z + P_zT^{-1}P_z-P_zT^{-1}P_zT^{-1}P_z \\ &= P_z-P_z(T^{-1}P_z) \\ &= P_z(I-T^{-1}P_z) \\ &= P_z(T^{-1}T-T^{-1}P_z) \\ &= P_z[T^{-1}(T-P_z)] \\ &= P_zT^{-1}P^- \\ &= P_z(P_z+P^-)^{-1}P^- \\ \end{aligned}$$

Let’s call $P=\mathrm{Cov}(\hat{x})$, and invert it to get precision of $\hat{x}$ instead of covariance.

$$\begin{aligned} P^{-1} &= (P_z(P_z+P^-)^{-1}P^-)^{-1} \\ &= (P^{-})^{-1}(P_z+P^-)P_z^{-1} \\ &= [(P^{-})^{-1}P_z+(P^{-})^{-1}P^-]P_z^{-1} \\ &= (P^{-})^{-1}+P_z^{-1} \\ \end{aligned}$$

Woah! What do we have here? $P^{-1}=(P^{-})^{-1}+P_z^{-1}$. Precision adds! We found the analogous way to make $K$ reproduces precision adding. That closes the loop. We now have a working Kalman filter on the full dataset :)

I will add experimental results to my dataset at a later date.