Comparing noisy neural population dynamics using optimal transport distances
Authors: Amin Nejatbakhsh
Victor Geadah
Alex H Williams
David Lipshutz
Venue: ICLR 2025 Oral (2025)
Keywords: Representational similarity
shape metrics
optimal transport
Wasserstein distance
Primary area: applications to neuroscience & cognitive science
TL;DR: We propose using optimal transport distances on stochastic processes to compare noisy neural trajectories.

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.

Published as a conference paper at ICLR 2025
COMPARING NOISY NEURAL POPULATION DYNAMICS USING OPTIMAL TRANSPORT DISTANCES
Amin Nejatbakhsh' Victor Geadah 1:2 Alex H. Williams" 1.3
& David Lipshutz"
Center for Computation Neuroscience, Flatiron Institute
2 Applied and Computational Mathematics, Princeton University
3 Center for Neural Science, New York University Department of Neuroscience, Baylor College of Medicine
anejatbakhshee flat roninstiiinn coninstitute.or.
ABSTRACT
Biological and artificial neural systems form high-dimensional neural represen- tations that underpin their computational capabilities. Methods for quantifying geometric similarity in neural representations have become popular tool for iden- tifying computational principles that 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 system's computational capabilities. Here, we demon- strate that existing metries can fail to capture key differences between neural sys- tems with noisy dynamic responses. We then propose a metric for comparing the geometry of noisy neural trajectories, which can be derived as an optimal trans- port 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 latent diffusion models for text-to-image synthesis.
INTRODUCTION
Biological and artificial neural systems represent their environments and internal states as patterns of neural activity, or "neural representations". In an effort to understand general principals governing these representations, large body of work has sought to quantify the extent to which they are simi- lar across systems (Klabunde et al., 2023). Guided by the intuition that the geometry of a systems' representation may be related to its function (Mikolov et al., 2013; Harvey al., 2024), several measures of geometric (dis)similarity have been proposed and applied to data, including Represen- tational Similarity Analysis (RSA, Kriegeskorte al., 2008), Centered Kernel Alignment (CKA, Kornblith et al., 2019), canonical correlations analysis (CCA, Raghu et al., 2017), and Procrustes shape distance (Williams al., 2021; Ding et al., 2021).
Most of these existing (dis)similarity measures assume that neural activities are deterministic and static. However, in many case, one or more of these assumptions does not hold. For example, the responses of sensory neurons to the same stimulus vary across repeated measurements (Goris al., 2014), and the structure of these variations is thought to be critical to system's 's perceptual capabilities (Averbeck et al., 2006). In addition, neural activities in brain regions such as motor cortex evolve according to complex dynamical motifs that correspond to current and future actions (Vyas et al., 2020). Similarly, responses of artificial neural systems can be noisy (c.g., variational autoencoders) or dynamic (c.g., recurrent neural networks) or both (c.g., diffusion- based generative processes).
Quantifying similarity in either the stochastic or dynamic aspects of neural systems has been in- dependently addressed in the literature. Duong et al. (2023) proposed Stochastic Shape Distance (SSD). a stochastic extension to Procrustes shape distance (Williams et al., 2021) for quantifying differences in trial-to-trial noise across networks without addressing recurrent dynamics. Ostrow

---

Published as a conference paper at ICLR 2025
A Schematic of noity neural dynamics
B Procrustes Shape Distance
C Stochastie Shape Distance (SSD)
D Consal Optimal Transpoort (OI)
Pegelation
Population B
compare aligned meons
compare alignsed murginals
compare aligned trafectories Syy
xy
xyy
Metrics for comparing the shapes of noisy neural processes. A) Spike trains for two populations of neurons with noisy dynamic responses. Each trial can be represented trajectory in neural (firing rate) state space. (In the main text we account for cases that the populations have different numbers of neurons and trials.) B) Procrustes shape distance compares the aligned mean trajectories (thick curves represent the mean neural responses during a trial). C) Stochastic Shape Distance (SSD) compares the aligned marginal statistics. The faint ellipses represent the aligned marginal distributions at different times points. D) Our proposed metric (Causal OT distance) di- rectly compares entire trajectory statistics, thus capturing across-time statistical dependencies.
Figure 1:
et al. (2023) utilized Koopman operator theory to develop Dynamical Similarity Analysis (DSA) to compare recurrent flow fields without directly considering the effects of noise. However, on their own, these methods can fail to capture differences between noisy neural dynamics.
Here, we propose novel metric between noisy and dynamic neural systems that captures differences in the systems' noisy neural trajectories (Fig. IA). The metric is naturally viewed as an extension of Procrustes shape distance which compares average trajectories (Fig. IB) and SSD which compares marginal statistics or noise correlations (Fig. IC). Conceptually, our metric compares the statistics of entire trajectories (Fig. ID). In particular, the metric can be derived as (bi- (bi-Jaausal optimal transport (OT) distance (Lassalle, 2018; Backhoff et al., 2017) in which two systems are compared by computing the cost "transporting" a set of trajectories generated by one system to match set of trajectories generated by another system, and the mapping between trajectories must satisfy temporal causality condition. The causality condition is especially relevant when comparing neural systems whose trajectories are close space and yet the trajectories vary significantly in the amount of information that available at different time points (see Sec. 4.1 for simple illustrative example). We apply our method to compare simplified models of motor systems (Sec. 4.2) and the dynamics of conditional latent diffusion models (Sec. 4.3).
In summary, main contributions are:
We demonstrate how existing metrics for comparing neural responses can fail to capture differences between neural population dynamics (Sec. 2).
We propose a metric for comparing noisy neural population dynamics that is motivated by causal optimal transport distances between Gaussian processes (Sec. 3 and Appx. A).
We provide an alternating minimization algorithm for computing the distance between two processes using their first- and second-order statisties (Appx. B).
We apply our method to compare models of neural responses in different regions of the motor system and to compare the dynamics of conditional latent diffusion models for text- to-image synthesis (Sec.4 4).
SET UP AND EXISTING METHODS
Consider a neural system that consists of neurons whose activities (e.g., firing rates) at time E (1,..., 7) are represented by an N-dimensional random vector x(t). The sequence x (x(1)) defines a random trajectory, or stochastic process, in N (Fig. IA). Given two or more such neural systems, we'd ''d like to develop meaningful measures of their (dis)similarity.
Source code: https: / //gi thub. com/ ami neiiinneejje n-nejat/netrep.
2

---

