Leveraging Generative Models for Unsupervised Alignment of Neural Time Series Data

We thank the reviewer for their feedback and are in the process of updating the manuscript accordingly. We address specific points below.

1. The experiment results are limited... In the updated manuscript, our experiments section includes a comparison to three more methods: Procustes, NoMAD [1], and Cycle-GAN [2]. We have also added an experiment based on the Lorenz attractor, a common baseline in the seqVAE literature. In the Appendix, we have included more supplementary experiments such as aligning across recording modalities. Moreover, we emphasize that although we only have data from one experiment, we display results when applying our alignment method on both smoothed spikes, which is a common pre-processing step, and raw spikes.
2. how do you model the latent dynamics pθ(xt|xt−1)? Do you add any smooth constraints? In figure 2b, it seems that the results of the proposed approach are much smoother than ERDiff. We model our latent dynamics as an MLP with sigmoid activations. We do not apply any smoothness constraint to the latent dynamics. Concerning ERDiff, we note that for all methods we used the same, fixed pre-trained latent dynamics and encoder to avoid confounders. Moreover, ERDiff does not use the latent dynamics for aligning and instead requires training a separate spatio-temporal diffusion model which could be a potential reason for the difference. Thus, the smoothness of our alignment procedure relative to ERDiff is due to our approach, not the dynamics. We will update the manuscript to make this clear.
3. what is FA + LDS method... We have removed the FA + LDS method as we found it was inappropriate for the experiments. As mentioned above, we have introduced three more comparisons instead.
4. what are the results showing in table 1? Table 1 demonstrates performance on decoding the hand velocity of the monkey after alignment. The updated manuscript will make this clearer.
5. in table 1, why the k-step MSE standard deviation... We have updated the table for the monkey reach dataset with new single-step and multi-step metrics for the new comparisons.
6. in figure 4, why the proposed method performs better than ERDiff for poisson observations but not for the smoothed observations, especially Figure 4b (Top, left)? There was a small error in our evaluation that will be fixed in the updated manuscript.
7. The proposed method assumes that prior dynamics p(xt | xt-1) is fixed, and the encoder q(x | y_embed) is fixed. I wonder if the authors have checked these two assumptions? For example, if you jointly train a model on two datasets and learn the latent dynamics and encoder, will they be similar to the results of training separate models on each dataset? While this is an interesting experiment, we found that training jointly one generative model for two datasets is highly non-trivial. This is especially true if the datasets are from different modalities. As a proxy, we have compared the performance of using our alignment method to re-training a seqVAE from scratch on the new dataset. We found that the proposed approach performs on par with training from scratch. Moreover, we see that when data is limited, using our alignment approach leads to improvements compared to re-training from scratch. These results will be available in the updated manuscript.

[1] Stabilizing brain-computer interfaces through alignment of latent dynamics. [2] Using adversarial networks to extend brain computer interface decoding accuracy over time.

We thank the reviewer for their feedback and are in the process of updating the manuscript accordingly. We address specific points below.

1. The approach uses a fairly strong assumption that the latent dynamics really are shared across data sets… We agree that this is a strong assumption as we would expect that the proposed approach would work well on data from animals doing very similar tasks. We will update the manuscript to acknowledge this weakness.
2. Given the large literature on data alignment/domain adaptation both within and without neuroscience, it's a bit surprising that there are only two approaches compared in Table 1 In the updated manuscript, we have included three more methods as comparisons: the Procrustes method, NoMAD [1] and Cycle-GAN [2].
3. How flexible is this setup to the specific architecture chosen? It's mentioned that using an ELBO that measures log predictive probability several steps ahead is important to achive a good embedding, but it's not entirely clear to me why. Our method makes no assumption on the architecture of the seqVAE. The only assumption that we make is that the pre-trained seqVAE achieved good performance on the original dataset. Thus, we would expect our method to perform regardless of the architecture. Regarding the k-step ahead prior, we found that during optimization, the standard ELBO used for seqVAE would sometimes get stuck in sub-optimal local minimum. Through further investigation, we found that the alignment procedure would produce latents that would respect the one-step ahead dynamics, $p_\theta(x_t \mid x_{t-1})$, but not the global dynamics. Thus, the $k$-step ahead prior was introduced as a regularizer to ensure that aligned latents would respect the global dynamics. We will update the manuscript to make this clearer.
4. How complex are the learned dynamics in cases that work versus don't work? Empirically, we found that the method performed well regardless of the complexity. In the updated manuscript, we have included a synthetic experiment based on the Lorenz attractor, a chaotic dynamical system, and found that our proposed method still performs well.

[1] Stabilizing brain-computer interfaces through alignment of latent dynamics. [2] Using adversarial networks to extend brain computer interface decoding accuracy over time.

