Extraction and Recovery of Spatio-Temporal Structure in Latent Dynamics Alignment with Diffusion Models

Summary

The authors address the problem of aligning behavior-related neural population dynamics, either within-subject but across different experimental sessions, or between subjects. This problem exists due to inter-subject variability in terms of which neurons are recorded or drift in the recording and is an important problem for systems neuroscience and for the development of brain-machine interfaces. Importantly, neural activity dynamics in many regions of the brain during behavior lie on a low-dimensional manifold and therefore have a defined spatio-temporal structure. While many state-of-the-art alignment approaches do not take into account this structure and therefore do not preserve it, the authors introduce a diffusion-guided method that does. The approach first uses a diffusional model to discover the manifold on which neural activity evolves (ie the spatiotemporal structure), then it uses this model to guide the alignment, which is done using MLE.

Strengths

Originality + quality: The authors developed a novel approach for time-series alignment that also discovers latent structure (the low-D manifold on which the neural activity evolves). Quality: The authors validate their model against a number of state-of-the-art alignment methods on both synthetic and real-world data, demonstrating the practical applicability of their method. Clarity: The authors clearly state the problem and its details, as well as how their method differs from existing ones, and its advantages. Significance: The method has significance to the alignment of timeseries with latent spatiotemporal structure, which is broadly relevant in systems neuroscience. Although the authors mainly focus on behaviorally relevant neural data, timeseries in other fields also often possess lower-dimensional latent structure, so this method could be broadly applicable. The authors approach could also be relevant to identifying latent structure outside of alignment context, although I am not familiar with how it compares to existing approaches to do so.

Weaknesses

Clarity: Figure and table legends should be more clear (see Questions section) Originality: Authors should state whether existing methods for discovering latent structures using diffusional models exist, and how their methods (e.g. architecture) differs from existing methods. Quality: The method is validated in data with strong, low-dimensional latent structure (monkey reaching tasks). The generalizability to other types of neural dynamics time series of varying dimensionality should be evaluated to determine limitations.

Questions

Figure 3B: Can the authors describe better what the takeaway from this panel is? In what ways does ERDiff improve on JSDM. The legend should tell the reader what the pink dots are (presumably critical points?). Figure 4: Please note in the legend or figure that the R-squared value is in %. I was confused at first as to how the R-squared could be negative before reading Table 1. Table 1: It is unclear to me what the denominator for the R-squared % calculation is. Can the authors make it more apparent in the table legend? If the time-series do not belong on a well-defined low-dimensional manifold, will the method hallucinate something? How well does the method work for more unstructured data? As the dimensionality of the latent structure increases, how much more badly does the method perform? The authors should quantify this.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The authors should address in more detail the kinds of latent structures/time series that would pose more of an issue for their method. Otherwise, limitations are adequately addressed.

Summary

Brain-computer interfaces require recalibration to accommodate drifts in the recorded neural populations over time. While there has been some success in aligning neural recordings based on their latent dynamics, deep learning-based alignment methods have also gained popularity since they ignore some of the implicit assumptions made by latent space methods, providing additional flexibility. However, deep learning neural alignment often ignores the temporal structure of the dynamics. The proposed model overcomes these limitations by first extracting the spatio-temporal structure in the source domain via a diffusion model and then aligning the target domain to the source dynamics. The authors demonstrated the success of the method in simulated and neural data, where it outperforms alternative alignment methods.

Strengths

The paper is adequately written and technically sound. The method was tested and shown to work well when applied to a simulated dataset and neural data. The method introduced here uses deep learning-based alignment while still exploiting the temporal dynamics that are critical in neural datasets. Robust alignment of neural recordings is crucial for the success of BCI applications, and in this work, they showed how this method outperforms existing alignment methods in both synthetic and neural datasets. Moreover, they also showed that the alignment can be performed not only across sessions of the same animal but also across animals.

Weaknesses