Published as a conference paper at ICLR 2025
2.1 PROCRUSTES SHAPE DISTANCE FOR COMPARING MEAN RESPONSES
Perhaps the simplest way to compare two dynamical neural systems living in different coordinate systems is to compute the Procrustes shape distance on their mean trajectories (Fig. IB). For pro- cesses x and y with mean responses m²(() == E[x(()] and my == E[y()), the (squared) Pro- crustes shape distance is:
&procmstes (x, == min
mz (t) (t,
QEO(N)
where O(N) denotes the set of orthogonal matrices. This metric conceptualizes each process as a deterministic trajectory through neural state space and equates two processes if their mean responses are equal up to a rotation and/or reflection of neural state space. A useful property of equation is that it defines mathematical metric on neural processes. In particular, the function dpmncostes is symmetric and satisfies the triangle inequality, which is helpful to establish theoretical guarantees for many statistical analyses such as neighborhood-basse clustering and regression (Cover & Hart, 1967; Dasgupta & Long, 2005).
A common goal in analyzing neural systems is to distinguish between the impacts of recurrent inter- actions and feedforward input drive (Sauerbrei et al., 2020), so a useful measure should distinguish systems with different recurrent interactions. Galgali al. (2023, Figure Ie) shows a simple but illuminating example for why this is challenging. They construct three different flow fields and ad- versarially tune tine-vayyigg inputs for each system so that the trial-average trajectories match. By construction, distances based on comparing mean trajectories (e.g., Procrustes shape distance) will fail to capture differences between the three systems.
It is easy, even in linear systems, to construct examples where the mean trajectories are insufficient to distinguish between candidate dynamical models. For example, consider the case that the process satisfies a linear stochastic dynamical system of the form
x(t) (tt)x(t 1) + b(t) - E(t)
where A() is an N x N dynamics matrix, b(l) is the N-dimensional input-drive, E() is an N x N matrix that determines the noise structure, and w(t) is an independent N-dimensional Gaussian random vector with zero mean and identity covariance matrix. Taking expectations on either side, we see that its average trajectory evolves according the deterministic dynamical system
m. (t) = A() (t 1) + b(t).
From this dynamics equation, we see that the evolution of the mean trajectory depends on the net contributions from the (mean) recurrent interactions A(1) (1 1) and the input-drive b(l). Con- sequently, it is impossible to distinguish between recurrent dynamics and input-driven dynamics of a linear system of the form in equation only using the mean trajectory. Galgali et al. (2023) showed that such differences sometimes be accounted for by considering stochastic fluctuations about the mean trajectories, which suggests using metrics that explicitly compare stochastic neural responses.
2.2 STOCHASTIC SHAPE DISTANCE FOR COMPARING NOISY RESPONSES
Duong et al. (2023) introduced generalizations of Procrustes shape distance to account for stochas- tic neural responses. While their method was originally developed for networks with noisy static responses in different conditions, it can be directly applied to compare noisy dynamic responses by treating each time point as a different condition. Here we focus on SSDs that compare the first and second-order marginal statistics at each time point. Specifically, define the marginal covariances of a process x with mean trajectory my by
(a( = E (x(() (I) (x(1) m. m.,(()) (I)'
t= 1, I......
(4)
The a-Stochastic Shape Distance (a-SSD) between processes x and y (Fig. 1C)
dz.ss((xy))= == min
Σ (2 m) Qm,(t))² aB² (P.(t). (P)(().QPP((Q)))) QP)()) (5)
QEO(N)
1=1
3

---

Published as a conference paper at ICLR 2025
Saodie
LineAitracer
ProrAmmcter
24
Saiti
0.00
on
Polet
Stadi:
8
-015
0.0
-1.5
Pole
Saadil
Live
Point
Demonstration that causal optimal transport (OT) distance disambiguates recurrent flow fields while distances based on the mean trajectories and the marginal noise correlations do not. On the left are simulations of three dynamical systems-a nonlinear dynamical system with two stable fixed points, a linear dynamical system with a line attractor and a linear dynamical system with a point attractor-with identical mean trajectories (top) and noise correlations (bottom) due adversarially chosen input drives and input noise. On the right are the pairwise distances between the three systems.
Figure 2:
where O(N) is the set of orthogonal matrices, B(,,) denotes the Bures metric between positive semidefinite matrices:
B(A, min UEO(N)
1/2/22 B1/UU
and 0 S a determines the relative weight placed on differences between the mean trajectories or the marginal covariances, When a = 0, the distance reduces to the Procrustes shape distance between the mean trajectories m, and my. so a-SSD can be viewed as generalization that ac- counts for second-order fluctuations of the marginal responses. Furthermore, if a = =1 and the marginal distributions are Gaussian, i.e., x()) (mm,( Pa (1)) and y() ~ (1), for all / ..... T, then equation 5 is the sum of squared Wasserstein distances (with p 2) be- tween these marginal distributions (Peyré Cuturi, 2019). This interpretation of SSD as an optimal transport distance between Gaussian distributions is relevant to the metric develop in section to compare full dynamical systems.
Lipshutz al. (2024) showed that a-SSD with (1 can succesfully differentiate systems with the same mean trajectories. To understand this, consider the special case that x satisfies equation 2. Then its marginal covariance matrix evolves according to the dynamics
(t) (t) 1) Ett)
From this equation, the input-drive b(l). so it can be used to compare the recurrent dynamics between two systems that share a common noise structure; that is, when Pa (0) Py (0) and the input noise structure E(-) is the same in both networks.
that the evolution of the covariance matrix P(L) does depend
On the other hand, equation also suggests that the marginal covariances alone cannot distinguish between contributions from the recurrent dynamics A(t) and contributions from the noise covari- ance Et)). We show this Fig. 2 where we extend the example of Galgali et (2023) by ad- versarially tuning the input noise (() in addition to the input drive b(l) to achieve systems with different underlying recurrent dynamics, but the same mean and marginal covariance trajectories (see Appx. D for details). This demonstrates that only comparing marginal distributions is insuffi- cient for comparing Gaussian processes the form in equation 2. In particular, these metrics do not account for across-time statistical dependencies.
2.3 DYNAMICAL SIMILARITY ANALYSIS FOR COMPARING RECURRENT FLOW FIELDS
Ostrow et al. (2023) proposed DSA for comparing recurrent neural dynamics. Their method is based Hankel Altemative View of Koopman (HAVOK) analysis (Brunton al., 2017), which is
4

---

Published as a conference paper at ICLR 2025
data-driven method for obtaining linear representations of strongly nonlinear systems that lever- ages Koopman operator theory (Koopman, 1931) and implicitly accounts for across-time statisti- cal dependencies. Essentially, DSA involves constructing Hankel matrix by delay-embedding system's 's observed trajectories and then fitting a reduced-rank regression model to obtain a linear estimate (i.e., fixed dynamics matrix A) for the evolution of the singular vectors of the Hankel ma- trix. Even though the approximation is linear, it can capture global nonlinear dynamics provided the delay is sufficiently long (Takens, 1981). However, this can also lead to difficulties in practice since the delay needs to be appropriately tuned to capture the nonlinear dynamics. The same procedure can be applied another system and the linear approximations of the two systems' dynamies can be compared. Specifically, given two such dynamics matrices A, and Ay. the distance between and y is computed as
ADSA (D.A(x,)) (x. == min
|
QA,Q
QEO(N)
where the similarity transform A, + QA the vector-field analogue the alignment step in Procrustes shape distance.
Ostrow et al. (2023) demonstrated that DSA can disambiguate dynamical systems with the same mean trajectories but different underlying recurrent dynamics. However, DSA is not explicitly de- signed to account for stochastic dynamics and does not distinguish systems with varying noise levels (Appx. F). In addition, DSA depends on hyperparameters: the number of delays used when structing the Hankel matrix and the choice of rank in the reduced-rank regression. When these hyperparameters are optimally chosen, DSA successfully differentiates the dynamical systems in the extended example with adversarially tuned input noise (Fig. 2); however, DSA can fail when the hyperparameters are not optimally chosen (Appx. G).²
CAUSAL OPTIMAL TRANSPORT DISTANCES
We define a metric on noisy dynamic neural processes that accounts for both stochastic and dynamic aspects of neural responses. The metric motivated by Causal OT distances on stochastic processes (Lassalle, 2018; Backhoff et al., 2017) and can conceptually be interpreted as the minimal -- for transporting a random set of trajectories generated by process x to random set of trajectories generated by process where the mapping (or coupling) between trajectories must satisfy a causal- ity (or adapted) property that can be parsed as "the past of x is independent of the future of y given the past of y" and vice versa. More specifically, the metric between coupled Gaussian processes & and ý that are driven by the same white noise process w moments respectively match the first moments of and y. The coupling
be derived the minimal cost
and whose first is flexible in the sense that at each time f, the white noise vector w(t) driving yt)) can be spatially rotated by an orthogonal matrix R, E O(N) because this does not change its distribution and preserves the temporal ordering of the white noise process. A detailed derivation of the distance is provided in Appx. A.
3.1 DEFINITION OF CAUSAL OPTIMAL TRANSPORT DISTANCE
Given a N N-dimensional process x = {x(t)} defined for t = I,..., ,, let C2 be the NT x NT covariance matrix following T x T block structure:
(1,1)) C₂ : (T
P2(1T))
m. (s))(x((1) m. (1))
(T,)
where the blocks (s, are x N matrices encoding the across-time covariances. The a- a-Caaaal OT distance between two processes x and y is defined to be
dov-oussal (x, y) ==
T
(8) (2-a)) =0 + (Ir (1
min QEO(N)
We implemented DSA using the code provided in the GitHub respository httpa://github..or mi tceee rowtro.
5

---

Published as a conference paper at ICLR 2025
where, with a-SSD, 0 f I 2 determines the relative weight place differences between the mean trajectories or the covariances, AB,,T(+..) is the adapted Bures distance on NT x NT positive semidefinite matrices defined by
min ABN,T N [A, B) = Ry RTE(N)
[L Lædiag( Rr)ll:
(9)
and A LALA A and B L B are the lower triangular Cholesky decompositions of NT NT positive semidefinite matrices A and B. Note that the Cholesky decomposition is unique for positive definite matrices. For matrix with rank R < N, we take L to be the unique lower-triangular matrix with exactly R positive diagonal elements and N RR columns containing all zeros (Gentle, 2012).
3.2 PROPERTIES AND COMPARISON WITH EXISTING METRICS
Causal OT distance is naturally interpreted as an extension of Procrustes shape distance and SSD. When the across-time correlations are zero for both x and y, then C2 and Cy are block-diagonal and a-Causal OT distance coincides with a-SSD defined in equation 5. Therefore, Causal OT dis- tance can be viewed as an extension of SSD that accounts for across-time correlations of stochastic processes. When applied to the dynamical systems with adversarially tuned input drives and input noise, we find that Causal OT can disambiguate the different dynamical systems that Procrustes shape distance and SSD cannot differentiate (Fig. 2).
also worth comparing Causal OT distance with Wasserstein (with p 2) distance between Gaussian processes (Mallasto & Feragen, 2017), which treats the processes Gaussian vectors without any notion of temporal ordering. The difference between the Wasserstein distance between two Gaussian process and equation 8 is that the adapted Bures distance AB,.T(; is replaced with the Bures distance
.(.,)) between the covariance matrices:
((A, B)
min UEO(NT)
|///22 B1/UU|/ min UEO(NT)
|L LBUIF.
(10)
In the definition of Bures distance (equation 10), the minimization is over all orthogonal matrices U O(TT), whereas in the definition of adapted Bures distance (equation 9), the minimization over the subset of block diagonal matrices (dlag(R)..... the adapted Bures distance always larger than the Bures distances: AB(A, B) Z B(A, B). The minimizations over rotations in equation 9 and equation 10 essentially amount to differences in the allowed couplings. In the definition of adapted Bures distance (equation 9), the rotations restricted to N x N spatial rotations, which implics that the couplings preserve temporal ordering. However, in the definition of Bures distance (equation 10), there are no restrictions on the rotations, the couplings stochastic processes may have sample trajectories that are close (so Wasserstein distance is small) and yet the processes are very different in terms of how much information is available early in the trajectory. In section 4.1, we provide a tractable example that illustrates the effect of preserving temporal causality when comparing two processes.
RT), R....., Rí O(N)}. Therefore,
change the temporal ordering. This is important because two
EXPERIMENTS
4.1 A SCALAR EXAMPLE
We first consider an analytically tractable scalar example that highlights the temporal causality prop- erty of Causal OT. which is essentially a Gaussian process version of (Backhoff et al., 2022, figure 1). Consider two mern-zero scalar Gaussian processes x = {x(t)} and y {y(t)} defined for t= 0. 1,2 (e.g., firing rates of two neurons over three time steps) that both start at zero and end with distribution N σ²) for some > 0. The difference between them is that the first step of is small stochastic and the second step of x is deterministic, whereas the first step of y is deterministic and the second step of y is stochastic (Fig. 3).
Specifically, for some e small,
z (1) ~ N(0, e) v(1) = 0,
x(2) = z(1)
v(2) ~ N(00

---

Published as a conference paper at ICLR 2025
x(())
v(t)
0-
-1
-2
Transition
Deterministic
Stochastic
Figure 3: Five samples of two scalar processes over two time steps. The two processes have the same marginal distributions at time = and = The difference between the process is whether the first second step is deterministic stochastic. In this example, 1.5 and 0.5.
When > 0 is small, the sample trajectories of these processes are very close (Fig. 3). However, x and are very different as stochastic processes: z(1) is known, then we have full information about z(2), whereas if (1) is known, then we have additional information about v2).. In the context of a neural system, the first process could represent a neuron whose activity is ramping in time, whereas the second process could represent neuron that is silent until the final time step.
The difference between the stochastic processes is captured by Causal OT, but not by Procrustes distance, SSD Wasserstein. In particular, the distances (see Appx. C for details):
dprocrustes (I, y) = O,
d d1 y) =6
dWassersteii (I. y) =
d (x. y) 0 +
Procrustes distance cannot distinguish the two systems. Furthermore, - O, both dss yx and dwassersten (z. y) converge to zero, whereas (1-causa((;, y) converges to a. In particular, Causal OT captures differences in the across-time statistical structure of the stochastic processes that is not captured by SSD or Wasserstein.
4.2 SYNTHETIC EXPERIMENTS MODELING MOTOR SYSTEMS
We apply Causal OT distance compare simple models of neural activity in motor systems. We take inspiration from experimental measurements of motor cortical dynamics in which preparatory and movement-related neural activity evolve into orthogonal subspaces (Kaufman et al., 2014; Church- land & Shenoy, 2024), Fig. 4A. This effectively creates gating mechanism by which the neural activities can organize future movements without being "read out" into immediate motor actions. Here, we consider simple set up comparing two "systems"-on system includes preparatory ac- tivity and motor activity that evolve orthogonal subspaces and the other system only includes motor activity. We show that Causal OT distance can differentiate these two systems even when preparatory activity is small (which has been observed experimentally, Churchland et al., 2012; Kaufman et al., 2016), while SSD cannot.
The first system is a two-dimensional process x = (x(1) = (1) x2 xt)) that a simple model of motor cortex dynamics during reaching task (Churchland al., 2012; Kaufman et al., 2014). The process is governed by the following time inhomogeneous linear system (explained below)
Preparatory dynamics
x(t) x(t 1) ue, + +w() x(t) x - 1) + w(t)
(11a)
Rotational dynamics
(11b)
where e1, e2 are the standard basis vectors in R² and w = {t) } is a two-dimensional additive Gaussian noise process (note that is aligned with the vertical axis and e2 is aligned with the hor- izontal axis in Fig. 4A). Over the first T, 2 time-steps, the system evolves along the preparatory dimension (also referred to the output-null direction in the literature) according to equation 1la, where u takes values between [-1, 1] that determines the future output the system. This is fol- lowed by T, 2 time-steps of rotational dynamics, equation l lb, where M is the rotation matrix

---

Published as a conference paper at ICLR 2025
State space
High preparatary activity
preporatary activity
Pepusaay
High
Prepantory dimension
10F
251
10
Distunce
10F
7,
Reainnt dimusiion vi()
10
30
-35-
16
3
30"
Outpul pates/aal ratio
12"
10'
30
Figure 4: Motor cortex preparatory activity in the nullspace. A) State space activity x(t), accu- mulating inputs along preparatory dimension, before entering rotational dynamics with non-zero projection onto the readout dimension, showing different trajectories for different input scales (see colorbar). We show two different scales of preparatory activity in solid and dashed lines. B) Pro- jections of state-space activity along the preparatory and readout dimensions. The readout y(!) is a noisy version of the bottom row. C) Same as in B for the low preparatory activity scale, with scaling 7 0.3. D) Distances as a function of the relative scale of the readout axis with respect to preparatory, normalized so that each distance is when the output potent/null ratio is 1 (Procrustes shape distance and SSD overlap when the output potent/null ratio between 10-1 and 10²).
defined by
cos (
sin(0) cos(O)
M
sin
((2T/(2T.)
At the end of the T T T,, time steps, the process lies along the readout dimension (also referred
to as the output-potent direction in the literature), Fig. 4A (trajectories are shown as solid curves).
The projections of x(t) onto the preparatory axis and read axis are shown in Fig. 4B.
The second system is scalar process y = {v() that represents the downstream motor readout of first system and is a noisy projection of x(t) onto the readout axis:
Readout projection
v(() = e, x(1) + v((),
(12)
where vtf) is additive Gaussian noise. Similar to the scalar example from the previous section, the two process {x(()} and {y({)} have an important distinction: if the values of x(I) for L E [0, Tp. are known, then we have much more information about the values of x(l) for L E [T,, Tp + T.]. On the
other hand, if the values of yt) for t E [0, Tpl are known, then we have no additional information about the values of y(t) for t [Tp. Tp Tr]-
We apply Procrustes shape distance, 2-SSD, DSA and 2-Causal OT distance to compare the process and (since y is one-dimensional process, we embed it in IR via the mapping y egy). When preparatory activity is large, we find that Procrustes shape distance, 2-SSD, DSA and 2-Causal OT distance can each differentiate the two systems (Fig. 4D, vertical red line). However, often prepara- tory activity is small relative to motor activity (Churchland et al., 2012; Kaufman et al., 2016). We can model this by scaling the preparatory axis by a constant 7 > 0: (x1, + (7=1, This results in smaller preparatory activity (Fig. 4A, trajectories are dashed curves), consistent with ex- perimental observations. In this case, when preparatory activity small, we see that Procrustes shape distance and SSD are now much smaller whereas Causal OT and DSA distances do not sig- nificantly change (Fig. 4D, vertical dashed blue line versus vertical red solid line). Moreover, as preparatory activity shrinks 7 + 0 (output potent/null ratio grows), then Procrustes shape distance and SSD converge to zero, while Causal OT and DSA distances plateau (Fig. 4D). Importantly, this implies that Procrustes shape distance and SSD cannot capture differences between a system with small preparatory activity and motor activity and a system with only motor activity, while Causal OT and DSA distances can distinguish these two systems. In this way, both Causal OT and DSA distances capture important differences in the across-time statistical structure of the two processes that is not captured by Procrustes shape distance and SSD.
4.3 LATENT DIFFUSION MODELS FOR IMAGE GENERATION
Denoising diffusion models (Sohs-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020; Kadkhodaie & Simoncelli, 2021) are powerful conditional and unconditional generative models
8

---

Published as a conference paper at ICLR 2025
Tescrunio
D5A
2550
2 Caroel OT
bel ianple injoden
Ne rea Tapetts
Figure 5: Comparisons of sample trajectories generated from text-to-image stable diffusion models v1-1, v1-2 for 10 different prompts. Top row: Distance matrices show the estimated pairwise distances between the conditional processes. Bottom row: We show one example of random trajectory from the "primate" and "bird" prompts and their corresponding mean trajectories decoded into the pixel space. While individual trajectories for each prompt category generate an image with features corresponding to that category, the mean trajectories decoded into the pixel space do not contain prompt-related visual features.
with wide range of applications (Yang et al., 2023). While there are well-established methods for probing internal representations deterministic and static neural networks, these methods have not (to our knowledge) been widely applied to diffusion models (Klabunde et al., 2023). Furthermore, the forward and reverse diffusion processes leveraged by these generative models are stochastic dynamical systems, making them ideal candidates for the methodology we have developed.
Here, we investigate the extent to which the metrics (Procrustes, DSA, SSD, Causal OT) can dis- tinguish stochastic dynamical trajectories generated by diffusion model. Given the nonlinearity of diffusion processes, it's certainly possible that none of these measures can capture meaningful similarities and differences between these trajectories. It's also possible that these methods can cap- ture meaningful similarities and differences given enough samples, but the high-dimensionality of these systems means that estimation of the statistics would require an impractical number of sample trajectories.
To test our framework, we consider two pretrained text-to-image latent diffusion models (v1 1 and v1-2 -2' trained to generate text-conditional images from noise (Rombach et al., 2022). Given text prompt, each model generates 24 dimensioonal stochastic latent trajectory model. prompt whose terminal point is decoded into an image (with dimensions 256 x 256 x 3). In particular, each model and prompt corresponds to distinct stochastic dynamical system. Therefore, we ex- pect that for each model, modell prempt model. prompt) O for the same prompts, and potentially d prompt A model, prompt B) > 0 for different prompts (though this is less clear due to the nui- sance transform). Moreover, the training dataset used for vl-i for v1-2, so we might expect their dynamics to be similar, in which case we should find that d1x prompe x=1 pyompp x 0.
subset of the dataset used
We prompted each diffusion model with 10 different prompts from both animate and inanimate sources: "bird", "car", "fish", "human face", "invertebrate", "machine", "plant", "primate", "robot", "toy". These prompts were inspired by the stimulus categories used to probe image representations in humans (fMRT) and non-human primates (single-cell recordings) (Kriegeskorte et al., 2008). For each prompt and each model, we generated 60 latent trajectories and decoded those trajectories into the image space (examples
two decoded trajectories are shown in Fig. 5 and for several other trajectories in Fig. 7 of the supplement). We repeated this process for 3 random seeds to use the cathiny distances baseline. This provided 60 datasets (2 diffusion models, 10 prompts, 3 seeds per prompt) each containing 60 latent trajectories.
Models were taken from http://tugiingfaec.ooCmCmMiii

---

Published as a conference paper at ICLR 2025
For each pair of processes, estimated the Procrustes shape distance, DSA distance, SSD and Causal OT distance (Fig. 5). Estimating the distances in the 2 2 latent space is in- tractable, so we first projected each set of latent trajectories onto its top 10 principal components (PCs) before computing the distances between the 10-dimensional stochastic trajectories. To ensure that 10 PCs are sufficient for recapitulating the data, we recomputed the pairwise distances as the number of PCs vary from 2 to 20 and showed that the distances indeed converge (Supp. Fig. 6).
Visual inspection of the pairwise distance matrices reveals that SSD and Causal OT capture struc- tured relations between the trajectories that are not captured by Procrustes distance or DSA. First, both the SSD and Causal OT distance matrices have 3 x 3 block structure, while the Procrustes and DSA distance matrices do not. The 3 x 3 block structure reflects the expected smaller within- category distances than between-category distances; that is, d(xA,, x²) < d(xx, x8) ) for A & B. This structure also suggests that the number of trials used for computing distances provides sufficient statistical power for capturing meaningful distances. Second, SSD and Causal OT appear to capture similarities in the representations between models v1 11 and v1-2, where as Procrustes does not. Specifically, SSD and Causal OT distance appear to mainly depend on the prompt (i.e., they are invariant to the choice of model). This is reflected in the fact that the four quadrants of the distance matrices have similar apparent structure. On the other hand, this structure is not apparent in the Procrustes or DSA distance matrices. Together, these results suggest that it is important to account for correlation structure when comparing the similarities of these latent stochastic trajectories.
To better understand why Procrustes distance fails to capture structured differences between pro- cesses, we investigated whether the mean latent trajectories contain any visual information about the prompt categories. To this end, for each model and category. we passed the mean latent tra- jectory through the decoder and visually inspected the resulting image (bottom row of Fig. 5, also Supp. Fig. 7). The generated images are largely indistinguishable from each other suggesting that Procrustes distance fails to capture important stochastic aspects of the internal representations.
There are some other interesting aspects related to the distance matrices. First, there an apparent hierarchy of representational distances encoded in the block diagonals which opens up interesting questions about the differences between generative and discriminative models of vision for the cog- nitive neuroscience community (e.g., see Supp. Fig. 8). Causal OT distances are larger and appear less variable in the block diagonals suggesting that the full covariance structure is informative for computing shape distances. Importantly, distances that capture the dynamic and stochastic aspects of representations may be useful tools for investigating diffusion models.
5 DISCUSSION
In this work, we introduced Causal OT distance for comparing noisy neural dynamics. The distance is natural extension of existing shape metrics (Procrustes and SSD) that accounts for both the stochastic and dynamic aspects of neural responses, while also respecting temporal causality. We applied our method to compare simple models of neural responses in motor systems and found that Causal OT distance can distinguish between neural responses encoding preparatory dynamics and motor outputs versus neural responses that only encode motor outputs. We also applied our method to compare representations in latent diffusion models whose generation process is a stochastic dy- namical system. We found that Causal OT distances can differentiate generative processes while Procrustes distance and DSA cannot (though SSD can), emphasizing the need to consider stochas- ticity when analyzing such systems.
There are important limitations. First, the Causal OT distance only captures first- and second-order moments, which sufficient for distinguishing Gaussian processes; however, is not guaranteed to capture differences between non-Gaussian processes. Second, estimating first and second moments can require a significant number of sample trajectories. For example, an N -ddimensional trajectory with T time points has NT x NT covariance matrix. When N or T are large, accurate estima- tion of this covariance matrix can require more sample trajectories than are available. Theoretical analysis of estimating shape distances and other representational dissimilarity measures from finite data remains an active area of study (Pospisil et al., 2023). One promising approach is to impose priors on the stochastic processes reduce the number of trials required to estimate the first- and second-order moments (Duncker et al., 2019; Nejatbakhsh et al., 2023; Geadah et al., 2025).
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Bruno B Averbeck, Peter E Latham, and Alexandre Pouget. Neural correlations, population coding and computation. Nature Reviews Neuroscience, 7(5):358-366, 2006.
Julio Backhoff, Mathias Beiglbock, Yiqing Lin, and Anastasiaa Zalashko. Causal transport in dis-
crete time and applications. SIAM Journal Optimization, 27(4):2528-22) 2562, 2017.
Julio Backhoff, Daniel Bartl, Mathias Beiglböck, and Johannes Wiesel. Estimating processes in adapted Wasserstein distance. The Annals of Applied Probability, 32(1): 2022.
Julio Backhoff- Veraguas, Sigrid Källblad, and Benjamin A Robinson. Adapted Wasserstein distance between the laws of SDEs. arXiv preprint arXiv:2209.02243, 2022.
Rajendra Bhatia, Tanvi Jain, and Yongdo Lim. On the bures-wasserrtein distance between positive definite matrices. Expositiones Mathematicae, 37 2019.
Steven L Brunton, Bingni Brunton, Joshua Proctor, Eurika Kaiser, and J Nathan Kutz. Chaos as intermittently forced linear system. Nature Communications, 8(1):19, 2017.
Mark M Churchland and Krishna V Shenoy. Preparatory activity and the expansive null-space. Nature Reviews Neuroscience, 25(4):213-236 2024.
Mark M Churchland, John P Cunningham, Matthew T Kaufman, Justin D Foster, Paul Nuyujukian, Stephen Ryu, and Krishna V Shenoy. Neural population dynamics during reaching. Nature, 487 (7405):51-555 2012.
Thomas Cover and Peter Hart. Nearest neighbor pattern classification. IEEE Transactions on Infor- mation Theory, 13(1):21-27, 1967.
Sanjoy Dasgupta and Philip M Long. Performance guarantees for hierarchical clustering. Journal of Computer and System Sciences, 70(4):555-569, 2005.
Frances Ding, Jean-Stanislas Denain, and Jacob Steinhardt. Grounding representation similarity through statistical testing. Advances Neural Information Processing Systems, 34:1556-1568, 2021.
Lea Duncker, Gergo Bohner, Julien Boussard, and Maneesh Sahani. Learning interpretable continuous-time models of latent stochastic dynamical systems. In International Conference on Machine Learning, pp. 1726-1734. PMLR, 2019.
Lyndon Duong, Jingyang Zhou, Josue Nassar, Jules Berman, Jeroen Olieslagers, and Alex H Williams. Representational dissimilarity metric spaces for stochastic neural networks. In The Eleventh International Conference on Learning Representations, 2023.
Aniruddh R Galgali, Maneesh Sahani, and Valerio Mante. Residual dynamics resolves recurrent contributions neural computation. Nature Neuroscience, 26(2):326-33) 2023.
Victor Geadah, Amin Nejatbakhsh, David Lipshutz, Jonathan W Pillow, and Alex Williams. Modeling neural activity with conditionally linear dynamical systems. arXiv preprint
1833,7,502..8347, 2025.
James E Gentle. Numerical Linear Algebra for Applications in Statistics. Springer Science & Business Media, 2012.
Robbe L. T. Goris, Anthony Movshon, and Eero P. Simoncelli. Partitioning neuronal variability. Nature Neuroscience, 7766):858-653 Jun 2014.
Madhu Gunasingam and Ting-Kan Kam Leonard Wong. Adapted optimal transport between gaussian processes in discrete time. arXiv preprint arXiv:240.006625, 2024.
Sarah E Harvey, David Lipshutz, and Alex H Williams. What representational similarity measures imply about decodable information. In UniReps: 2nd Edition of the Workshop on Unifying Rep- resentations in Neural Models, 2024.
11

---

Published as a conference paper at ICLR 2025
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 6851, 2020.
Zahra Kadkhodaie and Eero Simoncelli. Stochastic solutions for linear inverse problems using the prior implicit in denoiser. Advances in Neural Information Processing Systems, 34:13242- 13254, 2021.
Matthew Kaufman, Mark M Churchland, Stephen Ryu, and Krishna Shenoy, Cortical activity in the null space: permitting preparation without movement. Nature neuroscience, 17(3):440- 448, 2014.
Matthew T Kaufman, Jeffrey S Seely, David Sussillo, Stephen Ryu, Krishna Shenoy, and Mark M Churchland. The largest response component in the motor cortex reflects movement timing but not movement type. eneuro, 3(4), 2016.
Max Klabunde, Tobias Schumacher, Markus Strohmaier, and Florian Lemmerich. Similarity of neural network models: A survey of functional and representational measures, arXiv preprint (rXiv:2305.06329 2023.
Bernard O Koopman.
Hamiltonian systems and transformation in hilbert space. Proceedings of the National Academy of Sciences, 17(5):315-318 1931.
Simon Korbllith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International Conference on Machine Learning, pp. 3519- 3529. PMLR, 2019.
Nikolaus Kriegeskorte, Marieke Mur, and Peter A Bandettini. Representational similarity analysis- connecting the branches of systems neuroscience. Frontiers in systems neuroscience, pp. 4, 2008
Rémi Lassalle. Causal transport plans and their Kante-Kantooovich problems. Stochastic Analysis and Applications, 36(3):452 484, 2018.
David Lipshutz, Amin Nejatbakhsh, and Alex H Williams, Disentangling recurrent neural dynamics with stochastic representational geometry. In ICLR 2024 Workshop on Representational Align- ment, 2024.
Anton Mallasto and Aasa Feragen. Learning from uncertain curves: The 2- 2-ssssern metric for gaussian processes. Advances in Neural Information Processing Systems, 30, 2017.
Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representa- tions of words and phrases and their compositionality. Advances in neural information processing systems, 26, 2013.
Amin Nejatbakhsh, Isabel Garon, and Alex Williams. Estimating noise correlations across continu- ous conditions with Wishart processes. Advances Neural Information Processing Systems, 36, 2023.
Mitchell Ostrow, Adam Eisen, Leo Kozachkov, and Ila Fiete. Beyond geometry: Comparing the temporal structure of computation neural circuits with dynamical similarity analysis. In Ad- vances in Neural Information Processing Systems, 2023.
Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data sci- ence. Foundations and Trends@ in Machine Learning, 11(5-6)335--60 607, 2019.
Dean A Pospisil, Brett W Larsen, Sarah E Harvey, and Alex H Williams. Estimating shape dis- tances on neural representations with limited samples. In The Twelfth International Conference on Learning Representations, 2023.
Maithra Raghu, Justin Gilmer, Jason Yosinski, and Jascha Sohl-Dickstein. SVCCA: Singular vector canonical correlation analysis for deep learning dynamics and interpretability. Advances in Neural Information Processing Systems, 30, 2017.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
12

---

Published as a conference paper at ICLR 2025
Britton A Sauerbrei, Jian-Zhong Guo, Jeremy D Cohen, Matteo Mischiati, Wendy Guo, Mayank Kabra, Nakul Verma, Brett Mensh, Kristin Branson, and Adam Hantman. Cortical pattern generation during dexterous movement is input-driven. Nature, 577(77900:386-39, 2020.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing. pp. 2256-2265. PMLR, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Floris Takens. Detecting strange attractors in turbulence. In Dynamical Systems and Turbulence, Warwick 1980: proceedings of a symposium held at the University of Warwick 1979/80, pp. 366- 381. Springer, 1981.
Saurabh Vyas, Matthew Golub, David Sussillo, and Krishna V Shenoy. Computation through neural population dynamics. Anmual Review of Neuroscience, 43:249-275, 20200.
Alex H Williams, Erin Kunz, Simon Kornblith, and Scott Linderman. Generalized shape metrics on neural representations. Advances Neural Information Processing Systems, 34:4738 4750, 2021.
Ling Yang, Zhilong Zhang, Yang Song. Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: comprehensive survey of methods and applications. ACM Computing Surveys, 56(4):1-39, 2023.
13

---

Published as a conference paper at ICLR 2025
A CAUSAL OPTIMAL TRANSPORT DISTANCE BETWEEN GAUSSIAN PROCESSES
In this section, we provide first-principles derivation of the Causal OT distance between two processes defined in equation The distance can be derived as an optimal transport distance between two Gaussian processes and where the transport cost is Li-distance and the ad- missible couplings correspond the set of so-called causal synchronous couplings. These are syn- chronous couplings because both Gaussian processes
driven by common white noise process w = {t] and they are causal in the sense that at any given time point, the current values of the Gaussian processes x()) and (t) do not depend on the future values of the white noise process w(1 +/++00 1), w((+ 2),....
This distance is closely related bi-causal optimal transport distances (also referred as adapted Wasserstein distances) between Gaussian processes (Lassalle, 2018; Backhoff et al., 2017; Gu- nasingam & Wong, 2024). These are optimal transport distances in which the couplings between the processes are required to satisfy "bi-causal" property, which can be interpreted as "the past of process independent of the future of process conditioned on the past of process 2" and, conversely. "the past of process 2 is independent of the future of process 1, conditioned on the past of process 1", Here our couplings are defined terms of causal synchronous couplings rather than bi-causal couplings. These two notions of coupling are equivalent in the scalar setting (Backhoff- Veraguas et al., 2022); however, we are unaware of proof of this in the multi calti-dimensional setting.
{w((t) Suppose x and y are N-dimensional Gaussian processes defined for t L...., T. Let be a sequence i.idd N-dimensional standard normal vectors (i.c., white noise process). We define causal synchronous coupling between and y pair of processes (&, y) that satisfy the following criteria:
Synchronous coupling condition:
there is function that maps the white noise process w to (x, 9); (ii) is equal in distribution to x and ý is equal in distribution to y.
Causality condition:
(iii) for each < T, there a function that maps (w(1)..... w()) to (()).
We now construct family of causal synchronous couplings. Let X be the NT-dimensional Gaussian random vector obtained by concatenating xI1),...,x((1
x(1)
X
N (mi Ca):
x(T)
where m. and C₂ denote the mean and covariance of the concatenated vector X. Let C. = LT be the Cholesky decomposition of the covariance matrix C., C2 full rank, then the decompo- sition is unique. If C. is rank R < N, then we take L, to be the unique lower-triangular matrix with exactly positive diagonal elements and - R columns containing all zeros (Gentle, 2012). Given N -ddimensionall driving white noise process (w(1).... ................................ w(T)}-is -ic., w(t) are i..d. N dimeensiioaal Gaussian random vectors with identity covariance define the process & = {x(t))
by
x(1
w(1)
= (w) = m. + L W
W
x(T)
Since the distribution of Gaussian process be completely characterized in terms of its first two moments, the process & is equal in distribution to the process x. Furthermore, by construction & is a continuous function of the driving white noise process w and, since L, is lower triangular, for any t= 1= 1,
x(1
w(1) : w(t)
fF()) = -ma,, + Lu,
W
(13)
(())
14

---

Published as a conference paper at ICLR 2025
where mz,t denotes Nt-dimensional vector that is equal to the first Nt coordinates of mzg and Lg denotes the Nt x Nt matrix that is equal to the first Nt rows and Nt columns of Lg. We say that & is adapted 1, the vector x(() is function of the first steps of the white noise process w. As corollary, x(t) is independent of the future values of the white noise process w(t +1).-.;:(()) t t= 1,
the driving white noise process since equation 13 holds; that is, for each
T 1,....
w(T).
Let f*(-) and (-) be defined as above but with x's replaced by y's. Then = f"(w) and f is causal synchronous coupling between and y. In general, consider causal synchronous couplings of the form:
x ff (w).
y = (R.
where Ri...., RT are N x N orthogonal matrices. Since the distribution of white noise is invariant these transformations, these define synchronous couplings of x and Furthermore, the rotations do rotate the white noise process time, so they define causal synchronous couplings of x and y.
Having defined our admissible couplings, we now define the Causal OT distance
deausal (x. y) = min R...
[If((w) (Ir Q) ROw(T))1 (14)
where the minimization is over nuisance transformations Q E O(N) and over spatial rotations R..... Rr E O(N) of the noise. Substituting in with the formulas for f and f", we see that
E (w) - (Ir @ (R) ..... RTw(T)))
lm (Ir mm 11 +
| (Ir ,,, Ri....., Rt) w[l
lm (Ir Q m L - (Ir 8 L, RI..... R1
Substituting this expression back into equation 14, we get
L y) min QEO(N) {Im= (Ir & Q) AB (Ir & (I
where ABNT((..)) is the adapted Bures distance between positive semidefinite matrices defined
AA,
min
L Lndiag(
R EN((
where A = LAL] and B = LALA are the Cholesky decompositions. In the scalar setting, this reduces to the adapted Bures distance defined (Gunasingam & Wong, 2024). While this metric is motivated as distance between Gaussian processes, it defines proper (pseudo-)metric between any stochastic processes terms of their first- and second-order statistics.
It's worth noting that if we relax the causality condition then we can consider larger set of (acausal) synchronous couplings:
& = f (w),
ÿ = (f((((())
where U E O(NT) is any NT NT rotation and
w(1)
(w) =
w(T)
is rotation of the white noise process that leaves its distribution invariant. Minimizing the L distance over all such couplings results in the following distance, which is equal to the Wasserstein distance (with p 2) between Gaussian random vectors (Mallasto & Feragen, 2017) up nui-
15

---

Published as a conference paper at ICLR 2025
sance transformation:
Wasserstein (x, = min min QEO(N) UEOINT)
E [wl - (Ir @ ())
L (Ir @ Q)
min QEO(N)
mm - (Ir 2 Q) m
min UEO(NT)
E
min QEO(N)
Im, (Ir 00 Q) m)
min UEO(NT)
| - (Ir s Q)L
min QEO(N)
|m (IT 2 Q m
min UEO(NT)
(C.) (Ir & (
[3]] }
min QEO(N)
m= (Ir Q m,
1|
(C,, (Ir @ (I,
where (..,)) is the Bures distance on positive semidefinite matrices:
B(A, B) = min UEO(NT)
|| - /l/ r.
16

---

Published as a conference paper at ICLR 2025
B ALTERNATING MINIMIZATION ALGORITHM
Suppose x (x()) an Ng-dimensional Gaussian process and y (y( an Ng Ny-dimensional Gaussian process. If N₂ * Ny we can pad the lower-dimensione. process with zeros so that they are both N-dimensional processes, where N := max(N,, Nyy. For example, Ng < N, then we can embed y() into IN via the linear transformation
y()) 0
y()
where 0 is (N Nyy)-dimensional vector of zeros. For the remainder of this section, we assume that N = N₂ Ny.
From equation 8, we have that causal OT distance equal
min {Im (Ir @ m), L. (Ir @ L,igd R...... R
(15)
Q. .i.... Rr
We solve this via an alternating minimization algorithm. For an NT NT matrix and S 8, 8,15 T, let Le denote the (s, t)t N N block of L. We first rewrite the second term
12= (I L, Ry ..... R
Lu,st ot
Lir,st
Qly.st
D
QLy
Rr
1=1
Lz,Te
Q
Then optimizing over R, yields:
L
Q L
2
R. arg min R
R
L
QL
Alternatively, we can express the sum in equation 15 as
mm (I @ mm
12
+ L2
(I
@
L, R
Rt) 1
2
R L y
m x,
m y.1
s,11
+
Q
m T
m T.T
TT
R M.T
m
T
m
m yr R 9,11
QT T
L
L TT
R
g.Tr
F
Then optimizing over Q yields
T
T
m
g.
v.1
m Z.
m T
Q
arg min Q
L L
R L
9,11
LT
R L 9.Tr
17

---

Published as a conference paper at ICLR 2025
We summarize the previous steps into algorithm. Given
L11
m
L
mT
LTi
LTT
define the functions
mi
Lut
m
F((L) =
G(m, L) =
LT
L
L LTTJ
Let Cholesky(-) be function that maps positive semidefinite matrix to its Cholesky decomposi- tion. When the matrix is full rank, this decomposition unique. When the matrix is rank R N, select the unique decomposition with R. positive diagonal elements and N R columns whose entries are all zero.
Algorithm 1: Alternating minimization for computing causal distance
1: input: 0 S a 2, my E NT, my E RTT C2 E S47 and C2 E SNT.
2: initialize: Q ONN)
3: Lg < -Cholesky(C) C₂)
4: Ly + Cholesky
5: while not converged do
6: for 1,..., do
7:
Mz, + Fi (L))
8:
My.1 + Fi( & R))
9:
R, arg miggiiREOON M R.B
10:
end for
11:
N. + G(m,, Lx)
12: Ng + g(m Lydiag( Ri..... R2r)
Qy arg argmingEO(N) |N, N2Q T
14: end while
15: dist + (22 a)) m (Ir @ @)m,, + L2
(Ir @ Liagg Ri,..., Rr)
16: return dist
18

---

Published as a conference paper at ICLR 2025
C SCALAR EXAMPLE
We derive the distances reported in Sec. 4.1. The means of x and are identically zero so the Procrustes distance is zero. The covariance matrices of (æ(1), x2)) and (v(11, v()) satisfy
2
60
C2
L
o²
0
0
0
C
by
0
0²
From these expressions, we see that 1 distance (Bhatia et al., 2019), the Wasserstein distance is
Using the explicit formula for the Bures
daasserstein (I. y) = B(Cr, Cyy)
Tr())) + TT(C 2Th (VC,CoVOC,) 1/2 v(e+ + + -222
=6
where we have used the fact that VCy = Ly in the third equality. Finally, by equation 8, the Causal OT-1 distance is
deausat.1 (x,, v) = AB(C2, Cyy)
min
[ diag((a ()1 B)
(=11(=1)
min $E(11)
0 - Ba
=V
19

---

Published as a conference paper at ICLR 2025
D ADVERSARIAL TUNING OF NOISE CORRELATIONS AND INPUT
In this section we describe how to build adversarial examples of systems with distinct recurrent dynamics that share their marginal statistics. To achieve this, will the imput and noise correlations to these systems. Denoting a reference system that has marginal means m, and marginal covariances P. our goal is build linear systems for arbitrary recurrent dynamics matrices A such that their marginal statistics are equivalent to the reference system.
To begin, we rewrite the dynamics of linear system driven by inputs b, and noise correlations ΣΣ
L+41 = 2; + At + dtb, + VdtE,G [[ [I + A 2 + dtb;
Cov(2 (I + dt. Cov + dtA dIE,
Therefore in order have R[x: m, and Cov( P. need to set b, E, according the following equations.
b. m++1 (I+ m)
by = mo
dt
E, P1+1 (I + dt + ddA))
Σg / Po
dd
In Fig. 2 we first generated data from the Saddle dynamics since it is the only nonlinear model among the 3. We then used its marginal means and covariances to tune the inputs and noise correlations for other two systems (Line and Point attractor).
For completeness, we include the dynamical equations for the systems below:
Saddle (assuming that {-0.1, 0.1) the coherency level):
(x) 0. + 2x 0, f2(2;) - 82 + 0
dx
f f(x) Et
dt
Point Attractor (assuming the linear dynamics mentioned above with A dynamics matrix):
0.5
0
A =
-1
0
Line Attractor (assuming the linear dynamics mentioned above with A dynamics matrix and Rot(ő) is the rotation matrix for angle 6):
0
A =
Rot
Φ o= -= 45"
0.7 -1
20

---

Published as a conference paper at ICLR 2025
E DETAILS OF THE DIFFUSION MODEL EXPERIMENT
Stable diffusion for text-to-image generation are latent diffusion models that consist of few com- ponents. We denote images by random variable s, text by random variable p. and latent diffusion process by random variable x1:7. An encoder & maps images into the latent space and a decoder D maps the latent diffusion back into the image space. Text prompts are mapped into a fixed-length embedding To (p) that can be used for conditioning the latent diffusion model. The text embedding uses an architecture that well-suited for the text modality.
Separately, time-conditioned UNet architecture is used to run the backward diffusion process for denoising images in the latent space. We denote the denoiser by co (x,x t). Conditioning on text performed via mapping the text representation to intermediate layers of the denoising UNet using cross-attention mechanism (see Rombach et al. (2022) for details). Finally, the conditional diffusion process cost function L is used for the joint optimization of the denoiser and text embedding:
C = p.
||e ex t, To
Once the components of the model are trained, the function Eg provides an approximation of the score function V log p(x p) for large t.
During the image generation, first the encoding of the text prompt is computed. Then starting from random iid noise in the latent space, the following model run forward:
x++1 x, + (x (o xx) + 7/2--
In our experiments, we generated samples from this stochastic process conditioned on different prompts and computed pairwise distances between different conditional processes using Procrustes, DSA, SSD, and Causal OT distances. For DSA we chose the hyperparameters n. n.delle delays rank = We fixed all the other hyperparameters to the following for this and other DSA experiments: delay delay-inted -aayyiineevalll distances we showed that both SSD and Causal OT capture our desired properties. Here, we present more results validating our findings.
9,
1, 1r 0.01, iters = 1000. By computing these pairwise
The size of the images are 256 x 256 (3 RGB channels) and the size of the latent space is 64 x 64 4. Therefore the latent representation provides computational gain with factor of approximately 16. Since the dimension of the latent space is still very large, we performed PCA the latent trajectories normalized distances on k: E (2.4, 6, 8, 10, 15, 203. We normalized the distances to the median value of each distance matrix focus on the relative changes. The normalized distance changes are very small after k = 6 confirming that the distances are not an artifact of the lower-dimensional projection.
map onto the top PCs. In Fig. 6 show the dependence of the
In addition, we include the decoded sample trajectories and mean trajectories from all prompt cate- gories for both models vl-1 and v1-2 Fig. The decoded mean trajectories show very subtle signatures of the image category confirming that distances based on mean trajectories (e.g. Pro- crustes) are not sufficient to capture differences the stochastic processes.
Finally, we included the MDS projection of all three pairwise distances in Fig. 8. This figure shows that both SSD and Causal OT the conditional stochastic processes corresponding to the same prompt category for different seeds and diffusion models cluster together. This suggests that MDS embedding of the stochastic process distances is meaningful and potentially relevant for the cogni- tive neuroscience community.
21

---

Published as a conference paper at ICLR 2025
Intent dim 2
latent
Iatent dim =
latent dim
Intent dim 10
latent
latent dim 29
Canacal oT
SSD
Procrustes
Figure 6: Convergence of normalized shape distances with small number of principal components (PCs). Pairwise distances from three different methods (rows; Causal OT, SSD, Procrustes) between 60 conditional stochastic processes shown where the latent trajectories projected onto top PCs for varying e {2,4,6, 6, 10, 15, 20) (columns). All normalized shape distances converge after d = 6.
22

---

Published as a conference paper at ICLR 2025
sample trajectory
trajectory
bird
car
human face
invertebrate
machine
plant
primate
robot
fish
human face
invertebrate
machine
plant
primate
robat
toy
Figure 7: Decoded sample and mean trajectories for all categories in models v1-1 and v1-2.
Procrustes
DSA
SSD
Causal OT
Viil
but
B000
prinaic
nbvt
-0015
plax
Figure 8: Multi-dimensional scaling of the pairwise shape distances for Procrustes, DSA, SSD, and Causal OT distances corresponding to the distance matrices shown in Fig. 5.
23

---

Published as a conference paper at ICLR 2025
F COMPARISON OF METRICS ON A SCALAR SYSTEM
Here, using simple scalar dynamical system model (e.g., the dynamics of a single neuron's is re- sponses), we demonstrate that DSA can distinguish systems with varying dynamics (i.e., different underlying vector fields) but the same marginal statistics, while 2-SSD cannot. Conversely, we show that 2-SSD can distinguish systems with varying noise levels but the same dynamics, while DSA cannot. Finally, Causal OT can distinguish the systems in both cases while Procrustes cannot distinguish the systems in either case.
Specifically, consider the scalar autoregressive process (of order I)
xx1)) -ar(l 1) ow((),
where 1 is scalar that determines the autocorrelation of (a(t)), 8 > is the input noise level and {w(tt} is driving white noise process (i.e., a sequence of i.id. standard normal random variables). The process {a(()} has a unique stationary distribution given by ( o²), where o² := 80/(1 a²). Therefore, if the process (x(()) is initialized so that zx(0) ( then its marginal distribution satisfies z(t) (O, o²) for all =
F.1 SYSTEMS WITH DIFFERENT DYNAMICS, SAME MARGINAL STATISTICS
We first show that when the dynamics are different across systems but the marginal distributions are constant across systems, then Causal OT distance and DSA can discriminate between the systems but Procrustes and 2-SSD cannot. For each E {0.1, 0.3, 0.5, 0.7, 0.9] we let 5² = 1 - a² so that the stationary distribution of (z(t) } is NO, 1) provided (0) (O, 1). As increases, the sample trajectories become smoother and have larger autocovariances (Fig. 9).
0.1
a 0.3
0.5
0.7
0.9
2
-2
5
0.8
0.6
0.4
0
10
10
10
0
10
10
Figure 9: Five autoregressive processes with varying dynamics but constant marginal variances. The top row shows five sample trajectories from each process along with the empirical mean (thick red line) and marginal variance (dashed purple lines). The bottom shows the autocovariance matrices.
Based on visual inspection of the pairwise distance matrices (Fig. 10), it is clear that Causal OT and DSA differentiate between the systems with varying autocorrelations, but Procrustes and SSD do not.
F.2 SYSTEMS WITH THE SAME DYNAMICS, DIFFERENT NOISE STATISTICS
Next, show that when the vector field is fixed across systems but the noise varies, then Causal OT and SSD can discriminate between the systems but Procrustes and DSA cannot. Specifically, fix the vector field 0.3 and consider different noise levels ă E (..1, 0.3, 0.4, 0.5). As
24

---

Published as a conference paper at ICLR 2025
Figure 10: For each (0.1, 0.5, 0.7, 0.9}. we generated 3 sets of 1000 samples trajectories of length 7 10, from which we estimated the pairwise distances between the 5 x 3 = 15 sets of sample trajectories. To compute the DSA distance we used the following parameters: n_delays
3, delay-interval = 1, rank 30,
6=0.1
5=0.3
6=0.5
6=0.7
5=0.9
3
2-
1-
0
-1-
-2
-3
0-
5-
0.6
0.4
0.2
0.0
10
0
10
10
10 0
10
Figure 11: Five autoregressive processes with fixed dynamics and varying noise levels. The top row shows five sample trajectories from each process along with the empirical mean (thick red line) and marginal variance (dashed purple lines). The bottom row shows the autocovariance matrices. To compute the DSA distance we used the following parameters: n_delays 3. intervant = 1, rank = 30.
the noise level increases, the trajectories exhibit larger fluctuations about the mean however, the autocovariance structure remains about constant (Fig. 11).
It immediately clear from visual inspection that Causal OT and SSD differentiate between the systems with varying noise, but Procrustes and DSA do not (Fig. 12).
550
Canstor
DSA
Figure 12:
For each noise levels 8 E {0.1, 0.4, 0.5), we generated sets 1000 samples trajectories of length T = 10, from which we estimated the pairwise distances between the 5 3 15 sets sample trajectories.
25

---

Published as a conference paper at ICLR 2025
These experiments scalar linear systems demonstrate that SSD and DSA can respectively dis- criminate stochastic and dynamic aspects of stochastic processes; however, on their own they do not capture both elements. On the other hand, Causal OT distance discriminates both stochastic and dynamic characteristics.
26

---

Published as a conference paper at ICLR 2025
G SENSITIVITY OF DSA ON HYPERPARAMETERS
DSA (Ostrow et al., 2023) crucially depends on hyperparameters: the number of delays used when creating the Hankel matrix and the rank selected when fitting the reduced-rank regression model to linearly estimate the system's 's dynamics. We applied DSA (with different choices of delays and rank) the 2-dimensional dynamical systems from Fig. of the main text demonstrate the sentivity of the method to the choice of hyperparameters (Fig. 13). In Fig. 2, we presented the result when using the best performing hyperparameters.
rank=5
rank=10
rank=20
1I n_dellys==
Line
Point
Saddio
dedelas==1
Line
Point
Sudile
Line
Point
Swddlee
Line
Point
Snidile
Line
Point
Figure 13: Sensitivity of DSA to hyperparameter selection. select n_delays and rank on grid shown in rows and columns, and run DSA on the same toy dataset as in Fig. 2.
27


Reviews:
[{"review_id": "quBCwhRbKD", "reviewer": "Reviewer_4ecC", "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.", "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. \n- 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? \n- 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.\n- I would like to see experiments comparing Causal OT to DSA in some way.\n- 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "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. \n- The paper provided three sets of experiments: scalar task comparing the different metrics mathematically, synthetic for biological motor control, and latent diffusion models. \n- 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. \n- Estimating full covariance matrices for high-dimensional neural data (large N and T) is computationally intensive and probably requires a prohibitive number of samples.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IXsBPVpvlr", "reviewer": "Reviewer_Sz7A", "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).", "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._\n\nEDIT 11/21--The authors have included the analysis I asked for. I have accordingly raised my score. \n\n### Minor points:\n- L046-L047: The phrasing makes it sound like both studies addressed stochasticity and dynamics.\n- L050: You don't actually show this for DSA, as far as I can tell. \n- Equation 3: There is a missing time dependence in $A^\\top$.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6msPiKxLFZ", "reviewer": "Reviewer_mGZm", "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.", "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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The presentation throughout is very clear and a pleasure to read.\n* The motivation of the method through prior work is exceptionally good.\n* The method this paper proposes is intuitive but original, fairly simple, well-motivated, and seems to be a substantial extension of previous work.\n* 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?\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)