Paper: Leveraging Generative Models for Unsupervised Alignment of Neural Time Series Data (ICLR-2024)

Reviewer comment (Experiments):
The experiment results are limited. The paper only shows results on 1 real dataset.

Author response:
`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.

Reviewer comment (Reproducibility):
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.

Author response:
`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.

Reviewer comment (Reproducibility):
What is FA + LDS method that you compare to in figure 2? Can you list a bit more details on how it's implemented?

Author response:
`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.

Reviewer comment (Evaluation):
What are the results showing in table 1? Are they results for the real neural dataset? I found the r^2 number in table 1 doesn't match the number in figure 4.

Author response:
`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.

Reviewer comment (Evaluation):
In table 1, why is the k-step MSE standard deviation smaller than the stddev of MSE for ERDiff and the proposed method?

Author response:
`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.

Reviewer comment (Evaluation):
In figure 4, why does the proposed method perform better than ERDiff for poisson observations but not for the smoothed observations, especially Figure 4b (Top, left)?

Author response:
`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.

Reviewer comment (Experiments):
The proposed method assumes that prior dynamics p(xt | xt-1) is fixed, and the encoder q(x | y_embed) is fixed. Have the authors 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?

Author response:
`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.

Reviewer comment (Theory):
The approach uses a fairly strong assumption that the latent dynamics really are shared across data sets, which all but implies a shared task setup. That is, it doesn't appear to be the case that a sufficiently large task-free data set in one mouse would facilitate embedding of mice performing a task-based behavior. It would be surprising if true, but this weakness should be acknowledged, since this limits the range of applicability.

Author response:
`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.

Reviewer comment (Evaluation):
How complex are the learned dynamics in cases that work versus don't work? The monkey reach data typically have rotational dynamics. Do you see anything more complicated in other data?

Author response:
`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.

Reviewer comment (Experiments):
The choice of alignment function $g_\theta$ appears crucial to this paper, but there is very little discussion/evaluation about different choices. Specifically, if $g_\theta$ is a point-wise nonlinear function approximator $g_\theta(w_t)$ the authors are implicitly assuming that the latent dynamics are not just qualitatively similar, but **identifcal** between the new and old observations. However, if we make $g_\theta$ too flexible (e.g., another full bi-directional RNN), then we are effectively re-learning the encoder.

Author response:
`"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.

Reviewer comment (Writing):
I assume $q_\phi(x_t, x_{t-1}|y_{1:T}) = q_\phi(x_t|y_{1:T}) q_\phi(x_{t-1}|y_{1:T})$ in equation 3? Might be worth clarifying.

Author response:
`"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.

Reviewer comment (Theory):
Is the equation supposed to have $p_\theta(x_t|x_{t-1:t-k})$ instead of $p_\theta(x_t|x_{t-1})$?

Author response:
`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.

Reviewer comment (Experiments):
The baselines in the experimental part is too few, just the recently proposed SOTA method ERDiff [1]. There are many classical methods like [2] and [3].

Author response:
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.

Reviewer comment (Novelty):
The spatio and temporal structure has already been noticed by the SOTA method ERDiff, and ERDiff also employs a generative model (score-based model) for alignment. Thus what's the new motivations and insights of your work?

Author response:
We thank the reviewer for their feedback and are in the process of updating the manuscript accordingly. We address specific points below.
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.