Thank the author for the points in the response. I will hold my original score.

We thank the reviewer for their feedback and are in the process of updating the manuscript accordingly. We address specific points below.

1. We have updated all of the experiments to include more baselines. Namely, the new baselines are now the Procrustes method, NoMAD [1] and Cycle-GAN [2]. We have also updated the related works section to discuss more methods, including SABLE [3]. We didn’t include SABLE as a comparison as it requires both data and behavior for aligning, making it inapplicable for both synthetic experiments.
2. While there are similarities between the proposed approach and ERDiff—specifically, leveraging the spatio-temporal structure and the re-use of generative models—there are also significant differences between the two approaches. Namely, ERDiff requires training a spatio-temporal diffusion model, along with a seqVAE, on the source dataset to perform alignment on a new neural dataset. This adds overhead as if one wants to share their pre-trained seqVAE they must also train a diffusion model. Moreover, ERDiff does not explicitly use the learned latent dynamics when aligning. Instead, only the encoder is re-used and spatio-temporal transformer blocks are used to take into account the spatio-temporal nature of the data. In contrast, our approach is considerably simpler as it only requires training a lightweight alignment function that feeds into the encoder and does not need access to the source dataset. Moreover, the use of the pre-trained latent dynamics is explicitly used as opposed to ERDiff. We provide empirical evidence that—despite its simplicity—our method, which directly uses the pre-trained latent dynamics rather than a spatio-temporal diffusion model, outperforms ERDiff along with the other baselines.
3. We have added several experiments to the manuscript, demonstrating the success of the proposed approach. We have included an additional synthetic experiment testing all methods on the Lorenz system, a common benchmark for seqVAEs. We have included experiments demonstrating alignment across modalities (real-valued to spikes and spikes to real-valued).

[1] Stabilizing brain-computer interfaces through alignment of latent dynamics. [2] Using adversarial networks to extend brain computer interface decoding accuracy over time. [3] Robust alignment of cross-session recordings of neural population activity.

We thank the reviewer for their feedback and are in the process of updating the manuscript accordingly. We address specific points below.

1. "I am not sure I follow why Proposition 1 implies good alignment?..." There was an error in Proposition 1 that will be fixed in the updated manuscript. Namely, the alignment found by the proposed approach, $\vartheta_\star$, can be expressed as a linear transformation of the optimal alignment, $\vartheta_\dagger$, i.e, $\vartheta_\star = \left( I + Q (A A^\top)^{-1} \right) \vartheta_\dagger ,$ where $I$ is the identity matrix. From this, we can see that the difference between $\vartheta_\star$ and $\vartheta_\dagger$ is a function of $Q (A A^\top)^{-1}$. Thus for observations that have small noise $Q \approx 0$ and/or small $(A A^\top)^{-1}$ we would expect $\vartheta_\star$ to be similar to $\vartheta_\dagger$. Note, that we can loosely interpret $Q (A A^\top)^{-1}$ as the inverse of the signal-noise-ratio of $w$. Thus, the higher the SNR of $w$ is, the closer $\vartheta_\star$ will be to $\vartheta_\dagger$. The updated manuscript will contain the revised Proposition along with a revised proof as well.
2. "The choice of alignment function appears crucial to this paper,..." We first note that for experiments that operate on spikes, we parameterize $g_\vartheta$ using a multi-layer perceptron; we will update the manuscript to make this clear. Concerning the expressivity of $g$, we agree that if we allow $g$ to become very expressive then we could effectively re-learn the encoder. Thus, following standard practice in using pre-trained networks, we make $g$ relatively small compared to the generative model. In our experiments, both synthetic and monkey-reach data, we found that our method performed well even when we parameterized $g$ with an MLP. We will update the manuscript to discuss this point as well. We have also included an experiment comparing how the method performs if we set $g$ to be linear as opposed to an MLP.
3. "I assume $q_\phi(x_t, x_{t-1} \mid y_{1:T}) = q_\phi(x_t \mid y_{1:T}) q_\phi(x_{t-1} \mid y_{1:T})$..." Yes, that is correct! We will clarify this in the manuscript.
4. Is the equation supposed to have $p_\theta(x_t \mid x_{t-k})$ instead of $p_\theta(x_t \mid x_{t-1})$?. Yes, that is correct! $p_\theta(x_t \mid x_{t-k})$ is the k-step ahead prior that can be expressed as applying the dynamics k times in succession, i.e. $p_\theta(x_t \mid x_{t-k}) = \int p_\theta(x_{t - k + 1} \mid x_{t - k}) p_\theta(x_{t - k + 2} \mid x_{t - k + 1}) \ldots p(x_t \mid x_{t - 1}) dx_{t - k + 1} \ldots dx_{t - 1}$. We will update the manuscript to make this more clear.