Comparing noisy neural population dynamics using optimal transport distances

Authors: Amin Nejatbakhsh, Victor Geadah, Alex H Williams, David Lipshutz
Primary area: applications to neuroscience & cognitive science
Keywords: representational similarity, shape metrics, optimal transport, wasserstein distance

Abstract:
Biological and artificial neural systems form high-dimensional neural representations that underpin their computational capabilities. Methods for quantifying geometric similarity in neural representations have become a popular tool for identifying computational principles that are potentially shared across neural systems. These methods generally assume that neural responses are deterministic and static. However, responses of biological systems, and some artificial systems, are noisy and dynamically unfold over time. Furthermore, these characteristics can have substantial influence on a system’s computational capabilities. Here, we demonstrate that existing metrics can fail to capture key differences between neural systems with noisy dynamic responses. We then propose a metric for comparing the geometry of noisy neural trajectories, which can be derived as an optimal transport distance between Gaussian processes. We use the metric to compare models of neural responses in different regions of the motor system and to compare the dynamics of latent diffusion models for text-to-image synthesis.

Decision: accept (oral)
Reviewer scores: 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1730230788658

thread_type:
official_review

<thread_object>
creator: reviewer_mgzm

<content>
summary:
This paper proposes a new notion of distance between Gaussian processes based on optimal transport between them. The innovation is in considering correlation across time, such that the optimal transport incorporates information across time but respects causality. In particular, this allows the distance to differentiate processes which are indistinguishable when comparing their marginal moments. This distinction is shown to be relevant in a neuroscience-inspired synthetic problem, where distances which compare only marginal moments fail to distinguish small magnitude time segments which have high correlation with other time segments. These distances are also used to compare the output trajectories generated by denoising diffusion image models.

strengths:
* The presentation throughout is very clear and a pleasure to read.
* The motivation of the method through prior work is exceptionally good.
* The method this paper proposes is intuitive but original, fairly simple, well-motivated, and seems to be a substantial extension of previous work.
* The experimental evaluation is appropriate.

weaknesses:
* The experimental evaluation is a bit terse, and doesn't demonstrate a clear advantage for this paper's method over SSD across the board. The method is only evaluated on synthetic neural data even though neuroscience is the clearest application for these ideas. Why not try the it on the actual reaching data which inspired the synthetic dataset?
* I'm unsure how important the problem this method addresses is -- the extent to which it's useful to compute distances between stochastic processes for applications like comparing diffusion models or brains is not clear.

questions:
* I would appreciate some discussion of the derivation of the causal OT distance in the main body. This seems important to justifying how it can be viewed as an OT distance.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730331004704

thread_type:
official_review

<thread_object>
creator: reviewer_sz7a

<content>
summary:
This paper introduces a new metric for comparing noisy neural trajectories. The metric is based on computing optimal transport distances between gaussian processes. While the metric can in principle be used to compare any two stochastic processes, the authors justify the metric rigorously for linear-time varying stochastic processes. To support their theoretical claims, the authors provide several numerical examples where their new metric performs better than two previously proposed metrics (Procrustes and SSD).

strengths:
The paper is _very_ well-written, well-motivated, and timely. The mathematical exposition is especially clear. The numerical experiments are also nice.

weaknesses:
The major weakness of the paper is the lack of numerical comparisons to Dynamical Similarity Analysis (DSA), introduced by Ostrow et al, 2023. This omission is very puzzling, since the authors mention DSA early and often, and even use some of the same tasks as DSA (Fig 2).

questions:
Why not compare to DSA? Even if DSA does "better" on some of the tasks, that does not make your method less valuable. As you point out, DSA is only theoretically justified for deterministic systems. Therefore, _I will be happy to raise my score if these comparisons are included, regardless of comparative performance._

EDIT 11/21--The authors have included the analysis I asked for. I have accordingly raised my score. 

### Minor points:
- L046-L047: The phrasing makes it sound like both studies addressed stochasticity and dynamics.
- L050: You don't actually show this for DSA, as far as I can tell. 
- Equation 3: There is a missing time dependence in $A^\top$.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730799972770

thread_type:
official_review

<thread_object>
creator: reviewer_4ecc

