Extraction and Recovery of Spatio-Temporal Structure in Latent Dynamics Alignment with Diffusion Models
Authors: Yule Wang
Zijing Wu
Chengrui Li
Anqi Wu
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Neural distribution alignment
Diffusion model
Neuroscience
Neural decoding
TL;DR: We propose ERDiff, a method for neural distribution alignment, which preserves the trial-based spatio-temporal structure of latent dynamics with diffusion models.

Abstract: In the field of behavior-related brain computation, it is necessary to align raw neural signals against the drastic domain shift among them. A foundational framework within neuroscience research posits that trial-based neural population activities rely on low-dimensional latent dynamics, thus focusing on the latter greatly facilitates the alignment procedure. Despite this field's progress, existing methods ignore the intrinsic spatio-temporal structure during the alignment phase. Hence, their solutions usually lead to poor quality in latent dynamics structures and overall performance. To tackle this problem, we propose an alignment method ERDiff, which leverages the expressivity of the diffusion model to preserve the spatio-temporal structure of latent dynamics. Specifically, the latent dynamics structures of the source domain are first extracted by a diffusion model. Then, under the guidance of this diffusion model, such structures are well-recovered through a maximum likelihood alignment procedure in the target domain. We first demonstrate the effectiveness of our proposed method on a synthetic dataset. Then, when applied to neural recordings from the non-human primate motor cortex, under both cross-day and inter-subject settings, our method consistently manifests its capability of preserving the spatio-temporal structure of latent dynamics and outperforms existing approaches in alignment goodness-of-fit and neural decoding performance.

Extraction and Recovery of Spatio-Tempone Structure in Latent Dynamics Alignment with Diffusion Models
Yule Wang Georgia Institute of Technology Atlanta, GA, 30332 USA
Zijing Wu Georgia Institute of Technology Atlanta, GA, 30332 USA zwu3818gatech, edu
yulewang@gatech.. edu
Chengrui Li Georgia Institute of Technology Atlanta, GA, 30332 USA
Anqi Wu Georgia Institute of Technology Atlanta, GA, 30332 USA anqivuêgatech. edu
malichengrui@gateee. edu
Abstract
In the field of behavior-related brain computation, is necessary align raw neural signals against the drastic domain shift among them. A foundational framework within neuroscience research posits that trial-based neural population activities rely on low-dimesiional latent dynamics, thus focusing on the latter greatly facilitates the alignment procedure. Despite this field's progress, existing methods ignore structure during the alignment phase. Hence, their
the intrinsic solutions usually lead to poor quality in latent dynamics structures and overall performance. To tackle this problem, we propose an alignment method ERDiff, which leverages the expressivity of the diffusion model to preserve the spatio- temporal structure of latent dynamics. Specifically, the latent dynamics structures of the source domain are first extracted by a diffusion model. Then, under the guidance well-recovered through maximum
of this diffusion model, such structures likelihood alignment procedure in the target domain. We first demonstrate the effectiveness of proposed method on synthetic dataset. Then, when applied to neural recordings from the non-human primate motor cortex, under both cross- day and inter-subject settings, method consistently manifests its capability preserving the spatio-temporal structure of latent dynamics and outperforms existing approaches in alignment goodness-of-fit and neural decoding performance. com/ EllevanggNT//RDDii.
Codes are available at:
1 Introduction
key challenge severely impeding the scalability of behavior-related neural computational appli- cations is their robustness to the distribution shift of neural recordings over time and subjects M. Given a behavior model trained on previous neural recordings (e-g., velocity predictor for human with paralysis (2). it usually suffers performance degradation when applied to new neural recordings due to the neural distribution shift [3.4). Thus. for long-term usability and stable performance the trained neural decoding model, high-quality alignment between the neural recordings used for training (i.e., source domain) and new recordings for testing (i.e., target domain) is of vital importance.
Distribution alignment an important task at the heart of unsupervised transfer learning 6. The goal is to align the target domain to the source domain that the trained model in the source domain can be applied to the target domain after eliminating the distribution shift. However, due to issues such instabilities and low signal-to-noise ratio 77, raw neural spiking activities are noisy and ambiguous g 9, causing difficulties in aligning the distributions of these high-dimensional signals directly.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

One promising research direction [10] points out that the trial-based neural activities can always be understood in terms of low-dimensional latent dynamics [TT 12 13. Such latent dynamics manifest coordinated patterns of evolution constrained to certain "neural manifolds' [1] [5] Hence, early studies focusing on the alignment of latent dynamics reach comparably satisfactory results [T6] [7]. Generally, most previous methods [16] 18 [9] are based on pre-defined metric for optimization during latent dynamics alignment, i.c., minimizing the difference evaluated by the metric, between source and target domains within the low-dimensional latent space. However, those metrics are usually non-parametric and handcrafted, which are not guaranteed to suit specific neural recordings or problems well. Methods based on adversarial-leamiin (0 21] thus have been introduced since they can implicitly find an adapted metric [22] However, they suffer from mode collapse and instability
issues in practice
23
Moreover, during the alignment process, we note that the above-mentioned works lack the necessary awareness of the latent dynamics structure, especially when aligning non-linear and lengthy trials. Through an empirical study on the motor cortex of non-human primate (NHP) g (shown in Figure D}. we can observe that: state-of-the-an alignment method JSDM [24] (minimizing the symmetric Jensen-Shannon divergence between distributions) fails to recover the latent dynamics structures of the source domain since JSDM neglects those structures during alignment. From another perspective, in the alignment phase, existing methods fail to effectively model and leverage the information-rics correlations between each time bin and each latent dimension within latent dynamics.
In this paper. we focus on preserving the temporal evolution of each individual latent dimension and the spatial covariation between latent dimensions of the source domain during alignment. The main idea that we first extract the spatio-tempora structure of latent dynamics from the source domain; and then, we align the target domain by recovering the source domain's 's underlying structure. However, such a workflow non-trivial since the underlying spatio-temporal structure both implicit and complex.
Source
JSDM Aligned
To tackle this problem, we propose novel alignment method that is capable of Extracting and Recovering the latent dynamics structure with Diff uiffus model (ERDiff). Firstly, given the source-domain neural observations, we use a diffusion model (DM) 25 26 to extract the spatio- temporal structure of latent dynamics. Then, in the align- ment phase, we propose a maximum likelihood align- ment procedure through the guidance of DM, by which the spatio-temporal structure of source-domain latent dy- namics can be recovered well in the target domain. The proposed extract-and-recove method nicely encodes and preserves the spatio-temm temporal structure of latent dynam- ics, which are significant inductive biases for neural latent dynamics alignment. Furthermore, from the perspective of core machine learning, ERDiff introduces an approach of extracting structure knowledge from one distribution and imposing it as the prior to constrain the alignment of another distribution. Note that although we have been emphasizing extraction and recovery of the source-domain structure, ERDiff not performing a sopy-and-paste of the source domain distribution to the target domain. As ERDiff preserves the dynamics structure of the source domain, it also maintains the original characteristics of the target domain. We present experimental results support this argument. Finally, extensive experiments to verify the effectiveness of ERDiff on a synthetic dataset and two real-world neural datasets &8 (7). Visualization of latent dynamics also demonstrates that ERDiff is capable of preserving the spatio-temporal structure consistently during the alignment phase.
Spatio-temporal Structure
Spatio-temporal Structure Missing
Figure 1: Empirical study. Latent dy-
namics (3D visualization) of the source domain and the aligned target domain by JSDM primary motor cortex dataset.
conduct
2 Preliminary
Distribution alignment. We denote the source-domain observations of single-trial neural population activities as (*)
E |R where I is the trial length (i.e., number of time bins), and is the number of observed neurons. We denote its low-dimensional latent dynamics as
x . (a)
T
Rixd where d is the latent dimension size. Generally, we build a variational autoencoder (VAE) to estimate the latent (+) given the observations (*). The VAE

---

consists of a probabilistic encoder gZ X (
Φz)
and probabilistic decoder p p( ((7), xb)).
Φ, and 3 are the parameters of the encoder and decoder. The encoder also serves as an approximated posterior distribution to the intractable true posterior ((Z(*)
(x))). Then in the target domain, given x E IR we perform distribution alignment
the neural population activities X = x
by linear probing the probabilistic encoder q(Z X: Φ). This alignment phase is conducted by minimizing certain probability divergence (D) between the two posterior distributions:
min (z
(+); Φ (l)(((((()))
x ())).
(I)
Diffusion (probabilistic) model (DM). Given / x d-dimensional i.i.d. samples Z from unknown data distribution, DM (28) aims to approximate such distribution by fitting the parameters of neural network pe (Z). DM is composed of a forward process followed by a reverse process. In the forward process, isotropic Gaussian noise is added to diffuse the original data, which can be defined
linear stochastic differential equation (SDE) form:
dz fZz t)dt g())dw
where f()) Rixd x R + Rixd is the drift coefficient, g(.) R R is the diffusion coefficient, and w is the standard Wiener process. The solution of the SDE is a diffusion process (Z, (Ze)e in
which [0, 2] is fixed time zone. In this paper, we implement them with VP-SDE (25]. {Ze (0,
approaches the standard normal prior distribution (Z) when i T. Under mild conditions drift and diffusion coefficients (25]. the denoising reverse process can be solved in the following closed-form SDE:
dz (f V log [ dt + g())dw
where Vz pe is the score function, and w is reverse-time Wiener process. We train parameterized network sZ, () to fit the score function Vz po not directly accessible and resort the denoising score matching (DSM) 29] for optimization:
pa (z)). However, Vz () is
CDSM (0)
(Zo) (xt)) ||Vz, . Pog
Zo) s sz t: 0)12]
(4)
where U represents the uniform distribution and x(t) is the weighting function. Under VP-SDE, the transition probability Por (Z, Zo) also follows a Gaussian distribution A (u,Zo, Σ,), in which H4, Σ, E x On the other hand, according to [28] we can define noise estimator with the score function as eZ, 0) = - K, T s (Ze, f; f), in which K, K, = Σ, Invoking these expressions, we can thus reformulate the form of DSM loss based the Fisher Divergence between noise terms:
CDSM (0) = EE [ 7. ((t)) ||c t(,,, 0)1
(5)
in which w(t) KtA(t) and Z, Zu + K.C.
3
Methodology
In this section, we introduce our proposed latent dynamics alignment method ERDiff in detail.
3.1 Maximum likelihood alignment
Given the source-domain neural activities ( we infer their latent dynamics (z)) by building VAE.
We use variational inference to find the probabilistic encoder ((Z(o) (x) X (8); (b)) and probabilistic
decoder pl X (z)); bba) through maximization of the evidence lower bound (ELBO) 30):
Φs. 2b argmax 6.
E (
[x)g(((()))
(z)), [b]]] {
DeL
(g(Z(()
X (Z))))())))
which q (z) is the normal prior. Note that we introduce ERDiff with this basic VAE architecture. But ERDiff can be combined with many variants of latent variable models (LVM) 3 32 The essence of FRDiff to tune the parameter set d of the probabilistic encoder, regardless of the model architecture of the encoder and decoder.
3

---