While the authors show the promise of their method to align neural responses, they overlooked other methods based on the alignment of latent dynamics that have been shown to be successful for BCI applications, as mentioned in the introduction. I believe that a systematic comparison to such methods is critical to fully grasp the significance of this work. For example, CCA, multiset CCA, hyperalignment, or Procrustes-based alignment. In the context of latent space alignment methods, another relevant piece of literature is the method introduced in (https://proceedings.neurips.cc/paper/2021/hash/aad64398a969ec3186800d412fa7ab31-Abstract.html), which also uses neural dynamics for alignment. Deep learning-based methods allow for more expressive functions, but they often come with additional computational costs, data demands, and the need for careful hyperparameter selection. None of these limitations are addressed in the manuscript, nor is there an explicit comparison across methods (latent space vs. deep learning-based), which could help demonstrate the promise of the method for practical BCI applications. The authors minimally showed the effect of dataset size on performance, but the lowest dataset size tested still had dozens of trials, which could be unrealistic in most practical settings. Additionally, it would be important to report the training times as a function of dataset size, as long training times would render the approach useless for real-time alignment. The method defines the alignment between a single source and target dataset, but ideally, one would pool data across all sessions for BCI decoding. It would be interesting to note if the proposed method also allows for multi-session alignment. The authors showed the success of the approach in the context of a single data simulation. However, to fully assess the robustness of the method, they could further test it under different conditions, such as measurement or latent noise, dimensionality, or tasks.

Questions

It is unclear from the text how the training and testing are performed. In the sentence "During testing, we align the test neural data to the training neural data so that we can directly apply the velocity decoder," it is not clear whether the authors include test data for alignment. Additionally, the authors mentioned that behavioral data is used for alignment, but it is also used to evaluate the success of the approach via decoding. This raises the question of whether there is a circular evaluation of the method. References 9 and 10 cite the preprint and peer-reviewed versions of the same article.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

The authors should include a section or provide a clear description of the limitations and assumptions of the method. They should also address the computational cost, data demands, and potential implications of the presented work.

Summary

* One of the key challenges in analyzing neural recordings is the scalability of models that link behavior and neural population activity across recording sessions or in inter-subject settings. • When analyzing single-trial neural population activity, past studies have pointed out that neural activity can be understood in terms of low-dimensional latent dynamics. Such low-dimensional latent dynamics are helpful when visualizing neural profiles across different task conditions and trials. • Generally, existing methods try to align latent dynamics by minimizing the difference evaluated by the metric between source and target domains. The paper proposes a method to align the source and target domains of multivariate neural data by learning/capturing latent spatiotemporal structure in the source domain with a diffusion model and applying it as a prior on learning/capturing spatiotemporal structure in the target domain. • The authors applied their model to the non-human primate motor cortex, testing both cross-day and inter-subject recordings.

Strengths

* The authors motivate their approach clearly by arguing that naively aligning time series using domain adaptation is ineffective as multivariate neural data has low SNR. Thus, leveraging low-dimension representation is a viable option. • The paper seeks to achieve a form of domain adaptation by aligning the spatiotemporal structure of latent dynamics of the target to the source using a novel alignment method (ERDiff). • The model and derivations are presented clearly. • An exhaustive comparison is provided showing that the their model outperforms standard models on both motor cortex and synthetic datasets.

Weaknesses

* The authors must clarify why diffusion models are necessary. How about considerably simpler two-step approaches -- like extracting latents with GPFA (Yu et al, 2009) and aligning them with the proposed ML alignment? * Alternatively how about comparisons with alternative approaches of comparable complexity -- e.g. adversarial alignment with DANN (Ganin et al, 2015)?

Questions

* Overall: there are many approaches for extracting latent structure from time series data (GPFA, CEBRA, T-PHATE, CILDS) -- one could readily realign the latents extracted from these approaches with the second stage alignment algorithm. * Apriori, why would one expect the diffusion model to be more effective at aligning the latents than these other approaches? * Because data from non-human primates is used, please clarify whether appropriate IRB approvals were obtained. Or if this is only a secondary analysis of existing datasets, the approvals obtained in the original studies could be mentioned.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

* No significant negative societal impact envisioned -- the potential use in BCI may suggest a positive social impact.

Summary

This paper proposes a distribution alignment method (ERDiff), which combines extraction of spatio-temporal structures in latent dynamics from the source distribution and maximum likelihood alignment procedure on the target domain. The proposed method was evaluated on both synthetic and real data (neural recordings from non-human primate motor cortex), outperforms other methods under both cross-day and inter-subject settings.

Strengths

- The proposed method is novel and is technically sold - Performance of the method is well demonstrated on real data under inter-session/subject setup suggest that it could be an important tool with potential broad use in many field not just in neuroscience.

Weaknesses

-I don’t have any major concerns. Although the methods has been shown to outperform some of the current techniques, advantage of the proposed approach is not well demonstrated. I would like to see some analysis on computational cost

Questions

- Although the authors clearly mentioned the limitation of alignment method based on pre-defined metric, it would be nice see how the proposed method performs compared to these.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

- Limitations is not addressed in the draft

Summary

The paper focuses on aligning highly variable neural population activities across days and subjects to stabilize the learning proposes and advance applications such as the brain computer interface. The idea is to train a VAE on one dataset that is hopefully self-consistent and then using a VAE trained on a dataset from another day or subject align the conditional distributions (encoders) for the latent spaces maximizing the likelihood of the source domain latent space under the new encoder distribution. The difficulty is the need to model the marginal source distribution of the latent space, which the work does using a diffusion model. The approach is demonstrated in comparison with alternative models on a synthetic dataset and actual neural recordings.

Strengths

1. A well written paper (but the abstract). 2. An interesting application of diffusion models. 3. Potentially impactful in practice of the BCI, more work, including further evaluation, is needed here however.

Weaknesses

1. A niche application and demonstration. A cellular neuroscience focused paper with no additional effort to demonstrate a general applicability of the approach in evaluations. 2. Poorly written abstract, especially in contrast to the rest of the paper.

Questions

1. Is the code going to be released publicly? Looks like the success of the approach depends less on the high level probabilistic descriptions in the paper than on the details of the actual implementation. 2. Is the data going to be released publicly for reproducibility?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

Potentially, the applicability of this work may be limited only to the demonstrated use in intracranial multicellular recordings, and it may not contribute to advancements in other areas of Machine Learning (ML). No demonstration was provided to counter this potential limitation.

Summary

Inter-individual and inter-session variability significantly complicate direct comparison of neural recordings collected over time, degrading trained behavioral models. This can be cast as a more general distribution alignment problem, shared across unsupervised learning. To address neural distribution alignment, the authors introduce a novel "ERDiff" method, which co-trains a variational autoencoder and a diffusion model to extract latent spatiotemporal structure in a source dataset. To align with a desired target dataset, parameter finetuning is performed on the read-in layer of the probabilistic encoder learned during training to match the target dataset. Simulation and experimental results suggest that ERDiff captures relevant spatiotemporal structure, performing competitively against other baseline methods including those based on overall minimizing distribution divergence as well as based on adversarial methods. Overall, the authors argue that ERDiff is better able to capture the important spatiotemporal structure of their trialwise data; for example, from a monkey center-out reaching task in the experimental results. This focus differentiates ERDiff from many other alignment methods, which ignore unfolding dynamics in their alignments.

Strengths

Incorporating spatiotemporal structure into the alignment of neural recordings is a novel approach, as these methods traditionally consider successive data points as independent samples or learn a set of low-dimensional latent dynamics which can then be aligned. These dynamics are particularly critical where the the trialwise dynamics strongly influence both behavior and neural activity over time. By directly learning and aligning the latent spatiotemporal structure, ERDiff shows stable performance even over relatively low sampling density, retaining relatively high decoding performance compared to other baseline methods even with decreasing numbers of trials in the target domain. These properties suggest that it is also likely that the general ERDiff approach may be useful in other cases where distribution shift between a source and target domain obscures but does not remove a shared latent structure. The general approach of combining variational autoencoders (VAEs) and diffusion models (DMs) has been previously introduced (e.g., Panday et al., TMLR, 2023); however, ERDiff is a significantly different formulation of the idea and represents a novel approach in leveraging the relative strengths of these methods through cooperative training.

Weaknesses

While the current experiments extensively compare inter-session and inter-subject differences in real neural recordings --- in addition to the simulated experiments --- it is not clear to what extent the presented findings might generalize to data sources without such a clear trial structure. For example, in recordings collected during unconstrained exploration or sleep, low-dimensional latent factors may still drive a successful alignment. Nonetheless, it is not clear how ERDiff would best be deployed in that context. This is particularly relevant as the described real-data experiments additionally incorporated velocity information, and the performance of ERDiff without a behavioral signal during training is thus unclear. In the current paper, my additional concern with the experiments is on the relative baselines against which ERDiff is compared. It would be informative to see a direct comparison with canonical correlation analysis (CCA) approaches, which have been used to date in aligning neural datasets with a strong temporal correspondence (e.g., Gallego et al., Nat Neuro, 2020).

Questions

Would the authors be able to directly comment on the relationship between their work and other field-standard methods to identify low-dimensional, latent factors such as LFADS (Pandarinath et al., Nat Methods, 2018) ? In particular, the relative benefits of learning the low-dimensional latent structure as part of the alignment --- as compared to other existing methods which learn low-dimensional dynamics which can then be aligned --- is not clearly explained. This would help to better situate the work in the literature.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Although I do not see a Broader Impact section included in the current submission, I do not anticipate potential negative societal impacts given the constrained focus of the work.