<content>
summary:
This paper presents a novel metric for comparing noisy neural population dynamics based on the notion of causal optimal transport, which respects temporal causality. Other methods in the field generally assume deterministic or static neural activity, missing key aspects of biological and artificial neural systems. They validate their metric in a biological task (synthetic motor control experiment) and artificial task (conditional image generation in diffusion) and find that their metric distinguishes different neural trajectories properly.

strengths:
- The paper is well written, providing a good motivation and background on the limitations of current metrics such as SSD and Procrustes Shape Analysis in distinguishing different types of neural dynamics. The mathematics of optimal transport is well presented and motivated with good explanations of the intuition behind the new metric. 
- The paper provided three sets of experiments: scalar task comparing the different metrics mathematically, synthetic for biological motor control, and latent diffusion models. 
- The discussion on disentangling recurrent dynamics from input-driven dynamics really highlights importance of considering temporal dependencies in the data.

weaknesses:
- As mentioned by the authors themselves, the metric lies on an assumption of Gaussian neural processes, which does not hold for biological and artificial neural systems. It may not apply to non-Gaussian data as it won't be able to capture the higher-order statistical dependencies between the trajectories. 
- Estimating full covariance matrices for high-dimensional neural data (large N and T) is computationally intensive and probably requires a prohibitive number of samples.
- The paper mentions Dynamical Systems Analysis by Ostrow et al. in the end of the paper as a possible future direction. It would be nice if the authors can actually include a comparison to that method in the current paper, as they compare Causal OT to Procrustes and SSD, but DSA seems to be a better method for studying neural dynamics than both of them.

questions:
- I would love to see theoretical or empirical analysis of the metric's robustness to deviations from Gaussianity. I know this was briefly mentioned in the section on Diffusion, but I would like to see a demonstration of the metric's performance as you increase the impact of higher order statistical moments such as skewness or kurtosis. 
- The metric is computational expensive. Can it be reduced using some dimensionality reduction techniques that preserves stochastic and temporal structure (maybe Dynamic Mode Decomposition or Tensor Factorization)? Can you impose structure on the covariance matrices, such as low-rank approximations or sparsity constraints, to reduce the number of parameters to estimate? 
- Can you comment on how this method is affected by the observations in Qian et al. 2024 (Partial observation can induce mechanistic mismatches in data-constrained models of neural dynamics) that points out challenges in identifying mechanisms in neural data (such as line attractors) based on partial observation of neural data.
- I would like to see experiments comparing Causal OT to DSA in some way.
- I would like to see an example applies to actual neural data such as calcium imaging or electrophysiological recordings available publicly and see if the metric can distinguish between neural trajectories coming from different experimental conditions. This might be hard to do within the review period, but it would greatly improve the paper's quality and would motivate me to increase the score substantially.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732215359750

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response (summary of changes)

comment:
Thank you for your careful reviews and thoughtful comments. We’re glad that you found the paper well written and well motivated. We sincerely appreciate your constructive critiques and have revised our paper accordingly. We list the changes we made to the manuscript and address broad issues here, and respond individually to each review below.

# Major additions

1. We now include comparisons to DSA in all of the numerical experiments (Figs. 2, 4, 5).
2. We added Appendix F which compares all of the methods (Procrustes, SSD, DSA and Causal OT) on a scalar autoregressive (AR) process.
3. We added Appendix G which tests the sensitivity of DSA to the following hyperparameters: **number of delays** used to construct the Hankel matrix and **rank** used when solving the reduced-rank regression problem.

# Minor changes

1. We moved the discussion about “recurrent dynamics” versus “feedforward drive” earlier in section 2 because it’s one of the main motivations for developing the metric.
2. We slightly generalized the distance to depend on a parameter $0\le\alpha\le2$, which adjusts the relative weight that is placed on the means versus the covariances. When $\alpha=0$, the distance reduces to Procrustes distance, whereas when $\alpha=2$ the distance only depends on the covariance structure.
3. Notation: we now use ${\bf x}$ and ${\bf y}$ instead of ${\bf x}^A$ and ${\bf x}^B$ to denote two *processes*. This is to avoid confusion with the *matrices* ${\bf A}$ and ${\bf B}$, which were not always related to the processes.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors present a novel method for estimating the similarity of two sets of neural data using causal optimal transport. While there was some concern about the lack of comparison on real experimental data, the authors agreed the method was novel and useful and the paper was well written, and there was a clear consensus that the paper should be accepted. I agree.