(A)
Extraction
(B)
Recovery
Domain
Target-Domain
E
ST Block
Temporal Structure
Estracted Spaño-
Unalggeed Labent Dyramis
Structure-away
MLA Guidance
Dynamics
x
Spatio
Temporal
Figure 2: A schematic overview of spatio-temporal structure extraction and recovery in ERDiff. (A) The architecture of DM for spatio-temporal structure extraction. (B) A descriptive diagram of structure recovery schematic. The left presents the extracted spatio-temporal structure of the source- domain latent dynamics; the right illustrates the structure-aware maximum likelihood alignment guidance in ERDiff.
Practically, aligmment methods that directly match the discrete samples from the source and target domains in a pair-wise fashion may lead to sub-optimal solutions 33 34 especially when the collected samples from the target domain are limited. Thus given the target-domain neural activity x()), we propose to perform neural distribution alignment via maximum likelihood estimation (MLE);
argmax E Xi [logPs (h(X; +))] argmax (z~g(Z(()) [logp, (Z)]:
in which p. (Z) represents the ground-truth probabilistic density of latent dynamics in the source domain and h(-) refers to the non-linear transformation from X to Z underlying the probabilistic encoder {Z) X; Φ). The objective in Eq. implies that, instead of minimizing a distance metric between source observations and target observations, we aim at maximizing the likelihood where the density comes from the source domain and the data comes from the target domain. The left-hand-side (LHS) is the MLE for observation X and the right- right-hand-side (RHS) is the MLE for latent Z. We will focus on the RHS in the following sections. We note that the RHS objective implies that we will optimize the encoder parameter & during alignment so that the latent encoder will map X (!) to proper latent ( who fits the source density p. (Z) well.
3.2 Spatio-temporal structure extraction and source domain learning
In order to calculate the objective function in Eq. we need to know two density functions: g( X; Φ) is defined in the original VAE model with the learnable parameter Φ; (Z) is the density of latent Z for the source domain. The latter is inaccessible by building a VAE alone. Therefore, the first step is to learn (Z) given only ((s), We propose to learn (Z) through training a DM.
To fully capture (Z). the DM should consider the overall spatio-temporal structure of latent dynamics. To extract such a structure, the DM can not treat each latent state or time bin within latent dynamics as mutually independent and feed them into the model sequentially. We thus take the entire trial of latent dynamics Z((') ~ q(: X()); &s) as input to the DM for training. Specifically, the DM fits P. (Z) through the training of denoiser e(Z, t: 0.) (Rixd x R + Rixd
Next, we describe the architecture of <: which is refined for extracting the global spatio- temporal structure of latent dynamics. Traditional architecture based on 2D-Convolution Layers [35] focuses on capturing the local features within latent dynamics, which can hardly extract its global spatio-temporal dependency or structure. Thus, we adopt an architecture mainly derived from Diffusion Transformer (DiT) 36 37). Specifically, we propose to use Spatio-T T'emporal Transformer Block (STBlock), shown in Figure A). Each STBlock is composed of Spatio Transformer layer followed by a Temporal Transformer layer, which are l-layer encoders based on multi-head self- attention. The Spatio Transformer layer takes latent states of each time bin as inputs to extract spatial structure, whereas the Temporal Transformer layer takes the entire latent trajectory of each latent space dimension as inputs to extract temporal structure. (see Appendix for details of the architecture of DM).

---

For the training objective of c(:;0),, we sample noisy targets 2,0) and minimize the following DSM
loss function:
0, = argmin 100 B2 (o)
utt) - ( --((
(8)
We note that Z,0 here are actually latent dynamics inferred via VAE in Eq. 6 To enrich the input samples and adequately estimate the source density P. (Z) as motivated earlier, we propose to learn the VAE objective (Eq. (6) and the diffusion objective (Eq. ( simultaneously. In each training iteration, conditioning on the current value of Φ, and v/., we obtain a set of Zo = (x) o.) and
use it as Z to optimize Eq. g We can also optimize VAE first, obtain an optimal Φs, and use it to optimize Eq. 8 Experimental results show that the former approach achieves higher density estimation performance compared the latter (see Appendix for details).
3.3 Spatio-temporal structure recovery and distribution alignment
Given the trained denoiser e(Z, t: 0g), go through the reverse process from f T to t 0 in Eq 3 and obtain the marginal distribution Pn (Z: 0,). We use Po (Z; 0.) to approximate (Z) in
Eq. The maximum likelihood estimation thus be written
argmax Zz [log Po Z; 0.)]
Φ
We perform alignment by tuning the parameter set & of the probabilistic encoder while keeping the DM Pa (Z: fixed. Note that we have already optimized the VAE objective obtain an optimal Φ, using source data. During alignment, we first set Φ as ds and then linear probe & (c.g.. neural observation read-in layer). Consequently, we not only initialize the model with good encoder but also make optimization during alignment much faster and more efficient.
In the reverse process, the computation of log Po (Z:0,) is tractable through the probability flow ODE (25) whose marginal distribution at each time step matches that of our VP-SDE. However, the direct computation of log Pu Z: 0&) will require invoking the ODE solver in each intermediate time step 3 39 Such complexity is prohibitively costly for online neural applications. circumvent
this issue, we can reform Eq. (9) as follows:
(z) (b) [log Po [Z: 0.]] = DKL. gg( ((); Φ) Po (Z: 0.)) + + (Z X :)) (10)
where the first term the KI. divergence from the DM marginal distribution po (Z; 0.) to the probabilistic encoder distribution (Z X()); Φ),, and the second term HI(-) denotes the differential
entropy. For the DKI. (-) term in Eq. 10 via the Girsanov theorem 40, 41 we have
DKI. (g(Z) X (1); h) (p [Z: 0.)) < LDSM (db 0.) + DKL (pr (pr (Z; (Z())
(11)
where LDSM is the denoising score matching loss in Eq. 5p. and pr(-) is the distribution final time step I of Eq. (2). Consequently, we could obtain an upper bound of the maximum likelihood objective, as follows (we provide detailed derivation in Appendix B):
-zz-g(2(x00 [log Pa [Z; 0.)l < DKL (pr (Z; 0, (Z))
Constant Torm
(12)
+ Bonnuoo.TBoonaZZZaaa (o0
wo mo(1) l|e (Z (Z,,,;0,511
- (Z,,t)
Weighted Noise Residual
Divergence
Since xZ) is fixed prior distribution, it does not depend on parameter Φb. Thus, our optimization objective will include only the latter two terms, which are more computationally tractable. The first objective simplifies to weighted noise residual for the parameter set & and the second divergence objective can be approximated using the Hutchinson-Skilling trace estimator [42]. We note that recovery of spatio-temporal structure is primarily conducted by the weighted noise residual part, in which the probabilistic encoder obtains alignment guidance awareness of the spatio-temporal structure from e e t; 0.). This procedure is illustrated in Figure B).
5

---

In distribution alignment, it is a common practice to directly leverage the ground-truth data samples by introducing a regularizer term the optimization function. To encourage the diversity of latent dynamics after alignment, here we further compute and penalize the Sinkhorn Divergence [43
between the latent dynamics samples of source domain Z!*) ~ q(- X and that of target domain (1) ~ q(- [
min (7, +)F++H(++ XH(())
(13)
where each value C[i = 2 2 in matrix C denotes the squared Euclidean cost move
probability mass from - to Z and (7) computes the entropy of transport plan 7. The total loss for distribution alignment composed of the term in (13) and the latter two terms on the right side of (I2). We note that the total loss is minimized only with respect to the probabilistic encoder parameter set Φ (see Appendix C for the total loss formula and the detailed algorithm of ERDiff.)
4 Experiments
Datasets. We first train and evaluate ERDiff with a synthetic dataset. Then we apply ERDiff to non-human primate (NHP) dataset with neural recordings from the primary motor cortex (MI), in which the primates are performing a center-out reaching task in 8 different directions. The NHP dataset contains rich cross-day and inter-subject settings that provide us with ideal test bed.
Baselines for comparison. We compare ERDiff against the following two strong baselines proposed for the neural distribution alignment task:
JSDM [2] a metric-based method that leverages discrete samples from both the source and target domains. The alignment is performed through the symmetric Gensen-Shanoon divergence [44]. Cycle-GAN [2I] a of-te-of-the-ar GAN-based method that uses cycle-consistent adversarial networks to align the distributions of latent dynamics.
Considering the neural observations and latent dynamics in the format of multi-variate time series, we also compare ERDiff with the following methods aiming at distribution alignment for general time series data:
SASA [45] metric-based distribution alignment method for time series data regression task through the extraction of Nomain-invariaat representation.
DANN [46]: an adversarial learning framework in which domain classifier is followed by feature extractor through gradient reversal layer. This layer adjusts the gradient by multiplying with predefined negative constant during the training process.
RDA-MMD [7]: a distribution alignment method via minimizing MMD Loss between the latent dynamics extracted from LSTM.
DAF [88: : an adversarial learning framework that uses a transformer-basse shared module with domain discriminator. During the adaptation step, the domain-invariant features are invariant (Q,
K of self-attention); the domain-spectie -speiff fcatures (V of self-attention) keep tuning.
4.1 Synthetic dataset
Data synthesis and evaluation metrics. We first generate simulated latent dynamics dataset to illustrate the effect of our ERDiff method on spatio-temporal grructure-preservii and distribution alignment performance. In this setting, we consider modeling the nonlinear latent dynamics to follow conditionally Continuous Bernoulli (CB) [49] distribution. For each single-trial latent dynamics, generate 2-dimensional latent variables Z {zı.L} and their 32- 22-dimensional observations X = {xI:L}, }. where = 32. We use the following synthesis process and parameter settings to generate samples for the source and target domains, respectively:
p
(x)
2
= II
CB
(2(00)
Z (+()) tanh tanh(z( (12)))
(4)
p
x
(5)
(3)
= =N
x
R
R())(())
.).
p
2 = II
CB
((1) (anh(z) (2))
x
(t)
= N
K
d
(14)
where E (1,...,
[i....,LJ. L... and (W(+), (())), (W(), R0)) are the specific parameter sets of the source and target domains. To compare and evaluate the latent dynamics alignment performance, we estimate the trial-average log density of the aligned latent dynamics evaluated at the optimal generation

---

(A)
N
KLD
(B)
Source
ERDiff (Ours)
JSDM
22
Oycle-
RRp MMD
ERCH
Z1
Figure 3: Experimental results on the synthetic dataset. (A) Performance comparison on trial- average negative log-likelihood (NLL) and KL Divergence (KLD). means the lower the better. ERDiff achieves the recond-lowees NLL and the lowest KLD. (B) True continuous Bernoulli dynamics in the source domain compared to the latent dynamics aligned by ERDiff and JSDM in the target domain (blue dots denote the fixed points). ERDiff preserves the spatio-temporal structure of latent dynamics much better.
1/- logg" (2) and the trial-averaged KL Divergence to the optimal latent dynamics distribution: 1/L El DKL (po- (z/2)
distribution:
zi (c) Plpooo (zi+1 z (
Results on synthetic dataset. We repeat the simulation experiment five times and report the mean and standard deviation of each method in the above two quantitative evaluation metrics, shown in Figure AA. We observe that ERDiff achieves higher alignment performance on both two evaluation metries compared to baseline methods. For further analysis, we plot the phase portrait of the true source domain and those inferred by ERDiff and JSDM in Figure B Compared JSDM, ERDiff can extract and recover the spatio-temporal structure of the synthetic latent dynamics more precisely and be much closer the ground truth. These results mainly due the fact that ERDiff obtains structure-awee aware alignment signals from the DM while JSDM neglects this structural information.
4.2 Neural datasets
We conduct extensive experiments on two real-world neural datasets: the non non-human-- -rimate (NHP) primary motor cortex (MI) [8] and Rat hippocampal CAI [27]. Experiments on these two datasets use nearly identical setup. Here we primarily discuss the experimental approach and results related the NHP motor cortex dataset. (See Appendix D for the detailed results on the rat hippocampal
CA1 dataset.)
Motor cortex dataset description. We conduct experiments with datasets collected from the primary motor cortex (M1) of two non-human primates ('C' & 'M') [ The primates have been trained reach one of eight targets at different angles (Figure 4A). Neural recordings from these two primates have been widely studied [20 [50]. During such a process, their neural spike activities (signals) in the primary motor cortex (MI) along with the reaching behavior velocity were recorded. They performed the center-out reaching task multiple times in each direction and only successful trials were saved. For our experimental evaluation purpose, we select the trials from three recording sessions for each primate per day. In total, we have days for each primate. We will perform cross-day (recordings of the same primate performing the task on different days) and inter-subject (recordings of different primates) experiments.
Data processing and evaluation metries. The neural recordings of each day and each primate consist of about 180-220 trials across 3 sessions. For each trial, about 200 neurons are recorded and the number of time bins is 39 with 20ms intervals. We also bin the velocity of the primate's 's behavior into 39 bins. Therefore, we have time-aligned neural data and behavioral data. When training with the source data, we optimize the VAE model together with the DM. One thing we need to emphasize here is that we also include velocity- latent the velocity values, which is ridge regression model. Therefore, the inferred latent contains a rich amount of velocity information. During testing, we align the test neural data to the training neural data so that we can directly apply the velocity decoder to the latent the test data without performance degradation. In the training session, the ratio of training and validation set is split as 80%:20% through 5-fold cross-validation. The post-processed dataset of primate 'C' contains 586 trials in total while that of primate 'M' contains 632 trials. For the evaluation protocol, since
loss to the VAE loss. The decoder maps the neural
7

---

(A)
(C)
ERDiff (Ours)
DAF
JSDM
Cross
Day
R 0.257
R² = 0.062
R² = 0.032
(B)
Inter Subject
Source
R 0.234
R² = 0.049
R² = 0.048
Figure Motor cortex dataset and experimental results. (A) Illustration the center-out reaching task of non-human primates. (B) The 3D visualization of trial-averaged latent dynamics corresponding to each reaching direction the source domain. (C) The 3D visualization of trial trial-averaged latent dynamics corresponding to each reaching direction aligned by ERDiff, DAF, and JSDM given the target distribution from cross-day and inter-subject settings. We observe that ERDiff preserves the spatio-temporal structure of latent dynamics well.
the ground- truth source domain distribution of latent dynamics inaccessible, we use the behavior decoding performance evaluate the performance of latent dynamics alignment. Here compare true behavior velocity with the decoded one in the test data using coefficient of determination values (R22, in %) and noot-mean-square error (RMSE). To verify the generalization of each method in latent dynamics alignment, we make full use of the dataset collected in chronological order. We perform 6 sets of cross-day experiments and 10 sets of inter-subject experiments, all repeated over 5 different random seeds.
Experimental setup. The VAE is based on a sequential architecture [51] in which recurrent units are applied in both the probabilistic encoder and decoder. We also add domain knowledge of our alignment task into the model structure: behavior regression decoder is cooperatively trained from the latent dynamics so that the behavior semanties information is complementarily provided during the neural manifold learning. Poisson negative log-lik likeeikeeiood loss is used for firing rate reconstruction and mean squared error used for behavior regression. We use the Adam Optimizer [52) for optimization and the learning rate chosen among (0.005, D. The batch size is uniformly set as 64. Despite the varying size the input dimension (due to the varying number of recorded neurons in different sessions), the latent space dimension size is set 8 for all the methods for fair comparison. We use the dropout technique [53] and the ELU activation function (54) between layers in our probabilistic encoder and decoder architectures. During latent dynamics alignment, we perform linear probing only on the read-in layer of the probabilistic encoder while keeping the remaining layers fixed.
Neural manifold analysis. Considering the interpretability (50) and strong latent semantics con- tained in neural manifold [8], we conduct case study based on the visualization of neural manifolds to verify the spatio-temporal structure preserving capability and alignment performance ERDiff. In Figure B)), plot the averaged latent dynamics of each direction in the source domain, which based on one recording on primate 'C" using demixed Principle Component Analysis (dPCA) [55]. The parameters of dPCA fit with the source-domain latent dynamics while being fixed when applied to perform the transformation
the target domain. In Figure 4CC), we plot the averaged latent dynamics of each direction aligned by ERDiff and two representative baseline methods (DAF and JSDM) under both cross-day and inter-subject settings.
Under both experimental settings, the overall observation is that the spatio-temporal structure of the aligned results ERDiff is much more coherent with that of the source domain. The results of DAF and JSDM roughly recover the direction of dynamics but fail to preserve the spatio-temporal

---

(A)
Source
ERDiff (Ours)
JSDM
(B)
Sampling Density
02-
05-
-32
ERDiff
JSDM
-17
1fº = 0.257
Rº 0.032
535 45%
45%
Sampling Density
30%
30%
25% 27%
33%
10%
Figure 5: (A) True source nource-domain trial velocities and behavior decoding trajectories inferred from a ridge regression model given the latent dynamics aligned by ERDiff and JSDM, respectively, We can observe that ERDiff not only preserves the apatio-tempore structure but also decodes the direction more accurately. (B) We compare the decoding performance of ERDiff, DAF, and JSDM with decrease in the sampling density of trials on the target domain. We can observe that ERDiff maintains a relatively high accuracy under low sampling densities.
Table I: The R-squared values (R²2, in %) and RMSE of the methods on the motor cortex dataset. ERDiff w/o S is short for variant of our proposed method that removes the spatial transformer layer in the DM. ERDiff w/o T is short for variant of our proposed method that removes the temporal transformer layer in the DM. The boldface denotes the highest score. Each experiment condition repeated with runs, and their mcan and standard deviation are listed.
Cross-Day
Inter- Subject
Method
(%)) t
RMSE
R2(%) t
RMSE &
Cycle-GAN
-24.83 (#3.91)
11.28 (+0.44)
-25.47
(+3.87
12.23
(#0.46)
JSDM
-17.36
(+2.57)
9.01 (+0338)
-19.59
(+2.77)
11.55 (+0.52)
SASA
-12.66 (#2.40)
8.36 (+0332)
-14.33
(+3.05)
10.62 (=0.40)
DANN
-12.57
(4:328)
8.28 (#0.32)
-18.37
(4:324))
10.66
(+.057)
RDA-MMD
-9.96
(+2263)
8.51 (+0.31)
-6.31 (42.19)
10.29
(+0.42)
DAF
-6.37
(+3.72)
8.17 (+0.48)
-11.26
(+:364)
9.57(++0.8)
ERDiff w/o S
-12.69
(#2.64)
8.57 (+0550)
-14.60
(+2.88)
10.85 (00)7
ERDiff w/o
-14.61
(+2.33)
8.93 (40.50)
-17.10
(43.233
10.94 (+0.59)
ERDiff (Ours)
18.81(42.22
7.99 +0.43)
10.29
22.66)
9.78( +0.50)
structure tightly. That is because JSDM neglects the is difficult for adversarial networks to capture such structure implicitly in DAF. Additionally, the averaged latent dynamics of each direction are much more clearly separated through ERDiff. We owe this outcome to the fact that providing extra guidance on the spatio-temporal structure would also facilitate the model to align directions properly. Additionally, without any mean offset alignment, the starting points (from the bottom center) and ending points of latent dynamics are also aligned well with the source domain, further verifying the structure recovering ability of ERDiff.
structure during alignment and
Decoding performance comparison. Table shows a comparison of the r-squared value (R²) and average RMSE on both the cross-day and inter-subject settings. While Figure A) depicted the decoded velocity trajectories of a subset of trials on the cross-day setting given the method. We have traditional alignment methods specially designed for
the following observations: (1) Compared neural data, deep learning-bast based methods additionally model the sequential information the latent dynamics, thus achieving better alignment results, which reflects the importance of spatio-temporal structure modeling. In most cases, ERDiff achieves the highest decoding accuracy and alignment performance among all methods. (2) From the ablation study shown the bottom of Table D find that both the Spatial Transformer layer and Temporal Transformer layer are key components in ERDiff, verifying the effectiveness of spatio-temporal structure modeling. (3) As shown in Figure A)A), the spatio-tempont structure of the latent dynamics is well pellerreeer in the result of DRfi Compared to baselines, the smoothness and continuity of the trajectory decoded by ERDiff also more satisfying.
Impact of sampling density in the target domain. We verify the consistent performance of ERDiff in few-shot target-sample circumstances. In Figure we analyze the impact of sampling density of the target domain on decoding performance. The setting is that we sample portion

---

of target-domain data to learn the alignment and apply the alignment to the entire target domain. Despite the sampling density drops from 50% to 10%, our results demonstrate that ERDiff continues to produce fairly consistent decoding accuracy with a small drop. This result validates our argument that ERDiff both preserves the dynamies structure underlying neural activities and maintains the characteristics of the target domain. In comparison, the performance of baseline methods shrinks drastically because they lack prior knowledge of the spatio-temporal structure.
We can conclude that the DM in ERDifT is capable of extracting the spatio-temporal structure in the source domain latent dynamics, providing a valuable inductive bias in recovering such structure during distribution alignment.
Discussion
In this work, we propose a new method named ERDiff, for solving the neural distribution alignment issue in real-world neuroscience applications (e.g.. brain-computer interfaces). Firstly, with the source domain, we propose to use a diffusion model to extract the ppati-ttmporal structure within the latent dynamics of trials. Next, in the alignment phase with the target domain, the spatio-temporal structure of latent dynamics is recovered through the maximum likelihood alignment based on the diffusion model. Experimental results on synthetic and real-world motor cortex datasets verify the effectiveness of DRDiff in the enhancement of long-term robustness and behavior decoding performance from neural latent dynamics.
To be in line with the conventions of previous studies on neural distribution alignment 4 16 the behavioral (velocity) signals of the source domain present during the VAE training. These signals do help in learning more interpretable neural latent space. However, we emphasize that ERDiff does not incorporate any behavioral signals of the target domain during the distribution alignment phase. Hence, ERDiff is entirely an unsupervised neural distribution alignment (i.e., test-time adaptation) method. As for the computational cost analysis, In the source domain training phase, the additional computation cost of ERDiff primarily comes from the diffusion model (DM) training. We note that the DM is trained in the latent space Z, which is significantly lower in dimensionality than the raw neural spiking signal space X. Therefore, the computational overhead of this one-time training phase acceptable, especially given that can be conducted offline real-world BCI applications. In the alignment phase, we would like to emphasize that ERDiff maintains a comparable computational cost with baseline methods. Please refer to Appendix E for a comprehensive analysis of ERDiff's computational cost and time complexity.
Limitation and future work. (1) Currently, ERDiff can align well with a single source domain neural latent distribution. An intriguing direction for future work would be learning a unified latent space across multiple source domains using the diffusion model. Thus the method would be applicable to domain generalization problems. (2) Generalization on alternative latent variable models (LVM). In this paper, ERDiff identifies the latent variables of raw neural spiking signals with a canonical version of VAE. However, the architecture of the LVM within ERDiff is actually disentangled from the diffusion model training or MLA procedure. Future work includes validating ERDiff given more advanced implementations of LVM, e.g., LFADS (56) STNDT
32
Broader impact. Not confined to practical applications in systems neuroscience, the maximum likelihood alignment (MLA) with diffusion model algorithm proposed in ERDiff has great potential to apply broader domain adaptation tasks across general time-series datasets (e.g., weather forecasting, and seismology). We also expect that our method can be applied or extended to other real-world scenarios and the broader field of neuroscience/AI.
References
[1] F. Briggs, G. R. Mangun, and W. M. Usrey, "Attention enhances synaptic efficacy and the signal-to-noise ratio in neural circuits." Nature, vol. 499, no. 7459, pp. 476-480, 2013.
[2] A. Fourtounas and S. J. Thomas, "Cognitive factors predicting checking, procrastination and other maladaptive behaviours: prospective versus inhibitory intolerance of uncertainty." Journal of Obsessiie-Comppusiv and Related Disorders, vol. 9, pp. 30-35, 2016.
[3] J. R. Manning, J. Jacobs, I. Fried, and M. J. Kahana, "Broadband shifts in local field potential power spectra are correlated with single-neuron spiking humans," Journal of Neuroscience, vol. 29, no. 43, pp. 13613-13620, 2009.
10

---

[4] J. Jude, M. G. Perich, L. E. Miller, and M. H. Hennig, "Robust alignment of cross cross-session recordings of neural population activity by behaviour via unsupervised domain adaptation."" " arXiv preprini arXiv:2202.06/59 2022.
[5] Z. Wan, R. Yang, M. Huang, N. Zeng, and X. Liu, "A review on transfer learning in eeg signal analysis." " Neurocomputing, vol. 421, pp. 1-14, 2021.
[6] S. Niu, Y. Liu, J. Wang, and H. Song. "A decade survey of transfer learning (2010-2020)"" IEEE Transactions on Artificial Intelligence, vol. 1, no. 2. pp. 151-166, 2020.
[7] I.. Hu, A. Mouraux, Y. Hu, and G. D. lannetti, "A novel approach for enhancing the signal-to- noise ratio and detecting automatically event-related potentials (erps) in single trials," Neuroim- age, vol. 50, no. 1. pp. 99-1 1, 2010.
[8] J. A. Gallego, M. G. Perich, R. H. Chowdhury, S. A. Solla, and L. E. Miller, "Long-term stability of cortical population dynamics underlying consistent behavior," Nature neuroscience, vol. 23,
no. 2, pp. 260-270, 2020.
[9] J. A. Gallego, M. G. Perich, R. H. Chowdhury, S. A. Solla, and L. E. Miller, "A stable, long-term cortical signature underlying consistent behavior," BioRxiv, p. 447441, 2018.
[10] M. M. Churchland, J. P. Cunningham, M. T. Kaufman, J. D. Foster, P. Nuyujukian, S. I. Ryu, and K. V. Shenoy, "Neural population dynamics during reaching," " Nature, vol. 487, no. 7405, pp. 51-56, 2012.
[11] H. Sohn, D. Narain, N. Meirhaeghe, and M. Jazayeri, "Bayesian computation through cortical latent dynamics," Neuron, vol. 103, no. 5, pp. 934 001-947,222 2019.
[12] S. Saxena, A. A. Russo, J. Cunningham, and M. M. Churchland, "Motor cortex activity across movement speeds is predicted by network-level strategies for generating muscle activity," Elife, vol. 11, p. e67620, 2022.
[13] R. Liu, M. Azabou, M. Dabagia, C.-H. Lin, M. Gheshlaghi Azar, K. Hengen, M. Valko, and E. Dyer, "Drop, swap, and generate: A self-supervised approach for generating neural activity." Advances in neural information procexsing systems, vol. 34, pp. 10587-10599, 2021.
[14] A. Gallego, M. G. Perich, S. N. Naufel, C. Ethier, S. A. Solla, and L. E. Miller, "Cortical population activity within a preserved neural manifold underlies multiple motor behaviors," Nature communications, vol. 9, no. 1, p. 4233, 2018
[15] R. Mitchell- Heggs, Prado, G. P. Gava, M. A. Go, and S. R. Schultz, "Neural manifold analysis of brain circuit dynamics in health and disease," Journal of Computational Neuroscience, vol. 51, no. pp. 1-21, 2023.
[16] B. M. Karpowiez, Y. H. Ali, L. N. Wimalasena, A. R. Sedler, M. Keshtkaran, K. Bod- kin, X. Ma, L. E. Miller, and C. Pandarinath, "Stabilizing brain-computer interfaces through alignment of latent dynamics," " bioRxiv, 2022.
[17] A. D. Degenhart, W. E. Bishop, E. R. Oby, E. C. Tyler-Kabara, S. M. Chase, A. P. Batista, and B. M. Yu, "Stabilization of a brain-computer interface via the alignment of low-dimensional spaces of neural activity," Nature biomedical engineering, vol. 4, no. 7, pp. 672-685, 2020.
[18] J. Lee, M. Dabagia, E. Dyer, and C. Rozell, "Hierarchical optimal transport for multimodal distribution alignment," Advances in Neural Information Processing Systems, vol. 32, 2019.
[19] R. Liu, M. Azabou, M. Dabagia, J. Xiao, and E. Dyer, "Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers." Advances in neural information processing systems, vol. 35, pp. 2377-2391, 2022.
[20] A. Farshchian, J. A. Gallego, J. P. Cohen, Y. Bengio, L. E. Miller, and S. A. Solla, "Adversarial domain adaptation for stable brain-machine interfaces," " arXiv preprint 1810v:1810.0004 2018.
[21] X. Ma, F. Rizzoglio, E. J. Perreault, L. E. Miller, and A. Kennedy, "Using adversarial networks to extend brain computer interface decoding accuracy over time." bioRxiv, pp. 2022-08, 08, 2022.
[22] W. Zhao, L. Mou, J. Chen, Y. Bo, and W. J. Emery, "Incorporating metric learning and adversarial network for seasonal invariant change detection," IEEE Transactions Geoscience and Remote Sensing, vol. 58, no. 4, pp. 2720-2731, 2019.
[23] 7. Zhang, M. Li, and J. Yu, "On the convergence and mode collapse of gan," in SIGGRAPH Asia 2018 Technical Briefs, pp. 1-4, 2018.

---

[24] Y. Wu, E. Winston, D. Kaushik, and Z. Lipton, "Domain adaptation with asymmetrically-relaxe distribution alignment," in International Conference on Machine Learning, pp. 6872-6881, PMLR, 2019.
[25] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, "Score-based generative modeling through stochastic differential equations," arXiv preprint arXiv:2011.13456 2020.
[26] H. Xue, A. Araujo, B. Hu, and Y. Chen, "Diffusion-based adversarial sample generation for improved stealthiness and controllability." arXiv preprini arXiv:2305. /6494, 2023.
[27] A. Grosmark, J. Long, and G. Buzsáki. "Recordings from hippocampal area cal, pre, during and post novel spatial learning," " CRCNS. org, vol. 10, p. K0862DC5, 2016.
[28] Ho, A. Jain, and P. Abbeel, "Denoising diffusion probabilistic models," Advances Neural Information Processing Systems, vol. 33, pp. 6840-6851, 2020.
[29] P. Vincent, "A connection between score matching and denoising autoencoders," Neural compu- tation, vol. 23, no. 7. pp. 1661-1674, 2011.
[30] R. Wei and A. Mahmood, "Recent advances in variational autoencoders with representation learning for biomedical informatics: A survey," leee Access, vol. 9, pp. 9939,4956, 2020.
[31] C. Pandarinath, D. J. O'Shea, J. Collins, R. Jozefowicz, S. D. Stavisky, J. C. Kao, M. Trautmann, M. T. Kaufman, S. I. Ryu, L. R. Hochberg, al., "Inferring single-trial neural population dynamics using sequential auto-encoders," Nature methods, vol. 15, no. 10, pp. 805- 815, 2018.
[32] T. Le and E. Shlizerman, "Stndt: Modeling neural population activity with spatiotemporal transformer," arXiv preprint arXiv:2206. 04727, 2022.
[33] N. Courty, R. Flamary, A. Habrard, and A. Rakotomamonjy. "Joint distribution optimal trans- portation for domain adaptation," Advances in neural information processing systems, vol. 30, 2017.
[34] T. Kerdoneuff, R. Emonet, and M. Sebban, "Metric learning in optimal transport for domain adaptation," in Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence, pp. 2162-2168, 2021.
[35] O. Ronneberger, P. Fischer, and T. Brox, "U-net: Convolutional networks for biomedical image segmentation." " in International Conference on Medical image computing and computer- intervention, pp.
Springer, 2015.
[36] W. Peebles and S. Xie, "Scalable diffusion models with transformers," arXiv preprint arXiv:2212.09744. 2022.
[37] Y. Tashiro, J. Song, Y. Song, and S. Ermon, "Csdi: Conditional score-based diffusion models for probabilistic time series imputation," Advances in Neural Information Processing Systems, vol. 34, pp. 24804-24816, 2021.
[38] R. T. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud, "Neural ordinary differential equations," Advances in neural information processing systems, vol. 31, 2018.
[39] W. Grathwohl, R. T. Chen, J. Bettencourt, I. Sutskever, and D. Duvenaud, "Ffjord: Free-form continuous dynamics for scalable reversible generative models," arXiv preprint 1rX1v:1810.01367, 2018.
[40] B. Oksendal, Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
[41] Y. Song, C. Durkan, I. Murray, and S. Ermon, "Maximum likelihood training of score-based diffusion models," Advances in Neural Information Processing Systems, vol. 34, pp. 1415-1428.
2021.
[42] J. Skilling, "The eigenvalues of mega-dimensional matrices," Maximum Entropy and Bayesian Methods: Cambridge, England, 1988, pp. 455-466, 1989.
[43] M. Cuturi, "Sinkhorn distances: Lightspeed computation of optimal transport," Advances in neural information processing systems, vol. 26, 2013.
[44] M. Menéndez, J. Pardo, I. Pardo, and M. Pardo, "The jensen-shannon divergence," Journal of the Franklin Institute, vol. 334, no. 2, pp. 307-318, 1997.
12

---

[45] R. Cai, J. Chen, Z. Li, W. Chen, K. Zhang. J. Ye. Z. Li, X. Yang, and Z. Zhang, "Time series domain adaptation via sparse associative structure alignment." in Proceedings of the AAAI Conference Artificial Intelligence, vol. 35, pp. 6859-686) 2021.
[46] Y. Ganin, E. Ustinova, H. Ajakan, P. Germain, H. Larochelle, F. Laviolette, M. Marchand, and V. Lempitsky, "Domain-adversariall training of neural networks," The journal of machine learning research, vol. 17, no. 1, pp. 2096-2030, 2016.
[47] C.-L. Li, W.-C. Chang. Y. Cheng, Y. Yang, and B. Póczos, "Mmd gan: Towards deeper understanding of moment matching network," Advances in neural information processing systems, vol. 30, 2017.
[48] Y. Park, D. Maddix, F.-X. Aubet, K. Kan, J. Gasthaus, and Y. Wang, "Learning quantile functions without quantile crossing for distribution-free time series forecasting." in International Conference on Artificial Intelligence and Statistics, pp. 8127-8150, PMLR, 2022.
[49] G. Loaiza-Ganens and P. Cunningham, "The continuous bernoulli: fixing pervasive error in variational autoencoders," Advances in Neural Information Processing Systems, vol. 32, 2019.
[50] D. Zhou and X.-X. Wei, "Learning identifiable and interpretable latent models of high- dimensional neural activity using pi-vae," Advances in Neural Information Processing Systems, vol. 33, pp- 7234-7247, 2020.
[51] O. Fabius and J. R. Van Amersfoort, "Variational recurrent auto-encoders," arXiv preprint arXiv:1412.6581, 2014.
[52] D. P. Kingma and J. Ba, "Adam: A method for stochastic optimization," arXiv preprint
142.6914 2014.
[53] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov, "Dropout: simple way to prevent neural networks from overfitting," The journal of machine learning research, vol. 15, no. I, pp. 1929-1958, 2014.
[54] A. D. Rasamoclina, F. Adjailia, and P. Sinčák, "A review of activation function for artificial neural network," in 2020 IEEE 18th World Symposium on Applied Machine Intelligence and Informatics (SAMI), pp. 281-286, IEEE, 2020.
[55] D. Kobak, W. Brendel, C. Constantinidis, C. E. Feierstein, A. Kepecs, Z. F. Mainen, X.-L. Qi, R. Romo, N. Uchida, and C. K. Machens, "Demixed principal component analysis of neural population data." Elife, vol. 5, p. e10989, 2016.
[56] D. Sussillo, R. Jozefowicz, L. Abbott, and C. Pandarinath, "Lfads-latent factor analysis via dynamical systems," arXiv preprint arXiv: (0::::608.000 603116 2016
[57] K. Han, Y. Wang, H. Chen, X. Chen, J. Guo, Z. Liu, Y. Tang, A. Xiao, C. Xu, Y. Xu, et al., "A survey on vision transformer," IEEE transactions on pattern analysis and machine intelligence, vol. 45, no. 1, pp. 87-110, 2022.
[58] A. Dosovitskiy, L. Beyer. A. Kolesnikov, D. Weissenbom, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, el al., "An image worth 16x16 words: Transformers for image recognition at scale," arXiv preprini arXiv:2010.1 1929, 2020.
[59] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, "Attention all you need," Advances in neural information processing systems, vol. 30, 2017.
[60] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, el al., "Pytorch: An imperative style, high-performance deep learning library," Advances in neural information processing systems, vol. 32, 2019.
[61] Major and P. Major, Multiple Wiener-Itó integrals. Springer, 1981.
13

---

Appendix to "Extraction and Recovery of Spatio-Temporal Structure in Latent Dynamics Alignment with Diffusion Model"
A Methodology details
A.1 DM architecture details
We adopt the architecture of DM mainly derived from Diffusion Transformer (DiT) [36]. The canonical DiT architecture is based on techniques like patchify [57] and transformer layer for tokens [58]. which are well-suited for image feature extraction. This is because the above techniques focus on local feature extraction and the global feature can also be implicitly captured through the stacking of token-based Transformer layers. However, considering the neural observations and latent dynamics in the format of nulti-variate time series, patchify and local feature extraction loses their semantic meaning. There doesn't t exist a theoretical guarantee or bio-plausible observation that adjacent latent dimensions of the data matrix have stronger connection than far-apart latent dimensions, Thus, directly adopting the traditional DiT architecture into this setting may lead to sub-optimal solutions.
To fully utilize the domain knowledge of our task, we propose to use the Spatio-Temport Transformer Block (STBlock). Each STBlock is mainly composed of a Spatio Transformer layer followed by Temporal Transformer layer, which are 1-layer encoders based on multi- head self-attention. Since there exists underlying dependency and structure between latent state dimensions, the Spatio Transformer layer takes latent states of each time bin inputs to extract their spatial structure. Whereas the Temporal Transformer layer takes the entire latent trajectory of each latent space dimension as inputs to extract its underlying temporal structure. We note that position embeddings [59] to encode the timestep (i.e., noise scale) into the deep neural network of DM. In each STBlock, the input sequentially goes through:
the sinusoidal
Spatio Transformer: Layer Normalization > Multi-head Self-attention Layer (along time bins)
Point-wise Feed-forward.
Layer Normalization Multi-head Self-attention Layer (along latent space dimensions) - Point-wise Feed-foree feed-forward.
Temporal Transformer:
We illustrate the main architecture of DM in Figure 2 A), and implement the DM in Pytorch 660
A.2 VAE and DM cooperative source domain learning details
In DM training, we note that Z,65) here actually latent dynamics inferred via VAE in Eq. Consid- ering the limited number of source-domain latent dynamics, we wish to perform data augmentation so that the DM can adequately estimate p. (Z). Here propose enrich the input samples by learning the VAE objective (Eq. 6) and the diffusion objective (Eq. g) cooperatively. Through the learning process of the VAE objective (i.e., ELBO), the optimization process with stochastic gradient descent (SGD) adds auxiliary perturbation to the original data samples Z rather than pure Gaussian noise.
This technique further fills the sample space of (z(²) Icading to better density estimation. Specifically, in each training iteration, conditioning on the current value of o.., we infer set of Zu = h h((((); Φx) and use as the temporal z
to optimize Eq. The traditionally sequential approach is that we fully
Table 2: The coefficient of determination values (R² 1. in %) and RMSE J of sequential source domain learning and cooperative source domain learning on the primate motor cortex dataset. Boldface denotes the highest score. Each experiment condition repeated with different random seeds, and their mean and standard deviation are listed.
Metric
Sequential
Cooperative
(2(%)
18.96 (=2.27)
20.47 (=2.71)
M1-M2
RMSE
7.76 (+0.40)
7.62 (0.4))
(²(%)
21.73
(1:2.46)
22.62
(+2.6
M2-M3
RMSE
7.94 (+:0.4)
7.73 (+0.50)
R²(%)
6.96
(12.88)
8.57
(12.96)
M2-C2
RMSE
11.85 (+0.42)
11.64
(=0.51)
14

---

(A)
(C)
Secuencil Cooperative
Sequentol
Cooperal
Cooperstive
80
Epocr
Epccs
Epoch
Figure 6: Training loss curves under different neural sessions. (A) Neural session: M-1. (B)
Neural session: M-3. (C) Neural session: C-2.
optimize VAE first, obtain an optimal Φ., and use it to optimize Eq. g Experimental results show that the former approach achieves higher density estimation and alignment performance. Figure manifests the training loss curve in three source-target neural recording session pairs. We observe that despite the relatively under-fitting at the early stage, the cooperative source domain learning paradigm converges to solutions with lower losses and better fits. Table 2 manifests that our cooperative source domain learning paradigm leads to higher distribution alignment and neural decoding performance.
B Detailed derivation of maximum likelihood alignment
B.1 Relationship between KL-Divergence and DSM Loss
Under assumptions in Appendix of [41], the KL divergence between the ground truth density q [z X(0; Φ) and the DM marginal distribution po (Z; 0.) can derived as:
DKL ((( ; (Z: 0.)
(6)
E Bon (100))))))) (A(t) (Vz lo; - s(Z,t; dw + DKL (pr (Z; 0.) [|+(Z))
E (1 A(t)) VV (z; (Z,, t: 0.).) dt
+ 2
(i8)
E 1 R B
A(1)2 | (Z, sZZ,4;0,) 0.) di
DRI.. (pr (Z; 0.) (l=(
(187)
Em
1 17 (000)))) (o)
A(1) IVz log Por (Z, Zo) - 8 (Ze, t; 0)
+ DKL (pr (Z: 0.) (Z))
COSM (Ø. 0.) + DEL (pr Z 0,) (Z))
(15)
in which (i) is due to Girsanov Theorem 40 in (ii) we invoke the martingale property of Itô
integrals 61 and in (iii) we use the denoising score matching (DSM) technique 229 Thus we can
draw to Eq. 11
B.2 Upper bound of maximum likelihood alignment objective
In Section by substituting Eq. into Eq. TO we have
3.3
-E EEzzg(l x() [oogPo 0.)] < LDSM (Φ. 0.) DEL (pr (Z.0.) +H (9(Z x(0; (16)
())
15

---

We note that the third term HI(-) depends on the parameter set & of the probabilistic encoder. As (Z X (0; Φ) x Po ( 0,), we have
H [( x x)) - H (2:0)) ~ l" a H (p: (Z; dt
(17)
0t
E, ((0,7 Ezza 2f(Z, t) Vz lzlogp 0.) X(t)² | log (2; 0,) dt
(1)
2Vz f(2,() + A(1)² Vz log 22 0 di
(80
2Vz, f(z,, t) + 2 V log Pus (Z, Za)ll dt
(18)
where in both (i) and (ii) integration by parts, and in (iii) denoising score matching (DSM) (29]. Putting the second term on the LHS of Eq. into RHS and then substituting the third term the RHS of Eq. I6. have
Ex 2 [log (Z;;0(Z & DEL (0)
(19)
(0, E Z (xx) (())) (Z, Zo)
A(t)2 Vz, log Poi (Z, Z
(20)
+ (0,22(0T77) (o, E. 7. ol x :.) /Za)
((t) Zz log Por (Z. Zo) (Z,
(21)
E Zo ~ (x(1 : 2. Zo)
2Vz (Z,,
(22)
Since the transition probability
[Z, Zo) is a fixed Gaussian distribution and it independent of
the parameter set Φ,, we can eliminate the term in Eq. 20 and rewrite the above Eqs
Ez~q(z|x00; [log Po 0z)l < DKL (pr ZZ 0g) | (Z))
(23)
+ E Bz a( ((( ;0) (2, Zn)
((t) Z2 log Por (Ze Zo) s (Z., t: 0)
(24)
- -Ev-esooTyy E 7 ( x Po: (Z, [z)
2vz f (Zz, ()]
(25)
By substituting denoiser function e [Z, t: 0) into score function sZ,,, t: 0) of Eq. 24 we have:
Ez (z)) [logPo(Z: (Z; 0.)] < DRL. (pT (Z; (.)0[[((())))
+ E (()))
w(t)² ||e <ZZ,, t; 0.) 2vz f (Zo,
C Detailed algorithm
C.1 Overall alignment loss function in ERDifT
Extracting the latter terms from Eq. 12 have the following main Maximum Likelihood Alignment (MLA) loss:
LMLA (Ø) E2-ejo (7000
(o(1)3ll ||e E <(Z,, 1;0,11 -V2, ff (Z,, 1)
(26)
Then, considering the Sinkhorn Regularizer term in Eq. 13 LSHD( min (r. C)F + AHl(y),
(27)
where each value C| = ( z in cost matrix C' denotes the squared Fuclidean cost
from z(d) to ((4) H()) computes the entropy of transport plan 7, and À refers to the weight of the entropy term. Then, we can find the optimal Φ. via the minimization the following total loss function:
Φi = argmin - (M + aLsHD),
(28)
where E [0, 1] is a trade-off parameter that weights the importance of Sinkhom Regularizer term.
16

---

C.2
Algorithm for source domain learning of ERDiff
Algorithm Source Domain Learning of ERDiff
Input: Source-domain neural observations X ( Learning rate 7j and all other hyperparameters.
Output: Parameter set (5(*) of source-domain probabilistic encoder; Parameter z/(*) of source- domain probabilistic decoder; Parameter set ((*) of source-domain diffusion model.
1: Initialize Φ sb, and 0:
2: while not converge do
3: LELBO < Compute the loss function based on evidence lower bound according to Eq
4: Parameter Update: & Φ "i OLELBO/Od.
5: Parameter Update: ab + 3/-
/CELBo
Inference Z)6') q(- x()); Φ), and Noise Sampling N' (0, Ixxa);
6:
7: LDSM < Compute the denoising score matching loss according to Eq.
8: Parameter Update: 0 + n ÖCDsM/00;
9: end while
return Φ, v/ and 0.
C.3 Algorithm for maximum likelihood alignment ERDif
Algorithm Maximum Likelihood Alignment of ERDiff
Input: Target-domain neural observations ( Learning rate y and all other hyperparameters.
Output: Parameter set (t()) of target-domain probabilistic encoder;
1: while not converge do
2: Inference Z1(¹
q(( (4; <b), and Noise Sampling ~ N (0, Tixa);
3: LMLA +- Compute the main Maximum Likelihood Alignment Loss according to Eq.
4: CSHD +- Compute the Sinkhorn Regularizer according to Eq. 27
5: Parameter Update: Φ f Φ 1) [[1 aMLL + aLSHD) /0
6: end while
7: return Φ
D Experimental results on rat hippocampal CA1 dataset
We further verify the effectiveness of ERDiff on publicly available rat hippocampus dataset (2]] In this study, rat navigated 1.6m linear track, receiving rewards both terminals (L&R) depicted in Figure 4 (A). Concurrently, neural activity from the hippocampal CAI region was captured, in which the neuron numbers across all recorded sessions
around 120. The neural spiking activities were binned into 25ms intervals. A single round trip by the rat from one end of the track to the opposite was categorized as one lap. We chose 5 sessions in total, and in each session, 80 laps were sampled. The 2 2-dimensional neural latent manifolds can be observed in Figure 7 (B) and (C). We also list the neural decoding results of the rat's position Table
(A)
Source
ERDiff (Ours)
DAF
JSDM
Z1
R²(%) = 34.64
R²(%) = 20.64
R²(%) = 2.58
Figure 7: Rat hippocampus dataset and experimental results. (A) Illustration of the experiment setting. (B) 2D visualization of inferred latents in the source domain, in which the bold lines manifest the mean location of latent states corresponding to the rat's position for two opposite directions. (C) On target domain, 2D visualization of inferred latents after alignment by ERDiff, DAF, and JSDM, respectively. We observe that ERDiff preserves the spatio-temporal structure of latent dynamics well.
17

---

Table Decoding results after neural distribution alignment on the rat hippocampus dataset across sessions. Each experiment condition is repeated with 5 different random seeds, and their mean and standard deviation are listed.
Method
R²
RMSE
Cycle-GAN
4.81
(+:2.47)
7.65
(+00344)
JSDM
-2.88
(=1.59)
7.65
(+0.23)
SASA
18.51
(+1.65)
6.98 (+0.20)
DANN
15.04
(+2.32)
7.30 (#0266)
RDA-MMD
21.73
(#2.45)
7.22
(+0.26)
DAF
20.20
(+2.26)
7.36
(+0.28)
ERDiff
32.69
(-2.19)
6.84 (++0.66)
Table 4: Comparative analyses of computational cost between ERDiff and baseline methods during
alignment. ERDiff has a comparable computational cost and maintains the stability of alignment.
Method
Cycle-GAN
JSDM
SASA
RDA-MMD
DAF
ERDiff
Add'l. Param
26K
OK
33K
65K
91K
28K
Add'l. Size
117KB
OKB
187KB
314KB
367KB
139KB
Align. Time
103ms
77ms
155ms
264ms
251
1 183ms
Stability
X
X
X
E
Computational cost analysis
Here we conduct time complexity analysis with respect to the batch size B for the alignment phase. The ERDiff's alignment objective function is composed of two main terms: Diffusion Noise Residual and Sinkhorn Divergence. We note that in the diffusion noise residual computation, it does not go through the entire T diffusion steps. Instead, it just samples a single time step (noise scale) and calculates the noise residual specific to that step. Thus, the total complexity of this part takes O(K} * B * d), in which the coefficient K₃ relates to the inference complexity of the DM denoiser (Z, t); d denotes the latent dimension size. For the Sinkhorn Divergence, has to compute the distance matrix, costing O(K₂ * B²); K2 is a relatively small coefficient in magnitude. By summing up, the total complexity of ERDiff is given O(K, B d + K₂ B²). This O (²2) complexity is applicable since the non-adversarial baseline methods we compared (i.e., JSDM, and SASA) require quadratic complexities well.
For any given target domain, ERDiff can stably align it to the source domain a comparable overhead with baselines. In Table 4 conduct comparative analysis between ERDiff and baseline methods in terms of additional parameter number, additional model size, stability, and alignment time. The demonstrated alignment time corresponds to the execution time for aligning one iteration (a batch of size 64) on a MacBook Pro (2019 equipped with 8-Core Intel Core and 4 GB RAM).
18


Reviews:
[{"review_id": "qWUBa5QGvC", "reviewer": "Reviewer_9MfT", "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.\nSimulation 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.", "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.", "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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "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.\nThe 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. \nIn 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). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and constructive comments. We would like to make the following clarifications. Hopefully these will resolve most of your concerns, and they can be taken into account when deciding the final review score.\n\n> **(1)** It is not clear to what extent the presented findings might generalize to data sources without such a clear trial structure. **(2)** 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. \n\n**(1)** This is a highly valid point. We would like to emphasize that the target of performing alignment between neural distributions is to adapt the neural-behavior mapping function from the source domain to the target domain. The consistency of such mapping function within every single domain is a prerequisite. Therefore, our paper and previous works [1, 2] in neural distribution alignment mainly focus on datasets where behaviors have been confirmed to maintain a consistent trial-wise dynamical structure. \n\nOn the other hand, when we aim to apply neural distribution alignment (i.e., ERDiff) to datasets without a clear trial structure [3,4], it is feasible to truncate such continuous behavioral data into discrete but meaningful segments and syllables (e.g., grooming, running). These behavioral segments typically exhibit clear structural patterns. Then, their overall latent distribution can be learnt through a VAE for subsequent alignment. We conduct additional experiments to validate this approach through an unconstrained rat’s hippocampus dataset [5]. Please refer to Section 1.1 of the general response for details. \n\n**(2)** Please refer to Section 1.2 of the global response for a detailed discussion of this point.\n\n> It would be informative to see a direct comparison with canonical correlation analysis (CCA) approaches\n\nWe notice that CCA is a well-known method for neural distribution alignment. However, practical approaches based on CCA [6] require supervised behavior signals of the target domain to pair the source-target domain trials during the alignment phase. In contrast, our proposed ERDiff is an **unsupervised** neural distribution alignment method, having the potential to generalize to broader datasets and applications where supervised behavior signals are inaccessible. Here, for a fair comparison, we remove the behavior signals in practical CCA approach [6] and report the results in the following:\n\n|  Method  |       Cross-Day-$R^2$       |       Cross-Day-RMSE       |     Inter-Subject-$R^2$     |     Inter-Subject-RMSE     |\n| :------: | :-------------------------: | :------------------------: | :-------------------------: | :------------------------: |\n|  U-CCA   |     $-25.56 (\\pm 2.13)$     |     $9.70 (\\pm 0.28)$      |     $-29.26 (\\pm 2.11)$     |     $10.43 (\\pm 0.46)$     |\n| **Ours** | $\\mathbf{18.81} (\\pm 2.24)$ | $\\mathbf{7.99} (\\pm 0.43)$ | $\\mathbf{10.29} (\\pm 2.86)$ | $\\mathbf{8.34} (\\pm 0.34)$ |\n\n*Table 1: Performance comparison with Unsupervised-CCA*\n\n> **(1)** 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 ? **(2)** 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. \n\n**(1)** This is a good point. While Latent Variable Models (LVM) such as LFADS are focused on identifying low-dimensional, latent factors from raw neural signals, the primary goal of ERDiff is quite different. ERDiff aims to perform unsupervised neural distribution alignment (domain adaptation) given these identified latent factors in the source and target domains. On the other hand, not limited to the current canonical VAE, ERDiff has the potential to perform alignment based on latent factors identified by alternative LVMs (e.g., LFADS). We thus include this topic as part of our future work (please refer to Section 1.4 of the global response). \n\n**(2)** We apologize for the confusion. In ERDiff, we would like to clarify that the learning of low-dimensional latents is **separate** from the alignment. In the outlined cooperative training procedure, during each iteration, the DM uses the inferred latents by the VAE as input data. However, meanwhile, DM does not affect the VAE's learning. This means the procedure of learning the low-dimensional latents is disentangled and separate from the procedure of DM learning and neural distribution alignment. \n\nWe note that the cooperative training of VAE and diffusion model (DM) assists the DM in accurately extracting the distribution of the source domain latents, which ultimately improves the alignment performance. This is why we include a description of the VAE training (learning the low-dimensional latent) phase in the source domain.\n\n> I do not see a Broader Impact section included in the current submission.\n\nPlease refer to Section 1.4 of the global response.\n\n\n\nWe look forward to further discussion, and are happy to answer any questions that might arise.\n\nRefs:\n\n[1] Robust alignment of cross-session recordings of neural population activity by behaviour via unsupervised domain adaptation. (Justin et al., 2022)\n\n[2] Stabilizing brain-computer interfaces through alignment of latent dynamics. (Brianna et al. 2022)\n\n[3] Hierarchical recurrent state space models reveal discrete and continuous dynamics of neural activity. (Scott et al., 2019)\n\n[4] The Striatum Organizes 3D Behavior via Moment-to-Moment Action Selection. (Jeffrey et al., 2018)\n\n[5] Recordings from hippocampal area ca1, pre, during and post novel spatial learning. (Grosmark et al.,2016)\n\n[6] Long-term stability of cortical population dynamics underlying consistent behavior. (Gallego et al., 2020)"}, {"review_id": "dhsdPPMFjj", "reviewer": "Reviewer_tfQ5", "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.\n", "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.\n2. Is the data going to be released publicly for reproducibility?\n", "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.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. A well written paper (but the abstract).\n2. An interesting application of diffusion models.\n3. Potentially impactful in practice of the BCI, more work, including further evaluation, is needed here however.\n", "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.\n2. Poorly written abstract, especially in contrast to the rest of the paper.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments. We would like to make the following clarifications. Hopefully these will resolve most of your concerns, and they can be taken into account when deciding the final review score.\n\n> A cellular neuroscience focused paper with no additional effort to demonstrate a general applicability of the approach in evaluations. 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). \n\nWe would like to emphasize that the focus on this work lies in behavioral neuroscience research and its related applications (e.g., brain-computer interfaces (BCI)) and we have selected \"Neuroscience and cognitive science\" as the primary area for this paper. Within the field of neuroscience, we note that the focus of this paper (i.e., neural distribution alignment) is of vital importance. Therefore, we mainly validate the effectiveness of the proposed ERDiff on various neural recording datasets. \n\nOn the other hand, we fully agree that the generalizability of each method is an important consideration. Therefore, we have conducted experiments on synthetic datasets (Section 4.1 of the manuscript) in which more general time-series data are simulated. Moreover, in Section 1.1 of the general response, we also provide additional experimental results on an unconstrained rat’s hippocampus dataset [1]. These results prove the performance enhancement by ERDiff in a more general context and will be provided in our revised manuscript.\n\nHere we further discuss the contribution of ERDiff to the broader Machine Learning field: \n\n**(1)** In the field of domain adaptation, we propose to use diffusion model (DM) and STBlocks (Fig.2(A) of the manuscript) to extract spatio-temporal structures of the source domain distribution for later alignment propose. We note that the preservation of spatio-temporal structures is a key component for accurate adaptation, and such structures are ubiquitous in dynamical time-series data outside the neuroscience field [2, 3]. \n\n**(2)** We propose the maximum likelihood alignment (MLA). The robust statistical properties of MLA ensure stability during the alignment phase and offer broad applicability across a range of time-series datasets. We also note that the optimization objective of MLA aligns well with the DM since the $p_s(\\mathbf{Z})$ inside MLA is tractable through DM, and the entire MLA formula can be expanded into noise residuals terms that are simple and efficient to compute. \n\nHence, we believe ERDiff has potential use in underlying structure extraction and domain adaptation tasks of general time-series data (e.g., weather forecasting [2] and seismology [3]). We appreciate that Reviewer 9MfT, wx95, and 1oNu have recognized the potential generalizability of ERDiff. In the revised manuscript, we plan to incorporate the aforementioned insights into the broader impact part of ERDiff (please refer to Section 1.4 of the general response).\n\n> Poorly written abstract, especially in contrast to the rest of the paper.\n\nWe thank the reviewer for this practical suggestion. The following is our re-written abstract and we will update it in our revised version of the manuscript.\n\n\"In the field of behavior-related brain computation, it is necessary to align raw neural signals against the drastic domain shift among them. A foundational framework within neuroscience research assumes that trial-based neural activities rely on low-dimensional latent dynamics. Focusing on such latent dynamics greatly assists the alignment procedure. Despite the great progress the field have reached, existing methods usually ignore the intrinsic spatio-temporal structures during alignment. Thus, those solutions lead to poor quality in dynamics structures and overall performance after alignment. To tackle this problem, we propose a method leveraging the expressivity of diffusion model to relieve such issues. Specifically, the latent dynamics structures of the source domain are first extracted by the diffusion model. Then, such structures are well-recovered through a maximum likelihood alignment procedure in the target domain. We first demonstrate the effectiveness of our proposed method on a synthetic dataset. Then, when applied to neural recordings from primate motor cortex, under both cross-day and inter-subject settings, our method consistently manifests its capability of preserving the spatio-temporal structure of latent dynamics and outperforms existing approaches in alignment quality.\"\n\n> 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.\n\nYes. We would release the codes in the supplementary material to public. The detailed implementations of ERDiff on all datasets are described in section A.1 and D.1 of the appendix. To further ensure the reproducibility of our experimental results, we will provide a more comprehensive listing of these implementation details in the appendix of our revised manuscript.\n\n> Is the data going to be released publicly for reproducibility?\n\nYes. The dataset is collected by [4] and it is available from the corresponding author upon reasonable request.\n\n\n\nWe look forward to further discussion, and are happy to answer any questions that might arise.\n\n\n\nRefs:\n\n[1] Recordings from hippocampal area ca1, pre, during and post novel spatial learning. (Grosmark et al.,2016).\n\n[2] Application of Domain Adaptation Approach for Weather Data Mining. (Yang et al., 2018)\n\n[3] Seismic Facies Analysis: A Deep Domain Adaptation Approach. (Quamer et al., 2021)\n\n[4] Long-term stability of cortical population dynamics underlying consistent behavior. (Lee et al., 2020)"}, {"review_id": "ma4ehMKNGA", "reviewer": "Reviewer_wx95", "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.\n", "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.\n", "limitations": "- Limitations is not addressed in the draft", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The proposed method is novel and is technically sold\n- 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. \n", "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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and insightful questions. We provide clarifications and new results that we have generated to address your questions below.\n\n> Although the methods have been shown to outperform some of the current techniques, advantage of the proposed approach is not well demonstrated.\n\nWe appreciate your attention to this important aspect. We'd like to emphasize that the advantages of our proposed ERDiff in neural distribution alignment comes from the following two key points.\n\n**(1)** Precise source domain learning: We note that in behavior-related neural applications (e.g., brain-computer interfaces), the trial latent dynamics are non-linear and complex. This complexity pose challenges when it comes to aligning neural distributions across sessions (domains). Owing to the strong **expressivity** of diffusion models (DM), we first use a DM to learn the overall latent distribution $p_s(\\mathbf{Z})$ in source domain. To learn $p_s(\\mathbf{Z})$ well, the DM focuses on the spatio-temporal structures of neural latent dynamics in source domain and extracts these structures using the specially designed STBlock. Fig. 2A of the manuscript visualizes the distribution learning and spatio-temporal structure extraction process of the DM in source domain. However, previous neural alignment methods [1, 2] ignore these crucial spatio-temporal structures.\n\n**(2)** Appropriate alignment objective: During the alignment procedure, we propose to use maximum likelihood alignment (MLA) as the optimization objective. We note that this objective aligns well with the DM since the $p_s(\\mathbf{Z})$ inside MLA is tractable through DM, and the entire MLA formula can be expanded into noise residuals terms that are simple and efficient to compute. Thanks to the stability and flexibility of MLA, the source-domain spatio-temporal structures can be well-recovered in target domain, as illustrated in Fig. 2B and Fig. 4 of the manuscript. In contrast, methods based on pre-defined metrics often make a strong assumption that the overall latent distributions follow a Gaussian, which significantly restricts the expressiveness of those methods. For methods based on adversarial training, even though their optimization objectives can be formed into Jensen–Shannon divergence (JSD), their practical training steps always lack stability [3], and are far from reaching the optimization target (JSD).\n\n> I would like to see some analysis on computational cost.\n\nIn source domain training phase, the additional computation cost of ERDiff primarily comes from the diffusion model (DM) training. We note that the DM is trained in the latent space $\\mathbf{Z}$, which is significantly smaller than the raw neural signal space $\\mathbf{X}$ in magnitude. Therefore, the computational overhead of this one-time training phase is acceptable, especially given that it can be conducted offline in real-world BCI applications. \n\nIn the alignment phase, we would like to emphasize that ERDiff maintains a comparable computational cost with baseline methods. Please refer to Section 1.3 of the global response for a comprehensive analysis of the computational cost and time complexity.\n\n> 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.\n\nIn the experiment section of the manuscript, among the group of methods based on pre-defined metrics, we show the results of JSDM since it achieves the highest performance. To provide a more thorough comparison, here we additionally conduct experiments with two more methods in this group: Kullback–Leibler divergence minimization (KLDM) and Wasserstein distance minimization (WDM). These two metrics form the backbone of previous methods on neural distribution alignment [4] and [5], respectively. We can observe that ERDiff consistently demonstrates superior performance compared to these two methods. These additional comparative results will be incorporated into the revised manuscript.\n\n|  Method  |       Cross-Day-$R^2$       |       Cross-Day-RMSE       |     Inter-Subject-$R^2$     |     Inter-Subject-RMSE     |\n| :------: | :-------------------------: | :------------------------: | :-------------------------: | :------------------------: |\n|   KLDM   |     $-31.63 (\\pm 1.85)$     |     $10.55 (\\pm 0.37)$     |     $-30.79 (\\pm 2.24)$     |     $10.42 (\\pm 0.37)$     |\n|   WDM    |     $-19.85 (\\pm 2.45)$     |     $8.61 (\\pm 0.36)$      |     $-17.74 (\\pm 2.71)$     |     $9.21 (\\pm 0.43)$      |\n| **Ours** | $\\mathbf{18.81} (\\pm 2.24)$ | $\\mathbf{7.99} (\\pm 0.43)$ | $\\mathbf{10.29} (\\pm 2.86)$ | $\\mathbf{8.34} (\\pm 0.34)$ |\n\n*Table 1: Performance comparison with alignment methods based on pre-defined metrics*\n\n> Limitations is not addressed in the draft\n\nWe delve into the Limitations, Future Work, and Broader Impact of our work. (Please refer to Section 1.4 of the global response for details.) This comprehensive discussion will be put into the revised manuscript in a new 'Discussion' section, including and replacing the existing 'Section 5: Conclusion'.\n\n\n\nWe look forward to further discussion, and are happy to answer any questions that might arise.\n\nRefs:\n\n[1] Robust alignment of cross-session recordings of neural population activity. (Justin et al., 2022)\n\n[2] Stabilizing brain-computer interfaces through alignment of latent dynamics. (Brianna et al. 2022)\n\n[3] Evaluation of Mode Collapse in Generative Adversarial Networks. (Sayeri et al., 2018)\n\n[4] Stabilizing brain-computer interfaces through alignment of latent dynamics. (Brianna et al. 2022)\n\n[5] Hierarchical Optimal Transport for Multimodal Distribution Alignment. (John et al. 2019)"}, {"review_id": "jAjPFRFkzg", "reviewer": "Reviewer_2uTS", "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.\n• 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.\n• 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.\n• The authors applied their model to the non-human primate motor cortex, testing both cross-day and inter-subject recordings.", "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.\n* Apriori, why would one expect the diffusion model to be more effective at aligning the latents than these other approaches?\n* 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.", "limitations": "* No significant negative societal impact envisioned -- the potential use in BCI may suggest a positive social impact.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The authors motivate their approach clearly by arguing that naively aligning time series using domain adaptation is ineffective as\nmultivariate neural data has low SNR. Thus, leveraging low-dimension representation is a viable option.\n• 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).\n• The model and derivations are presented clearly.\n• 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?\n* Alternatively how about comparisons with alternative approaches of comparable complexity -- e.g. adversarial alignment with DANN (Ganin et al, 2015)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and insightful questions. We provide clarifications and new results that we have generated to address your questions below. Hopefully these will resolve most of your concerns, and they can be taken into account when deciding the final review score.\n\n> The authors must clarify why diffusion models are necessary. How about considerably simpler two-step approaches -- like extracting latents with GPFA and aligning them with the proposed ML alignment?\n\nThis is an insightful question. We would like to emphasize that the diffusion model (DM) is an **essential** and **necessary** component in ERDiff. The traditional methods for identifying low-dimensional latents (e.g., GPFA, LFADS) are **not applicable** to the proposed maximum likelihood alignment (MLA) phase alone. This is because their models only learn a point-to-point mapping function $p(\\mathbf{Z} \\mid \\mathbf{X}^{(s)})$ from neural spike signal to latent factor, rather than explicitly capturing the overall latent distribution $p_s(\\mathbf{Z})$ in the source domain. Consequently, in the second step (alignment phase), given the MLA optimization objective: $\\underset{\\phi}{\\operatorname{argmax}} \\mathbb{E}_{\\mathbf{Z} \\sim q\\left(\\mathbf{Z} \\mid \\mathbf{X}^{(t)} ; \\phi\\right)}\\left[\\log p_s(\\mathbf{Z})\\right]$ (RHS of Eq.7 in the manuscript), the likelihood $p_s(\\mathbf{Z})$ is intractable through those methods alone. In contrast, in our proposed ERDiff, along with the learning of $p(\\mathbf{Z} \\mid \\mathbf{X}^{(s)})$ through VAE, we use DM to learn the overall latent distribution $p_s(\\mathbf{Z})$. Therefore, through DM, we can reformulate the MLA objective into Eq. (10) of the manuscript, which is tractable, and computationally efficient.\n\n> Alternatively how about comparisons with alternative approaches of comparable complexity -- e.g. adversarial alignment with DANN (Ganin et al, 2015)?\n\nIn the experiment section, the methods SASA, RDF-MMD, and DAF we compared are all of comparable alignment complexity with ours. For a detailed comparison of time complexity, please refer to Section 1.3 of the general response. Here we conducted further experiments (Table 1) using the Domain-Adversarial Neural Network (DANN) method employing 5 random seeds. We implemented the label predictor of DANN with the behavioral velocity predictor. The results of these experiments will be included in the revised manuscript.\n\n|  Method  |       Cross-Day-$R^2$       |       Cross-Day-RMSE       |     Inter-Subject-$R^2$     |     Inter-Subject-RMSE     |\n| :------: | :-------------------------: | :------------------------: | :-------------------------: | :------------------------: |\n|   DANN   |     $-12.57 (\\pm 3.28)$     |     $8.28 (\\pm 0.32)$      |     $-18.37 (\\pm 3.24)$     |     $9.29 (\\pm 0.33)$      |\n| **Ours** | $\\mathbf{18.81} (\\pm 2.24)$ | $\\mathbf{7.99} (\\pm 0.43)$ | $\\mathbf{10.29} (\\pm 2.86)$ | $\\mathbf{8.34} (\\pm 0.34)$ |\n\n*Table 1: Performance Comparison with DANN*\n\n> 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.\n\nAs we've explained in the first clarification of this response, approaches such as LFADS or GPFA that you mentioned are not applicable for the second MLA stage alone. A diffusion model is necessary to learn an overall $p_s(\\mathbf{Z})$ for MLA.\n\n> Apriori, why would one expect the diffusion model to be more effective at aligning the latents than these other approaches?\n\nWe appreciate your attention to this important aspect. We would like to clarify that the methods you listed do not serve as neural distribution alignment methods but rather latent variable models (LVM) that are designed to identify latents from raw neural signals.\n\nBesides the learning of an overall $p_s(\\mathbf{Z})$ for MLA, we note that the effectiveness of diffusion model (DM) comes from the following two key points.\n\n(1) Precise source domain learning: We note that in behavior-related neural applications, the trial latent dynamics are non-linear and complex. This complexity pose challenges for aligning neural distributions across sessions. Owing to the strong **expressivity** of diffusion models (DM), we first use a DM to learn the overall latent distribution $p_s(\\mathbf{Z})$ in source domain. To learn $p_s(\\mathbf{Z})$ well, the DM focuses on the spatio-temporal structures of neural latent dynamics in source domain and extracts these structures using the specially designed STBlock. Fig. 2A of the manuscript visualizes the distribution learning and spatio-temporal structure extraction process of the DM. However, previous neural alignment methods ignore these crucial spatio-temporal structures.\n\n(2) Appropriate alignment objective: During the alignment procedure, we propose to use maximum likelihood alignment (MLA) as the optimization objective. We note that this objective aligns well with the DM since the $p_s(\\mathbf{Z})$ inside MLA is tractable through DM, and the total MLA formula can be expanded into noise residuals terms that are simple and efficient to compute. Thanks to the stability and flexibility of MLA, the source-domain spatio-temporal structures can be well-recovered in target domain, as illustrated in Fig. 2B and Fig. 4 of the manuscript. \n\n> 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.\n\nThis is an analysis of existing datasets [1] and IRB approvals have been obtained in their studies.\n\n\n\nWe look forward to further discussion, and are happy to answer any questions that might arise.\n\nRefs:\n\n[1] Long-term stability of cortical population dynamics underlying consistent behavior. (Lee et al., 2020)"}, {"review_id": "J52MD35WTm", "reviewer": "Reviewer_xrSS", "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.", "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.\nReferences 9 and 10 cite the preprint and peer-reviewed versions of the same article.\n", "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. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "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.\nDeep 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.\nThe 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.\nThe 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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and constructive comments. We would like to make the following clarifications. Hopefully these will resolve most of your concerns, and they can be taken into account when deciding the final review score.\n\n> They overlooked other methods based on the alignment of latent dynamics that have been shown to be successful for BCI applications. E.g., CCA, multiset CCA, hyperalignment, or Procrustes, ..., another literature is  (amLDS [1]).\n\nWe appreciate your suggestion for this comparison. We would like to emphasize that our proposed method, ERDiff, is an **unsupervised** neural distribution alignment method, meaning **no** supervision signals or labels related to behavior are required during the alignment phase. In contrast, the listed methods are **supervised** learning methods. In linear algebra-based methods: CCA and multiset CCA, hyperalignment, and Procrustes-based alignment, they require the knowledge of supervised behavior signals for every trial in both domains to pair the source-target domain trials during the alignment phase. Therefore, these methods are heavily reliant on supervision. Thus, a direct comparison between ERDiff and CCA-based methods may not be a fair evaluation. Regarding amLDS, its probabilistic framework focuses on $K$ discrete stimulus conditions, making it impractical to smoothly expand to our continuous behavioral setting. Therefore, it's also not a direct counterpart to ERDiff.\n\nIn practical brain-computer interface (BCI) and neural behavioral applications, unsupervised neural distribution alignment methods are highly desirable. Our proposed ERDiff has the potential to generalize to many more scenarios and fields where supervised signals are not accessible. \n\n> 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. \n\n(1) In source domain training phase, the computation cost of ERDiff primarily comes from the diffusion model (DM) training. We note that the DM is trained in the latent space $\\mathbf{Z}$, which is significantly smaller than the raw neural signal space $\\mathbf{X}$. Therefore, the computational overhead of this one-time training phase is acceptable, especially given that it can be conducted offline in real-world BCI applications. \n\nAs for the alignment phase, the computational cost of ERDiff is comparable with those of the baseline methods. More detailed analyses and comparisons of computational costs are provided in Section 1.3 of our global response.\n\n(2) We emphasize that training the diffusion model (DM) used in ERDiff does not require any additional data compared to the baseline methods. The DM uses the same dataset as the baselines for source domain learning as well as distribution alignment.\n\n(3) The hyperparameter selection of our model, including the number of epochs, learning rates, and other relevant parameters, are listed in Section D.1 of our submitted appendix. \n\n> 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.\n\nWe infer the terms 'latent space' and 'deep learning-based' you mentioned might be referring to 'methods specifically designed for neural distribution alignment tasks', and to 'general time-series domain adaptation methods', respectively. Besides their performance comparisons in Table 1 of the manuscript, we additionally conduct comparisons on their computational cost and time complexity, please refer to Section 1.3 of the global response for details.\n\n> The authors minimally showed the effect of dataset size on performance, but the lowest dataset size tested still had dozens of trials \n\nWe infer the term 'dataset size tested' you mentioned might be referring to the number of trials used for alignment in the target domain. We would like to highlight that in Fig. 5(B) of the manuscript, we plot the performance of ERDiff and baselines given different numbers of alignment trials. From the plot, it can also be observed that ERDiff maintains a relatively high accuracy when only 25~ trials (representing a 10% sampling density) are used during alignment. \n\n> it would be important to report the training times as a function of dataset size.\n\nPlease refer to section 1.3 of the global response for training time cost functions.\n\n> It would be interesting to note if the proposed method also allows for multi-session alignment\n\nMulti-session alignment belongs to the field of multi-domain adaptation (MDA), which is a more complex problem. However, we believe integrating additional components [2] for MDA, ERDiff can generalize to multi-session alignment. We include this direction as part of our future work, listed in section 1.4 of the global response.\n\n> 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.\n\nPlease refer to Fig. 2 in the attached PDF for our robustness study investigating the impacts of latent dimensionality and Gaussian noise on ERDiff.\n\n> it is not clear whether the authors include behavioral data for alignment. This raises the question of whether there is a circular evaluation of the method.\n\nWe would like to clarify that no behavioral data or velocity information from the target domain is used in ERDiff. The alignment procedure of ERDiff is performed in an unsupervised manner, showing the broad applicability of our method.\n\n> The authors should include a section or provide a clear description of the limitations and assumptions of the method\n\nPlease refer to section 1.4 of the global response for limitations and assumptions.\n\n\n\nRefs:\n\n[1] Across-animal odor decoding by probabilistic manifold alignment. (Pedro et al. 2021)\n\n[2] Multi-Source Unsupervised Domain Adaptation via Pseudo Target Domain.  (Ren et al. 2022)"}, {"review_id": "H5PZPfywyX", "reviewer": "Reviewer_1oNu", "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.", "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?).\n\nFigure 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.\n\nTable 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?\n\nIf 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?\n\nAs the dimensionality of the latent structure increases, how much more badly does the method perform? The authors should quantify this.", "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.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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). \n\nQuality: 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.\n\nClarity: The authors clearly state the problem and its details, as well as how their method differs from existing ones, and its advantages.\n\nSignificance: 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)\n\nOriginality: 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.\n\nQuality: 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and insightful questions. We provide clarifications and new results to address your questions below.\n\n> Authors should state whether existing methods for discovering latent structures using diffusional models exist.\n\nTo the best of our knowledge, no existing methods use diffusion models (DM) for the discovery of neural latent structures. \n\n> The generalizability to other types of neural dynamics time series of varying dimensionality should be evaluated to determine limitations.\n\nThe generalizability of each method is a vital consideration. The monkey reaching dataset has a neural signal dimension of around 180 and an average trial length of around 40. Additionally, we conducted experiments on an unconstrained rat’s hippocampus dataset [1], where the neural signal dimension is 120 and the average trial length is around 110. The corresponding results are presented in Section 1.1 of the general response.\n\n> 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.\n\nWe appreciate your suggestion about Fig. 3(B). In this figure, the phase portraits represent the trajectories and structures of latent dynamics. We note that, after alignment, ERDiff's phase portrait is noticeably closer to that of the source domain in comparison to JSDM. The pink dots you pointed out represent the stable fixed points within the phase portraits, which correspond to steady positions or equilibria of the dynamics system. We will provide a more detailed legend in the revised manuscript.\n\n> 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?\n\nWe would like to clarify that $R^2$ refers to the *coefficient of determination* that quantifies the goodness-of-fit for regression models. Formally, it is calculated as follows: $R^2(y, f) = \\sum_{i=1}^{n}(y_i - \\bar{y})(f_i - \\bar{f})/\\sum_{i=1}^{n}(y_i - \\bar{y})^2$. The denominator is the total variance of ground-truth values.  We will incorporate the formula of $R^2$ calculation into the revised manuscript.\n\n>  **(1)** If the time-series do not belong on a well-defined low-dimensional manifold, will the method hallucinate something? **(2)** How well does the method work for more unstructured data? \n\n**(1)** We would like to emphasize that ERDiff isn't constrained to datasets containing a strong 'well-defined' low-dimensional manifold, such as the commonly depicted sphere-like or Swiss-roll-like manifolds. Instead, with the strong expressivity of DM, ERDiff can broadly fit distribution of datasets containing temporal dynamical structures. Meanwhile, it's supported by numerous neural studies [2, 3] that data from various brain regions inherently exhibit temporal dynamical structures. This manifests the wide applicability of ERDiff. \n\n**(2)** If the spatio-temporal structures in data are weak and the time-series resembles a random walk, the alignment task might not benefit as much from ERDiff, which is designed for more structured data. On the other hand, when we aim to apply ERDiff to datasets without a clear trial-structure [4], it is feasible to truncate such continuous data into discrete but meaningful syllables. These syllables typically exhibit clear structural patterns. Then, their overall latent distribution can be learnt through a VAE for subsequent alignment. We conduct additional experiments to validate this approach through an unconstrained rat’s hippocampus dataset [1]. Please refer to Section 1.1 of the general response for result details.\n\n> As the dimensionality of the latent structure increases, how much more badly does the method perform? The authors should quantify this.\n\nIn real-world datasets, the dimensionality of the latent structure is inherently fixed. Hence, we increase the dimensionality of latent structure ($\\mathbf{z}$) in synthetic datasets and conduct experiments in the following. From the results, we observe that the aligned distribution tends to deviate further from the source-domain distribution as the dimensionality increases. Despite this trend, ERDiff consistently demonstrates superior performance compared to the best baseline method.\n\n|    Method     |      NLL $\\downarrow$      |      KLD $\\downarrow$      |\n| :-----------: | :------------------------: | :------------------------: |\n| Best-Baseline |     $3.58 (\\pm 0.20)$      |     $7.74 (\\pm 0.43)$      |\n|   **Ours**    | $\\mathbf{3.38} (\\pm 0.16)$ | $\\mathbf{7.09} (\\pm 0.38)$ |\n\n*Table 1: Latent structure dimension $d_z$ = 4*.\n\n|    Method     |      NLL $\\downarrow$      |      KLD $\\downarrow$       |\n| :-----------: | :------------------------: | :-------------------------: |\n| Best-Baseline |     $6.15 (\\pm 0.21)$      |     $12.33 (\\pm 0.43)$      |\n|   **Ours**    | $\\mathbf{5.53} (\\pm 0.29)$ | $\\mathbf{11.39} (\\pm 0.47)$ |\n\n*Table 2: Latent structure dimension $d_z$ = 8*.\n\n> The authors should address in more detail the kinds of latent structures/time series that would pose more of an issue for their method. \n\n(1) For shorter time-series that may not exhibit strong latent structures, the alignment task becomes less complex. In such cases, even simpler methods like JSDM can have satisfactory results.\n\n(2) If the spatio-temporal structure is weak and the time-series resembles a random walk, the alignment task might not benefit as much from ERDiff, which is designed for more structured data.\n\n\n\nRefs:\n\n[1] Recordings from hippocampal area ca1, pre, during and post novel spatial learning. (Grosmark et al. 2016)\n\n[2] Context-dependent computation by recurrent dynamics in prefrontal cortex. (Valerio et al., 2014)\n\n[3] Network dynamics underlying OFF responses in the auditory cortex. (Giulio et al., 2016)\n\n[4] Hierarchical recurrent state space models reveal discrete and continuous dynamics of neural activity. (Scott et al., 2019)"}]

Author rebuttal:
We would like to express our sincere gratitude to all the reviewers for their insightful feedback and suggestions. We appreciate the positive comments which characterized our work as having a `"clear motivation"` (2uTS), `"novel"` methodological progression (9MfT, wx95), being `"technically solid"` (wx95, xrSS), with a `"well-written and clear presentation"` (tfQ5, 2uTS), conducting `"exhaustive comparison"` in the experiment (2uTS) and recognized our work to have `"potential broad use"` (wx95, 1oNu). In this place, we would like to first provide several general clarifications to enhance overall understanding of our work.

**1.1 Generalizability of ERDiff across Datasets**

The spatio-temporal structures are intrinsic to neural dynamics associated with behaviors, and have been extensively studied in neuroscience [1,2]. We would like to emphasize that with the powerful distribution learning ability of diffusion model (DM), ERDiff has great potential to generalize across a wide range of neural behavior datasets and applications. For validation, we run additional experiments on an unconstrained rat’s hippocampus dataset [3] and verify the efficacy of ERDiff on it (please refer to Fig. 1 and Table 1 of the attached PDF). These results will be provided in the revised manuscript.

**1.2 ERDiff is an unsupervised neural distribution alignment method**

To be in line with the conventions of previous studies on neural distribution alignment [4, 5], in the manuscript, we agree that behavioral (velocity) signals of the source domain are present during the VAE training. These signals do help in learning a more interpretable neural latent space. However, we emphasize that ERDiff does not incorporate any behavioral signals of the target domain during the distribution alignment phase. Hence, ERDiff is entirely an **unsupervised** neural distribution alignment (i.e., test-time adaptation) method.

We also note that the introduction of behavior signals in the source domain is an alternative choice. We conduct additional experiments to verify the efficacy of ERDiff when such behavioral signals are removed. Please refer to Table 3 of the attached PDF for detailed results. 

**1.3 (1) Computational Cost and (2) Time Complexity of ERDiff of Alignment Phase**

**(1)** In the alignment phase, for any given target domain, ERDiff can stably align it to the source domain in a comparable overhead with baselines. In Table 2 of the attached PDF, we conduct a comparative analysis between ERDiff and baseline methods in terms of additional parameter number, additional model size, stability, and alignment time. The demonstrated alignment time corresponds to the execution time for aligning one iteration (a batch of size 64) on a MacBook Pro (2019 equipped with 8-Core Intel Core i9 and 4 GB RAM). These analyses will be provided in the revised manuscript.

**(2)** Here we conduct time complexity analysis with respect to the batch size $B$ for the alignment phase. The ERDiff's alignment objective is composed of two main terms: Diffusion Noise Residual and Sinkhorn Divergence. We note that in the diffusion noise residual computation, it does not go through the entire $T$ diffusion steps. Instead, it just samples a single time step (noise scale) $t$ and calculates the noise residual specific to that step. Thus, the total complexity of this part takes $\mathcal{O}(K_1 *B * d)$, in which the coefficient $K_1$ relates to the inference complexity of the DM denoiser $\boldsymbol{\epsilon}\left(\mathbf{Z}, t \right)$; $d$ denotes the latent dimension size. For the Sinkhorn Divergence, it has to compute the distance matrix, costing $\mathcal{O}(K_2 *B^2)$; $K_2$ is a relatively small coefficient in magnitude. By summing up, the total complexity of ERDiff is given by $\mathcal{O}(K_1 *B * d + K_2 *B^2)$. This $\mathcal{O}(B^2)$ complexity is applicable since the non-adversarial baseline methods we compared (i.e., JSDM, and SASA) require quadratic complexities as well. 

These analyses will be provided in the revised manuscript.

**1.4  New 'Discussion' Section**

Here we delve into Limitations, Future Work, and Broader Impact of our work. These parts will be put into the revised manuscript in a new 'Discussion' section, including and replacing the existing 'Section 5: Conclusion'.

*Limitation and Future Work*: **(1)** Multi-domain Adaptation. Currently, ERDiff can align well with a single source domain latent distribution. An intriguing direction for future work would be learning a unified latent space across multiple source domains using the diffusion model. Thus the method would be applicable to domain generalization problems. **(2)** Generalization on alternative latent variable models (LVM). ERDiff currently identify the latent variables of raw neural signals with a canonical version of VAE. However, the architecture of the LVM within ERDiff is actually disentangled from the diffusion model training or MLA procedure. Future work includes validating ERDiff given multiple implementations of LVM (e.g., LFADS, pi-VAE). 

*Broader Impact:* Not confined to computational neuroscience, the cooperative training technique and the MLA in ERDiff have potential to apply into broader domain adaptation tasks across general time-series datasets (e.g., weather forecasting, and seismology). We also expect that our method can be applied or extended to other BCI applications and the broader field of neuroscience/AI.

Refs:

[1] STNDT: Modeling Neural Population Activity with Spatiotemporal Transformers. (Trung et al. 2022)

[2] Deep inference of latent dynamics with spatio-temporal super-resolution. (Feng et al. 2021)

[3] Recordings from hippocampal area ca1, pre, during and post novel spatial learning. (Grosmark et al. 2016)

[4] Robust alignment of cross-session recordings of neural population activity. (Justin et al., 2022)

[5] Stabilizing brain-computer interfaces through alignment of latent dynamics. (Brianna et al. 2022)

Meta-review: The reviewers agree that this is a novel, clear, and technically solid contribution. The additional work during the rebuttal and discussion phase further improved the manuscript. There is a growing need for methods that can align neural and behavioral activity across individuals and sessions, and ERDiff shows good potential for tackling this problem.
Decision: Accept (spotlight)