Leveraging Generative Models for Unsupervised Alignment of Neural Time Series Data
Authors: Ayesha Vermani
Il Memming Park
Josue Nassar
Venue: ICLR 2024 poster (2024)
Keywords: neural dynamics
transfer learning
distribution alignment
neuroscience
few-shot learning
Primary area: applications to neuroscience & cognitive science

Abstract: Large scale inference models are widely used in neuroscience to extract latent representations from high-dimensional neural recordings. Due to the statistical heterogeneities between sessions and animals, a new model is trained from scratch to infer the underlying dynamics for each new dataset. This is computationally expensive and does not fully leverage all the available data. Moreover, as these models get more complex, they can be challenging to train. In parallel, it is becoming common to use pre-trained models in the machine learning community for few shot and transfer learning. One major hurdle that prevents the re-use of generative models in neuroscience is the complex spatio-temporal structure of neural dynamics within and across animals. Interestingly, the underlying dynamics identified from different datasets on the same task are qualitatively similar. In this work, we exploit this observation and propose a source-free and unsupervised alignment approach that utilizes the learnt dynamics and enables the re-use of trained generative models. We validate our approach on simulations and show the efficacy of the alignment on neural recordings from the motor cortex obtained during a reaching task.

Published as a conference paper at ICLR 2024
LEVERAGING GENERATIVE MODELS FOR UNSUPER- VISED ALIGNMENT OF NEURAL TIME SERIES DATA
Ayesha Vermani', II Memming Josue Nassar
f Champalimaud Centre for the Unknown, Champalimaud Foundation, Portugal f RyvivyR, USA
(ayesha.vermann) vermani, memming, htaerrrrrrrerrc....eaaaaaaaaaae org
josue wwwu...asaarreeeeee.ee nassarrrvvrrrvyyyyy com
ABSTRACT
Large scale inference models are widely used in neuroscience to extract latent rep- resentations from high-dimensional neural recordings. Due the statistical het- erogeneities between sessions and animals, a new model is trained from scratch to infer the underlying dynamics for each new dataset. This is computationally expensive and does not fully leverage all the available data. Moreover, as these models get more complex, they can be challenging to train. In parallel, coming common use pre-trained models in the machine learning community for few shot and transfer learning. One major hurdle that prevents the re-use of generative models in neuroscience is the complex spatio-temporal structure of neural dynamics within and across animals. Interestingly. the underlying dynam- ies identified from different datasets on the same task are qualitatively similar. In this work, we exploit this observation and propose a source-free and unsuper- vised alignment approach that utilizes the learnt dynamics and enables the re-use of trained generative models. We validate our approach on simulations and show the efficacy of the alignment on neural recordings from the motor cortex obtained during a reaching task.
be-
INTRODUCTION
With advancements in recording techniques, we have access to a large number of simultaneously recorded neurons, exhibiting complex spatio-temporal activity. Consequently, significant efforts have been dedicated to the development of computational models that can infer the underlying structure from these recordings (Linderman et al. 2017 Pandarinath et al. 2018 Duncker et al. 2019 Schimel et al. 2022 Dowling et al. 2023). The progress in deep generative models, such variational autoencoders (VAEs) (Kingma & Welling 2013) and sequential variational autoen- coders (Bowman et al. 2015 Hafner et al. 2019). has further contributed to a proliferation of these latent variable models for neuroscience. These models trained to extract the latent dynamical process- - typically confined to a low-dimensional manifold - that drives the high- nigh-dimensional neural or behavioral observations.
Despite the abundance of latent variable models for neural data, there are some issues that prevent their widespread adoption by the experimental community. Firstly, training large models can be data intensive; although the number of simultancously recorded neurons continues to increase, the number of trials subject can perform during a single experimental session is still limited (Williams) & Linderman] 2021). Furthermore, there is growing interest in studying naturalistic behaviors the field, where trial boundaries are ill-defined and trial repetitions are few, if any (Rosenberg al. 2021 Kennedy) 2022; Minkowicz al., 2023). Secondly, training deep neural networks is computationally expensive and can pose several challenges. This is partly attributed to the complex relationship between training process and hyperparameter optimzation, which can considerably im- pact the model's 's performance.
In parallel, the use of pre-trained models has led to significant breakthroughs in natural language processing and computer vision Girshick et al., 2014 Rasmy et al., (021)). These are driven by the empirical observation that model re-use is highly data efficient and achieves comparable per-

---

Published as a conference paper at ICLR 2024
formance as model trained from scratch, with only fraction of the data (Goyal al. 2019 Moreover, re-using a pre-trained model allows us to circumvent challenges associated with training model from scratch. Recent evidence also suggests that pre-trained models are fairly generalizable and can be fine-tuned to perform a variety of tasks, even across domains (Parisi et al. 2022).
Inspired by the empirical success of pre-trained models in machine learning and the recent interest in training large models for neuroscience (Azabou et al.) 2023] Ye et al., (023), we investigate the case of using pre-trained sequential VAES (seq VAB)) for neural time series data. seq have been widely successful at inferring the underlying latent dynamics from high-dimensional neural time series data. However, due to statistical heterogeneities across datasets, arising from disparities in the number and tuning properties of recorded neurons, differences in recording modalities, etc., pre-trained seq VAEs cannot be re-used directly on new recordings. A potential approach tackle this problem is by learning an alignment that transforms the new dataset such that it is statistically similar to the data used to train the seqVAE. Previous approaches for learning an alignment between neural datasets require access to the original data used to train the model and/or the existence of paired samples between datasets (Degenhart et al. 2020; Chen et al. 2021 Williams et al. [02T]] Duong et al. 2023 Wang et al. 2023). The paired samples pairing dimulus-condittinn
commonly constructed by arbitrarily neural activity across the datasets. This entirely ignores rrial-to-rria variability and cannot be applied naturalistic task settings. Moreover, many of these methods do not explicitly model the temporal structure of data which can lead to suboptimal learning of the
alignment (Wang al. 2023
Alignment
Encoder
Fixed prior dynamics
masc
Inferred latenss
Alter training
min D
naurons
Ient 9139
sined reference dynamics
stoee
time
-
q(x Yaroma
p(x Ix,)
Interred latents follow reference dynamics
New Ocsorvations
Trakned
Figure 1: Schematic of unsupervised alignment scheme. We train a sequential VAE on some refer- ence observations to learn an encoder, along with the underlying low-dimensional dynamics. Given new observations Ynew generated from the same dynamical process, we learn a function g that tran- forms and implicitly aligns them to the reference, allowing for the re-use of the pre-trained model. The latent trajectories inferred after transforming the new observations, g g(Jnew), ), are constrained by the learnt latent dynamics.
In this work, we propose a simple unsupervised method for aligning neural time series that facilitates the re-use of pre-trained seq VAEs. Our approach stems from the observation that learning to perform inference on new recordings using pre-trained seqVAE implicitly results in learning an alignment between datasets. Moreover, our approach leverages the similarity in temporal dynamics across datasets to learn the alignment, as optimizing for inference in this framework encourages the inferred latents to be close the learned dynamics (see Figure D.. As our proposed method is unsupervised, it does not require the availability of paired samples nor access to the original training data, making it highly flexible and easy to use. We empirically validate our method on synthetic experiments and test it on neural recordings obtained from the primary motor cortex (M1) of two monkeys during center out reaching task (Dyer et al. 2017).
The main contributions of this paper are summarized as follows: (1) We propose novel unsuper- vised method for implicit alignment sequence data with low-dimensional dynamics that enables the re-use of trained generative models. (2) For the linear and Gaussian setting, we analytically demonstrate that the proposed approach recovers an alignment that close to the optimal. (3) On synthetic and real data, we provide empirical evidence that the proposed method outperforms other methods. The corresponding code is available at httpp:: /ttl::///ittai /ithub..tuuuuuug/// com/ ayeahav/ a align- seqvae
2

---

Published as a conference paper at ICLR 2024
2 RELATED WORK
There has been large body of work on aligning neural datasets for applications such computing
metries between neural representations (Williams et al. 2021 Duong et al. 2023); enabling the use of decoders across sessions (Sussillo et al. 2016bb Farshchian et al. 2019; Gallego et al. 2020 De-
genhart et al. 2020 Ma et al., 2023 Wang et al. 2023), across animals (Herrero- Vidal et al. 2021
Chen et al. 2021), and even across species (Rizzoglio et al. 2022). One popular alignment approach minimizes the error between the original and the aligned dataset, using methods such as Canonical Correlation analysis(CCA) (Gallego et al., 2020; Rizzoglio et al., 2022) or Procrustes (Degenhart et al. 2020: Duong et al. 2023). However, they require access to the original dataset and require the existence of a one-to-one correspondence between the two datasets. Moreover, they don't t lever- age the spatio-temporal structure in neural time series. A related approach works by minimizing a divergence between the original and aligned dataset, either explicitly (Dyer al., 2017 Karpowicz
et al. 2022 Wang et al. 2023). or implicitly via generative adversarial networks Farshchian et al., 2019
Ma et al. 2023).
The two most similar approaches to the proposed method are NoMAD Karpowicz et al. 2022) and ERDiff (Wang al.] 2023). In NoMAD, LFADS (Sussillo al. 2016a) popular seq VAE for neural time series- is first trained on the original dataset. New datasets are aligned by maximiz- ing the log-likelihood of data and minimizing the KLD between the distribution of latent states of the original and aligned dataset, where both are assumed to be Gaussian. Crucially, since LFADS models the latent dynamics to be deterministic, the log-likelihood accounts for the spatio-temporal structure of data. While similar to the proposed method, we make no assumption on the distribu- tion of latent states nor do we require the statistics of the original dataset. In ERDiff, a VAE is trained along with a spatio-tempora diffusion model that upproximates the latent distribution of the original dataset. Given new data, ERDiff optimizes the alignment function to produce latent states from the pre-trained encoder that are likely under this distribution. Although this is similar in spirit to the proposed approach, there are significant differences. Namely, ERDiff requires training spatio-temporal diffusion model, along with a VAE, on the source dataset to perform align- ment on new data. This incurs additional overhead as re-using or sharing the pre-trained VAE necessitates training diffusion model. Moreover, does not use the learned latent dynamics for aligning as it only re-uses the encoder, instead relying spatio-temporal transformer blocks to cap- ture the spatio-temporal structure the data. In contrast, our approach considerably simpler as only requires training an inexpensive alignment function that feeds into the encoder. Moreover, we explicitly consider the spatio-temporal structure by using the pre-trained latent dynamics.
3 BACKGROUND
3.1 SEQUENTIAL VARIATIONAL AUTOENCODER
In this work, we focus on learning an alignment between datasets that enables the re-use of state- space models (SSMs), class of latent variable models for spatio-temporal data. Let 2, E X C Rd, and y E y C Rdy be the low-dimensional latent state, and the observation at time I, respectively. A SSM can be described follows,
Dt Tz-1 po (Ft x2-1),
Iz (M Iz).
(2)
where equation [ is the latent dynamics distribution, parameterized by 0, that describes the temporal evolution of the latent state, ze; equation 2 the likelihood distribution, parameterized by o, that maps the low-dimensional latent state to the high-dimensional observation, y/. While there are many choices for the parametric form of equation 1] we follow standard practice (Krishnan et al., 2015 Hafner et al.) 2019) and parameterize it as Po (Tt (+-1)) N(xx, fo(= Q). where fo is a deep neural network (DNN). While one can also parameterize the likelihood distribution with a DNN, previous work has shown that making both the likelihood and the dynamics highly expressive can lead to optimization issues (Bowman et al. 2015). Thus, we parameterize the likelihood distribution to be linear function of Dr. Specifically. for spike data, parameterize equation 2 as (s | xx) = Binomial(y; 4, a (Cz, + D)) where a is the sigmoid function; for real-valued observation, such behavioral recording, we parameterize equation 2 as po(Ue (x)) = =V Cr., D2 R).
3

---

Published as a conference paper at ICLR 2024
Given neural time series dataset, V1-T [yi.... Mt,.... ".rl, we are generally interested in inferring the corresponding latent states, 21:7. and learning the parameters of the generative model, and Φ Exact inference and learning is difficult it requires computing the posterior, p(z1:r M1:r). and the log marginal likelihood, p()1r), which are both commonly intractable. We address this challenge by using the VAE model extension of VAEs for spatio-temporal data (Krishnan al. 2015). Similar to VAEs, seqVAEs are trained by maximizing a lower-bound of the morginaiiaal likelihood, commonly referred to as the evidence lower bound (ELBO). Specifically, given data, V1.T, the ELBO is defined as
T C(yi.T,0,0 Σ Egg [logPo 21) + log xx-1) (o (1rr)),
(3)
where 2. = go (21.rlar)) and (xr yı.r) commonyy referred as the encoder- a varia- tional approximation the posterior distribution, p(æ1:r (1:r). The parameters of the generative model, B, o and the encoder, v., are optimized jointly.
While there are various approaches for designing (2 y1:r), we follow the parameterization described in Krishnan et al. (2015) for simplicity, where qo (ir)) = II.m (0 M..T) and
(x Vi:r) N (xa ( (wn:)))
where pv() and o (+) are bidirectional recurrent neural networks.
3.2 ALIGNMENT OF NEURAL TIME SERIES
Now let's consider VeQ model trained on yı.r. which are interested in re-using for new dataset, Wi.T = [W1... ........ WT] where W/ E W c Rdw, In general, W1.T will not follow a similar distribution to 81:T. This can be due to several reasons there might be drift in the recording probes over sessions, the data might have been collected from a different animal, or using a different record- ing modality, and so on. The distribution mismatch between f1.r and w1.r prevents straightforward application of the trained seq VAE to WIT.
One approach for re-using this model for WI. is learning an alignment function between the datasets, go : W J/ that projects w, to y/ i.e. Dt g(w)). The projected data can subsequently be fed to the pre-trained encoder, i.e. (1 (w(w)))) go go (w(r)) = [go [go(w)),.... gol (wr))] thereby enabling us to re-use it for inferring the latent states from w1.7. Broadly, the main ob- jective for optimizing the alignment function, 9a. is minimizing the distance between the original data distribution, p(V1:r). and the distribution of projected data, p(g1-r) (Dyer al. 2017 Duong eet al. 2023). Directly minimizing the distance between the two distributions is usually infeasible, as we don't 't have knowledge about the marginal distributions of the datasets. Moreover, most common distance measures are tractable for a limited class of distributions, many of which are not able to effectively model complex spatio-temporal neural activity.
An alternative is using a supervised learning approach to learn the alignment. Specifically, suppose that we have paired samples from the two datasets, i.c., D {wr, yo ) We can then learn go by minimizing the error between Ye and go (w,), i.e., | - go (w, 2l Although this approach can recover the optimal alignment, requires the source dataset y, and has the restrictive requirement of paired samples between the datasets.
4 UNSUPERVISED ALIGNMENT OF NEURAL TIME SERIES
In order to re-use the pre- preitrdined seq VAE, we begin by assuming that the underlying latent dynam- ies for Wi. are the same as y1.7, allowing us to fix the parameters of the trained dynamics model, po(x; | x2-1). This is supported by empirical evidence that the inferred latent dynamics from different neural networks (both biological and artificial) performing the same task are similar (Ma-) heswaranathan et al. 2019] Safaie et al.) 2022. Brain et al.) 2023). The other components of pre-trained seq VAE consist of the likelihood function, P (Bt xt). and the encoder, qu (x1.7 (1:r).
Forr ease of presentation, set the length of V1.2 and wir to be the same, but the proposed approach does require this to he the case.
4

---

Published as a conference paper at ICLR 2024
The encoder and likelihood both assume that the observations have dimensionality of dy but in gen- eral, de * dy. Moreover, different recording modalities require different likelihoods, thus p- xx) may not be the suitable parametric form for w/. However, given that paramterize the likelihood as a linear function of Tt, re-training a likelihond model for WI.r will result in minimal compu- tational overhead. Thus, we train new likelihood distribution specific the new dataset, WI:r, i... Pew (we) Ft). In contrast, the encoders for seq VAEs are usually parameterized by large neural networks, such as bidirectional recurrent neural network, and would be preferable to keep its parameters fixed. As described in Section 3.2] way to avoid re-training the encoder is by learning an alignment between the two datasets. In this work, we propose an unsupervised algorithm based off simple observation.
In the VAE framework, the role of the encoder to infer the latent states given the observed data, where the optimal encoder corresponds to the true posterior distribution (Blei et al. 2017). Suppose that the pre-trained encoder provides good approximation to the posterior, i.e., qu (F1.T M1.T) Ro p(Z1.T VI.r). Intuitively, good alignment function should facilitate the re-use of the pre-trained encoder to obtain a reasonable approximation to the posterior on the new dataset, i.e., qo (xi:r go (w):r)) p(#1.r W1:T). To study the validity of this intuition, we consider simple linear model that affords analytical tractability; for ease of presentation, we drop the time index. Let p(xr) = (0, I), ply x) (Ax, Q). p(w x) (Cr,,, R) and ga(w) &w, where / is the identity matrix. Based on the previous intuition, we can optimize the parameters of the alignment, 0, by minimizing the expected Hullback-Leible divergence between q(< go(w) and p(æ w)
0. = arg min Ep(w)) [TxL. (q (rr go(w)) I|p(2 [w)]]
(5)
which is equivalent to maximizing the expected ELBO
0,- arg max
Ep(w) [ [logp(w x)) + p(x) - log q (z go(w))]]
(6)
Recalling that the optimal encoder is the posterior and that the linear and Gaussian model allows for a tractable posterior (Bishop 2007 we define q(> y) as
q y) = p( v) = -N)=)(((y))))))
p(y) A ΣATQ
(8)
Σ a A==''++++"
Thus, g(x (o(w) ) = N1 (go (w)
Σ). The tractability of this simple model allows us to directly 9.. with the optimal alignment with respect to the mean-squared
compare the solution of equation 6 error, 0,
0 = arg min E (((wa)) [(v ow) (y - Uw)]
(10)
In the following proposition, we demonstrate that 1,, can be expressed a linear transformation of 0 the proof can be found in the Appendix (A)
Proposition 1 Let 9, be the solution of equation 6 and 0, be the solution of equation T0 Then 9. -I+ Q(AAT)-1) 01. where I the identity matrix.
Proposition demonstrates that by optimizing equation 6 we obtain a linear transformation of the optimal alignment, 0+. Moreover, we see that the difference between 94 and 9, is a function of Q(AAT)-I. Thus, when the new observation noise, Q is small and/or when AAT is large, we expect for 0, x 0z. We emphasize that we are able to implicitly learn good approximation of the optimal alignment function in an unsupervised fashion without paired samples or the source data, y.
Inspired by Proposition[] [ we move to designing a general purpose algorithm for unsupervised learning of an alignment function, ga. straightforward approach is to jointly learn the parameters of the alignment, go. and of the dataset specific likelihood, Ps.. (w, | xx)) by optimizing the ELBO
T
[(WITT 2w. 0)
Σ Eq... [log Pou. 1=1
(w, x²) + logpe(zz xt-1) qu (Ire (o(w))))), (11)
where Eyu,o (e) and both the latent dynamics, pe (2t æイ-1), and the encoder, qo (x, "I.T), ) are kept tixed.
5

---

Published as a conference paper at ICLR 2024
While optimizing equation 11 is simple and can lead good empirical performance, we found that it was easy for the optimizer to converge to suboptimal local minimum. Further investigation revealed that the optimizer would produce latent states that are likely under the one-step ahead dynamics, log pa 1t-1), but would not respect the global dynamics; Fig presents an example (denoted as 1-step prior).
To regularize the optimizer to produce latent states that respect the global dynamics, we replace the one-step ahead dynamics, log pe (z-1),,
with K-step ahead dynamies term, ΣØ=1 log (2 2z-K),
K-stap
(Hafner et al.) 2019), which encourages the latent states to follow the dynamics over k-step horizon. Following (Hafner et al. 2019). al-
Figure 2: A case for K- step regularization.
though log ( Æ1-K) is intractable, it is straightforward to obtain an unbiased Monte Carlo estimate; in the Appendix we discuss
how unbiased estimate is obtained. This leads the K K-step ahead ELBO, which use going forward
K
LK (WIr, pw. 1) =
Σ
E
log P an (wo xx) +
log po (o [T-K) (z+ go (wI:r)
We note that, although log po (o K+j unbiased Monte Carlo estimate (Hafner et al.) (019). In Fig [ we see that using the K-step ahead ELBO leads to much better latents (denoted as K prior)
K(-K) is intractable, it is straightforward to obtain an
5 EXPERIMENTS
We validate our alignment approach synthetic datasets generated with via the Van der Pol dy- namics and the Lorenz system. Then, we test our method on neural recordings obtained from the primary motor cortex (MI) in two monkeys during reaching task (Dyer et al. 2017). We compare the proposed approach against the following methods:
ERDifT (Wang et al. 2023). This method uses pre-trained seq VAE along with diffusion model with ppapooltemporal transformer blocks to estimate the density of latent trajectories on the original dataset ps (. The alignment function is trained to maximize Eg,, [logPs (x(rr)) where the encoder is kept fixed. They additionally optimize a Sinkhorn divergence hetween the source and target latents.
NoMAD (aarpowizz et al. 2022). Given a pre-trained seq VAE, NoMAD fits a multivariate Gaus- sian the inferred latent states from the original dataset, Py (a)) (My, Ey). The alignment func- tion is trained to maximize Σ/= Equ. [log P (w,)] - Px ( (p)]lppp))) where pe - -M(H, Eg) is Gaussian distribution fit the latents from the new dataset.
Cycle-GAN (M et al. 2023). Cycle-GAN leverages adversarial training, via a generative adversar- ial network, to align new sessions to the original dataset.
Orthogonal Procrustes (Schönemann) 1966). An alignment learned via Orthogonal Procrustes. We note that this requires paired samples from the original and new datasets.
Re-training. We train a generative model from scratch on the new dataset as an upper bound on performance.
To isolate the benefits of each method, one seq VAE is trained and is then given to all methods. Due space constraints, we defer training and architecture details to the Appendix C
5.1 VAN DER POL OSCILLATOR
The Van der Pol oscillator is -o-ddimessiona nonlinear dynamical system. We consider noisy version of this system described as follows:
1
Φ = p(
r x2) +
$ =
+
(12)
3

---

Published as a conference paper at ICLR 2024
where μ 1.5 and (0(00.11) 0.1). For training the seq,AE, generated 1600 trajectories of length = 300 with spike observations where the number of neurons was set to 250, i.c., d, = 250. To avoid aligning the raw spikes (which are nigh-dimensional,, use non-linear embedding function that down-projects spikes to 64 dimensions using an MI.P before passing it into the encoder.
generated three W3.1:7 each of length T = 300-where each dataset has a different number of neurons (d = 200, due 250, dag 300). For the proposed approach, ERDiff and NoMAD, we parameterize Go as an MLP. For each dataset, all methods were trained using 500 trajectories and were evaluated on a held-out test set.
For evaluating the alignment methods,
datasets, W1,1:7++ 1221177, and
Dynamics leamt by pre-trained
True fiing procabilty
C. Likelihood of source dynamics
D.
Forecasting
Few-shot performance
performance
30 alep
reconstruction prediction Our approach
0.6
C.6
0.7
ERCH
10
20
30
40
50
20
40 50
prediction
rahnn trajactorias
Trajectories latent samples alter alignment
data
1-slep prior K-stop prior
Methods
Our
ERD
Procuusee
C.dd NaMAO
Resrain
timo
Figure 3: A. The vector field is generated from the learnt source dynamics. Sampled latent trajec- tories from the encoder after aligning the observations using the proposed approach, ERDiff and CycleGAN. B. The true firing probability of example neurons on a trial (top) and the reconstruction and prediction for the aligned data using the pre-trained model (below). C. Likelihood of source dy- namics given inferred latents for the source and aligned data. D. K-step prediction r² performance. E. K=30-step prediction performance for various approaches. We plot the median (solid) and the [20, 80] percentile (dotted) 2
In Fig. 3A we plot example latent trajectories sampled from the encoder using comparing our approach to ERDiff and CycleGAN. We see that our approach produces smoother latents that respects the pre-trained dynamics; this leads to better reconstructed firing rates and smoother predictions Fig. B. To quantify whether the alignment procedures lead to latents that respect the dynamics we compute the likelihood of the inferred latents on the trained dynamics, i.e.,
Σ. (go [log (0())(2(-1) (Fig. 3C). We see that the proposed method outperforms all comparisons. Moreover, opposed to the standard 1-step ahead prior.
that using K-step ahead prior leads to better performance as
We subsequently evaluate the methods on their forecasting performance. We use the first 250 time points to infer the latents after aligning and sample 50 steps in the future. We measure the per- formance by computing the ,d between the true and predicted trajectories (Fig 3D) Our approach performs close to a model trained from scratch on the new dataset. We also test the few-shot perfor- mance of these methods. In Fig we plot the forecasting performance for 30 as a function of the number of trajectories used for training. Even in the one-shot regime, our method consistently achieves high 2 performance and demonstrates low variance compared to other alignment methods.
To demonstrate that the method can also allow for alignment across recording modalities, we include an experiment in the Appendix where we align real-valued data to the pre-trained model Fig 6 Specifically, Wt N (Cr,,, a /), where d,, = 30 and 0.1. From Fig. 6 we see that the proposed method performs well and able to match the forecasting performance of a model trained from scratch.
5.2 LORENZ ATTRACTOR
The Lorenz attractor is three-dimensional system with chaotic dynamics described by the following set of equations,
21 a (12--z1) 21), z 2 ( x3) I2, where o = 10, 3 = 8/3, and p= 28. For training the seq VAE, we generated 1600 trajectories of length T = 500 with real-valued observations where d,y 40.
2/3 T1 .3r3,
(13)
7

---

Published as a conference paper at ICLR 2024
For evaluating the alignment methods, gen- erated two more datasets, W1.1:7, and W2,1:7, each of length T 500, where dw, 35 and day = 55. For the proposed approach, ERDiff and NoMAD,
A. Sample latent trajectories
B.
Forecasting
performance
0.6 prediction
parameterize go linear function. For each dataset, all methods were trained using 1,000 trajectories and were evaluated on a held-out test set. We evaluate the models on reconstruction and forecasting, where for forecasting we use 400 time points to infer the latents and sample steps in the future.
0.2
Source
10
20 30 40 50
Quit
BREWT
Prosusee
C-GAN
HiMAD
Reetain
prediction
step
Re- Re-ramiig with more data
Figure 4: A. Latent samples from new data using our approach are aligned with the source latents. B. Prediction performance for various alignment methods and re-training a model from scratch.
In Table T we display the reconstruction and forecasting y² for each of the methods. On reconstruction, we see that both the proposed approach and NoMAD perform very well and are able to match the performance of training a that for prediction horizon of only 5, all the baselines
model from scratch. In contrast, deteriorate. Only the proposed approach is able to match the forecasting performance of model trained from scratch. In Fig we see that the proposed approach allows for stable forecasting up to 30 time steps ahead.
Methods
Reconstruction y22
5-step ahead y
Re-training
0.99 + 0.0005
0.99 1 0.0006
Our approach
0.99 + 0.0008
0.99 + 0.0012
ERDiff
-0.08 + 0.0496
-0.23 + 0.0553
NoMAD
0.99 = 0.0005
-0.03 + 0.0714
Cycle-GAN
0.86 + 0.0131
-0.42 + 0.0937
Procrustes
0.68 0.052
0.18 + 0.0833
Table 1: Reconstruction and forecasting performance for the Lorenz attractor. The values indicate the median and standard error over the observations from new sessions. We report the prediction performance for 5-step ahead prediction.
Next, compare the alignment learnt from approach to the optimal alignment that can be obtained with respect to mean-squared error. In order to do this, we simulated 100 trajectories from the Lorenz attractor, and used the same likelihood models above to generate observations with paired samples. The alignment function from our unsupervised approach closely matches the optimal (Fig. 7B, RMSE: 0.0017 + 0.002).
5.3 NEURAL RECORDINGS
We applied our method to motor cortex recordings from two monkeys (M and C) during delayed center out reaching task(see (Dyer et al. 2017) for details). Briefly, the monkeys were trained to use manipulandum move a cursor to one of eight possible target locations on the screen (Fig. 5 As they performed the task, electrophysioloical activity was recorded from MI along with the hand position and velocity. For each monkey, two sessions of data were available where the number of correct trials per session ranged from 159 to 215 while the total number of neurons varied from 150 167. Following (Wang et al., 2023), we pre-process the data by first binning the neural activity into 20 ms bins. The binned spikes were then smoothed using 50 ms Gaussian kernel.
We trained VAE on session from monkey M as we observed that the recordings from this session were highly informative about the monkey's 's behavior relative to the other datasets. We set the latent dimension to be 30 and also learn an embedding that projects the smoothed spikes down to 64 dimensions before being passed into the encoder. To ensure that the latents were also informative of the behavior, we included likelihood terms for both the smoothed spikes and the monkey's 's hand velocity where Gaussian likelihood M along with sessions and 2 from Monkey C as new datasets and use them to investigate the performance of the methods. For the proposed approach, NoMAD and ERDiff, we parameterize
used in both cases. We treat session 2 from Monkey
8

---

Published as a conference paper at ICLR 2024
A.
Decoded hand trajectories
Center-out
Ours
ERDI
NoMAD
GAN
Procrustes
reaching task
Pa-tA
Pw
Cross subject
Ps 0.883
#f 4023
Ps 0.884
Pli -0,060
Ps 0.546
Figure 5: A. Schematic depicting the center out reaching tasks that the monkeys performed. B. True and decoded monkey hand trajectories after alignment.
Go with an MLP. For the proposed approach and NoMAD, the behavior likelihood term was also included the loss function, but its parameters were kept fixed.
Pig.B. 3. we demonstrate example decoded hand trajectories for some methods where we see that the proposed method, along with NoMAD and Procrustes are able produce good decoded hand trajectories, both across sessions and across monkeys. In Fig. 8 we plot example latent trajectories where a similar trend observed. In Table 2 we quantify the reconstruction and forecasting per- formance for each of the methods. We see that the proposed approach, NoMAD and Procrustes are able to reconstruct the hand trajectories, with NoMAD performing slightly better. forecasting, we see that only the proposed approach is able to forecast well while the other approaches struggle.
In Fig.
Methods
Reconstruction r²
5-step ahead y
Our approach
0.66 + 0.023
0.39 + 0.071
ERDiff
-0.32 + 0.38
-0.23 0.553
NoMAD
0.68 + 0.021
0.15 + 0.101
Cycle-GAN
-0.15 = 0.154
-0.81 0.121
Procrustes
0.61 + 0.051
0.07 = 0.141
Table 2: Reconstruction and forecasting performance for monkey hand trajectories. Values indicate the median and standard error over the observations from new sessions. We report the prediction performance for 5-step ahead prediction.
6 CONCLUSIONS AND LIMITATIONS
In this work, propose unsupervised alignment approach that leverages the temporal dynamics learnt from a source dataset to align new data. This enables the re-use of pre-trained generative model without access training data the restrictive requirement of paired samples. We demon- strate the efficacy of our approach by re-using a VeQ trained on neural recordings from MI of monkey predict behavior different sessions. This lends further credence the hypothesis that low dimensional neural representations play a crucial role in neural computation and behavior. The importance of studying these representations in common space has been previously high- lighted (Dabagia et al. 2022) and is naturally afforded by our approach. Moreover, the variability in recordings not explained by the common dynamics assumption can offer a complementary insight into individual differences.
While the proposed upproach is promising, there are limitations and room for improvement. Firstly, we assume that we have a good pre-trained model that has learnt the underlying dynamics well. An important direction for future research would be learning the underlying dynamics on task using multiple datasets to identify generalizable latent representations. Secondly, the proposed approach relies assumes that the latent dynamics are exactly the same across datasets. Thus, we would expect method to work well on recordings obtained during the same or similar cognitive task the data used for pre-training. Moreover, this assumption does not take behavioral variability into account. The dynamics tasks with different structure that require the same computation would introduce additional variability. Fine-tume tuning the model after aligning recordings is one possibility to get good performance across different contexts and would be interesting direction for future work.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
AV and MP were supported by NIH RFI DA056404 and the Champalimaud Foundation. We thank the anonymous reviewers for their helpful feedback and comments.
REFERENCES
Mehdi Azabou, Vinam Arora, Venkataramana Ganesh, Ximeng Mao, Santosh Nachimuthu, Michael J. Mendelson, Blake Richards, Matthew G. Perich, Guillaume Lajoie, and Eva L.. Dyer. A unified, scalable framework for neural population decoding, 2023.
Christopher M. Bishop. Pattern Recognition and Machine Learning (Information Science and Statis- fics). Springer, edition, 2007. ISBN 0387310738.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisti- cians. Journal of the American statistical Association, 112(518):859-93) 877, 2017.
Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Ben- gio. Generating sentences from a continuous space. arXiv preprint 151.035 2015.
Laboratory International Brain, Kush Banga, Julius Benson, Jai Bhagat, Dan Biderman, Daniel Bir- man, Niccolò Bonacchi, Sehastian A. Bruijns, Robert AA Campbell, Matteo Carandini, Gaëlle A. Chapuis, Anne K. Churchland, M. Felicia Davatolhagh, Hyun Dong Lee, Mayo Faulkner, Berk Gerçek, Fei Hu, Julia Huntenburg, Cole Hurwitz, Anup Khanal, Christopher Krasniak, Guido T. Meijer, Nathaniel J. Miska, Zeinab Mohammadi, Jean-Paul Noel, Liam Paninski, Alejandro Pan- Vazquez, Noam Roth, Michael Schartner, Karolina Socha, Nicholas A. Steinmetz, Karel Svoboda, Marsa Taheri, Anne E. Urai, Miles Wells, Steven J. West, Matthew R. Whiteway, Olivier Win- ter, and Ilana B. Witten. Reproducibility of in-vivo electrophysiological measurements in mice. bioRxiv, 2023.
doi: 10. 10.1200/2022.5.. .5..449999.
Hung- Tu Chen, Jeremy R. Manning, and Matthijs A. A. van der Meer. Between-subject prediction reveals shared representational geometry in the rodent hippocampus. Current Biology, 31(19), October 2021. ISSN 0960-9822. doi: 10. www.1//......2221..... 20.1007.661.
Max Dabagia, Konrad P. Kording, and Eva I.. Dyer. Comparing high-dimensional neural recordings by aligning their low-dimensional latent representations, May 2022. arXiv:2205.08413 [q-bio].
Alan D Degenhart, William E Bishop, Emily R Oby, Elizabeth C Tyler-Kabara, Steven M Chase, Aaron P Batista, and Byron M Yu. Stabilization of brain crain-computer interface via the alignment of low-dimensional spaces of neural activity. Nature biomedical engineering. 4(7):672 685, 2020.
Matthew Dowling, Yuan Zhao, and II Memming Park. Linear time gps for inferring latent trajectories from neural spike trains. arXiv preprint arXiv:2306.0 01802, 2023.
Lea Duncker, Gergo Bohner, Julien Boussard, and Maneesh Sahani. Learning interpretable continuous-timg time models of latent stochastic dynamical systems. In Proceedings of the 36th Inter- national Conference Machine Learning. PMLR, May 2019.
Lyndon Duong, Jingyang Zhou, Josue Nassar, Jules Berman, Jeroen Olieslagers, and Alex H Williams. Representational dissimilarity metric spaces for stochastic neural networks. In The Eleventh Internutional Conference Learning Representations, 2023.
Eva L. Dyer, Mohammad Gheshlaghi Azar, Matthew G. Perich, Hugo L. Fernandes, Stephanie Naufel, Lee Miller, and Konrad P. Körding. A cryptography- approach for movement de- coding. Nature biomedical engineering, 1(12):967-9) 976,, December 2017. ISSN 2157-846X. doi: 10. 10.10088441551-07--1091111
A Farshchian, JA Gallego, JP Cohen, Y Bengio, LE Miller, and SA Solla. Adversarial domain adaptation for stable brain-machine interfaces. In International Conference on Learning Repre- sentations, volume 7, 2019.
Juan A. Gallego, Matthew G. Perich, Raeed H. Chowdhury, Sara A. Solla, and Lee E. Miller. Long- term stability of cortical population dynamics underlying consistent behavior. Nature neuro- science, 23(2):260-277) 270, February 2020. ISSN 1097-6256. doi: 10. 10.1038/441533-019900 053-049--555544
10

---

Published as a conference paper at ICLR 2024
Ross Girshick, Jeff Donahue, Trevor Darrell, and Jitendra Malik. Rich Feature Hierarchies for Accurate Object Detection and Semantic Segmentation. In 2014 IEEE Conference on Computer Vision and Pattern Recognition, June 2014. doi: 10.1109//VPR.2014.811 ISSN: 1063-6919.
Priya Goyal, Dhruv Mahajan, Abhinav Gupta, and Ishan Misra. Scaling and Benchmarking Self- Supervised Visual Representation Learning. 2019 IEEE/CVF International Conference on Computer Vision (ICCV), Seoul, Korea (South), October 2019. IEEE. ISBN doi: 11.1100//111. (O////CC.22.900000000
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning Latent Dynamics for Planning from Pixels. In Proceedings of the 36th International Conference on Machine Learning. PMLR, May 2019. ISSN: 2640-3498.
Pedro Herrero-Vidde Dmitry Rinberg, and Cristina Savin. Across-animal odor decoding by proba- bilistic manifold alignment. In Advances in Neural Information Processing Systems, volume 34. Curran Associates, Inc., 2021.
Brianna M. Karpowicz, Yahia H. Ali, Lahiru N. Wimalasena, Andrew R. Sedler, Mohammad Reza Keshtkaran, Kevin Bodkin, Xuan Ma. Lee E. Miller, and Chethan Pandarinath. Stabilizing brain- computer interfaces through alignment of latent dynamics, 2022.
Ann Kennedy, The what, how, and why of naturalistic behavior. Current Opinion in Neurobiology, 74:102549, 2022.
Diederik Kingma and Max Welling. arXiv:1312.6114 2013.
Auto-encoding variational bayes.
arXiv preprint
Rahul G Krishnan, Uri Shalit, and David Sontag. arXiv:1511.05121, 2015.
Deep
kalman filters. arXiv preprint
Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian Learning and Inference in Recurrent Switching Linear Dynamical Systems. In Proceed- ings of the 20th International Conference on Artificial Intelligence and Statistics. PMLR, April 2017.
Xuan Ma, Fabio Rizzoglio, Kevin L Bodkin, Eric Perreault, Lee E Miller, and Ann Kennedy. Using adversarial networks extend brain computer interface decoding accuracy over time. eLife, 12, August 2023. ISSN 2050-084X. doi: .ff..//e....e.42....
Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo. Uni- versality and individuality neural dynamies across large populations of recurrent networks. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Samuel Minkowicz, Mychaela Alexandria Mathews, Felicia Hoilam Mou, Hyoseo Yoon, Sara Nicole Freda, Ethan Cui, Ann Kennedy, and Yevgenia Kozorovitskiy. Striatal ensemble activity in an innate naturalistic behavior. bioRxiv, 2023.
Chethan Pandarinath, Daniel J. O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D. Stavisky, Jonathan C. Kao, Eric M. Trautmann, Matthew T. Kaufman, Stephen I. Ryu, Leigh Hochberg, Jaimie M. Henderson, Krishna V. Shenoy, L. F. Abbott, and David Sussillo. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature Methods, 15(10), October 2018. ISSN 1548-7105. doi: 1..1038//41522-018-0110 Publishing Group.
Number: 10 Publisher: Nature
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The Unsurprising Effectiveness of Pre-Trained Vision Models for Control. In Proceedings of the 39th International Conference on Machine Learning. PMLR, June 2022.
Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008
Laila Rasmy, Yang Xiang, Ziqian Xie, Cui Tao, and Degui Zhi. Med-BERT: pretrained contextu- alized embeddings on large-scale structured electronic health records for disease prediction. npj Digital Medicine, 4(1), May 2021. ISSN 2398-6352.

---

Published as a conference paper at ICLR 2024
Fabio Rizzoglio, Ege Altan, Xuan Ma, Kevin L Bodkin, Brian M Dekleva, Sara A Solla, Ann Kennedy, and Lee E Miller. Monkey-to-humant transfer of brain-computer interface decoders. bioRxiv, pp. 2022-11, 2022.
Matthew Rosenberg, Tony Zhang, Pietro Perona, and Markus Meister. Mice in labyrinth show rapid learning, sudden insight, and efficient exploration. eLife, 10:e66175, jul 2021.
Mostafa Safaie, Joanna C. Chang, Junchol Park, Lee E. Miller, Joshua T. Dudman, Matthew G. Perich, and Juan A. Gallego. Preserved neural population dynamics across animals performing similar behaviour. bioRxiv, 2022.
Marine Schimel, Ta-Chu Kao, Kristopher T Jensen, and Guillaume Hennequin. iLQR-VAE: based learning of input-driven dynamics with applications neural data. In International Con- ference Learning Representations, 2022.
VAE control-
Peter H Schönemann. A generalized solution of the orthogonal procrustes problem. Psychometrika, 31(1):1-10, 1966.
David Sussillo, Rafal Jozefowicz, LF Abbott, and Chethan Pandarinath. Lfads-latent factor analysis via dynamical systems. arXiv preprini arXiv: 1e0..6608...1.. 2016a.
David Sussillo, Sergey D) Stavisky, Jonathan C Kao, Stephen Ryu, and Krishna V Shenoy. Making brain machine interfaces robust to future neural variability. Nature communications, 7(1):13749, 2016b.
Yule Wang, Zijing Wu, Chengrui Li, and Angi Wu. Extraction and recovery of spatio-temporal structure in latent dynamics aligmment with diffusion model. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Alex H Williams and Scott W Linderman. Statistical neuroscience in the single trial limit. Current opinion in neurobiology, 70: 193-20,, 2021.
Alex H Williams, Erin Kunz, Simon Kornblith, and Scott Linderman. Generalized shape metrics on neural representations. Advances in Neural Information Processing Systems, 34, 2021.
Joel Ye, Jennifer Collinger, Leila Wehbe, and Robert Gaunt. Neural data transformer 2: Multi- context pretraining for neural spiking activity. bioRxiv, 2023. doi: 00.100//002.09.18.5555 13.
A PROOF OF PROPOSITION
We assume that p(x) = ~O0 I), x = ((xx, Q) and p(w x)) = NC(x, R). Let go (z) be linear function of x. i.e. 90 dx. Using standard Gaussian identities, can derive the marginal distribution for
p(v) / p(u x)p(z)dæ,
(14)
N(w Ax, Q) 0, D)dx, =N(y 0, +000+A0 AAT),
(15)
(16)
and w
p(w)= = / p(w x)p(x)dxx
(17)
/ N(w Ca, RW 1,, Tdr,,
(18)
-N(w 0, R CCTT..
(19)
12

---

Published as a conference paper at ICLR 2024
A.1 DERIVATION OF 1
We begin by first solving for the parameters of the optimal alignment function by minimizing the expected mean squared error (MSE)
f = arg min Bp(ww)) [(V - -w) T (v - Owi]]
(20)
To solve the above optimization, we expand the expected MSE
Ep(w.y) ((v- (ww) T (y 0w)] Ep(w,o) [V y + w O' Ow w0 yyyy]
(21)
= EEpcn)l 1y T '1 + Ep(ue) Tool 2E ww T 0 TTyl
(22)
= =tr(Q + AAT) ) + tr (0 O'TORR CCT)) 2tr(C 0 A).
(23)
To find 0,, we differentiate equation 23
20(R + CCT) - 2AC
(24)
We set equation 24 to 0 and solve for dt
20, (R + CCT) 2AC =0
(25)
+ = ACT
(26)
01 ACT (R CCT)-
(27)
A.2 DERIVATION OF 9,
We now move on to solving for 0,
0. = arg max Ep(w) [C(0, w)],
(28)
arg max Eptw) [B((lla(w)) p(w | x)) + logp(x) - log q( go)))]]]
(29)
arg max Ep(w) [B((rlo(ww) [log p(w | x)] DRL [g( [o(w)]|[()]]]
(30)
where DKL. is the Kullback- divergence. Recalling that the optimal encoder is the posterior, we set g(x) y) to be the posterior p(x| y)
((x y) = p(x 1)) = (N((v), E), p(v) - ΣA T Q y.
(31)
(32)
Σ a (AT A
(33)
thus
q() go(w)) (N( E),
(34)
(w)) e = Iw., Σ (A (A
(35)
(36)
Due to the linear and Gaussian assumptions placed on the model, we can analytically evaluate L(O, w). For conciseness, we use the substitution, ý A dw. when appropriate and absorb all con- stants with respect to into C
C(O, w) (g(rlgo(ww) [log p(w x)] - DkL. lq(2 so(w))p(())
(37)
(w))
[lo p(u
x)) (=(()) m(0) + C,
(38)
1 (w - Cx)) R- Cr)
(x())
2 (u(i) (M()) C.
(39)
(w)))
-r- x CT CTR-1 --
2 =1()) m(0 +C,
(40)
1 TCT R CC((i)
1
w R C=(())
J1()) (M()) + C
(41)
13

---

Published as a conference paper at ICLR 2024
We now take the expectation of the analytical ELBO from equation 41 tion, we set T ΣA EA'Q!
where to ease nota-
1 2 -20 CT R CC( 24(0) ul))
1 +C,
(10) [2(0, w]] E (w))
w T (.TR-CC(D))
(42)
ga w CT-W = C -/w Two +C, (43)
= ("
1 2
TT CCT))
+
))
+
(44)
To find 0., differentiate equation 44 and set it equal 0.
p [C(0, [() = 0,
(45)
0,
(46)
TCC TCC CTR-'C00700=C"F"
(47)
(48)
(49)
() R-!
(50)
TV = (1+ CTR
(51)
We the following variant of the Woodbury identity (Petersen et al.
2008
(equation 158)
(T C R C) - R- C (R+CC")
(52)
and plug it into equation 51 to get
TV T0= CT CR+CC
(53)
Recalling that = ΣATQ-I and Σ = A+1)
ΣA
C (R+CCT)-
(54)
(ATQ-AA4 1) - CT (R (R+CCT)
(55)
Using the same Woodbury identity, we can simplify the above
AT (Q + AAT) 0 CT (R+ (R+CCT)=, (A AAT (Q
(56)
0= - (R+CCT)
(57)
(Q + AAT) 8 (AAT)-IAC" (R+CCT)
(58)
1 (Q+ AA++AA") (AAT)- (AA 'ACT (R+CCT)"
(59)
8 Q(AAT)- ACT (R+C0) +AAT(AATT) -ACT
ACT (R+CCT)-
(60)
& Q(AAT) 1R+00 -1 + +ACT (R+CCT) i (I Q(AAT)-) ACT (R+ CCT)-!,
(61)
(62)
8. = -I+ ([++(AATT-)) 0r
(63)
B EVALUATING K-STEP AHEAD PRIOR
The K-step ahead log prior, log (o 12).. is defined
logp(и:4 xx) = pa( ax)) ... ..))...ppoozrp (o (14K log logeo [=#K-))dd<+44
(64)
We can rewrite the above as an expectation
logp(#:4K xx) = Epo(xe_ (l))
[log (o [(4K-1]]
(65)
which allows for unbiased estimator by applying the dynamics K times Tz-
14

---

Published as a conference paper at ICLR 2024
C ADDITIONAL DETAILS
C.I BASELINES
We note that while the proposed approach and NoMAD learn a dataset specific likelihood while aligning, i..,, Pog (w,, 2t). the other methods do not. Thus, for the other meth- ods, we first train the alignment function, Ja- Next, we fit log Pow. (wy D4) by maximizing
E ( [log Po.. (w,) 21)]
C.2 MODEL ARCHITECTURE AND TRAINING
For all experiments, the seq VAE encoder was parametrized by a bi-directionnl GRU with 64 hidden units. The latent dynamics were modeled as po(2, x+-1) (NJo(Ere11, where So two-layer MIP with a width of 256 and tanh activations.
As spike data tends to be very nigh-dimensional, space. Instead, when training the seq VAE we also trained an embedding function that projects the spikes down 64 dimensions before being passed into the encoder. The embedding function two-layer MLP with width of 64 and relu activations.
avoid learning alignment in the original
We used Adam optimize the model. Unless stated otherwise, we used weight decay of 10-4 and a learming rate of le
C.3 EXPERIMENTS
EVALUATION METRICS AND COMPARISONS
We evaluated all methods on their reconstruction RMSE and the normalized r². We also report the forecasting performance using the pre-trained model. After aligning the datasets with each method, use T time points to sample latents from the encoder, To:7-1 and subsequently generate fr. using the dynamics. This is decoded to obtain the predicted observations fr.r+k- The rī is then computed as:
M (ve 1.) Ei (g)
r -
where ý is the mean activity during the trial, and M is the number testing samples. In the synthetic experiments, we additionally train likelihood function after learning the alignment for methods that do not optimize for reconstruction. We have not performed an extensive hyperparameter search for the reported methods and instead followed the values recommended in the papers and/or the code.
ADDITIONAL DETAILS
Van der Pol oscillator. We used an embedding function in this experiment with d3 = 64 since the number of simulated neurons for datasets was over 200. In this case, directly align to the intermediate embedding except for the approaches that directly work in the observation space (Cycle-GAM) Procrustes). We used an MLP with dh 64 align new observations our approach and used 10-step prior during optimization.
Lorenz Attractor In this experiment, we parameterized g as linear function, and used 1000 trajec- tories for training all the alignment methods. Due to the chaotic dynamics in the system, we found that using a small, stochastic value the k-step prior was better than using large As a result, we randomly sampled & = [1 during training. We used a weight decay of for the optimizer.
Motor cortex When training the reference model from neural recordings with different stimulus conditions, we included the stimulus input when passing data through the embedding function, as well as the prior dynamics. In order to encourage the model to learn smooth dynamics, we randomly jittered the spike counts in [- 2] bins during the first half of training. Additionally, the weight decay parameter for the Adam optimizer was set to 10 2 We did not include the stimulus input when training the alignment function for any method.
15

---

Published as a conference paper at ICLR 2024
D ADDITIONAL FIGURES
A.
B.
Forecasting
performance
predictin
NoMAD ERDIf Oue
0.6
10
20
30
40 50
prediction step
100
200
Figure 6: Cross-moda alignment. Aligning Gaussian observations to Seq trained on spiking data A. Example trajectories of true rates (top) and reconstructed and predicted trajectories after aligning. B. k-step prediction performance using the pre-trained model after aligning.
A.
g(w)
y
W
Ours
True
100
M
www
W
M
WV
WV
-10
0
400
400
400
400
B.
(ours)
(true)
0.15
0
-0.15
W
W
Figure 7: Lorenz attractor with paired samples. Example units with the same initial condition for w and the aligned s(w) obtained using our approach, as well as the optimal alignment. B. The optimal g and the one recovered from our method.
A. Source latents
B. Our approach
ERDIFf
NoMAD
Cycle-GAN
Procrustes
Croas sess on
%
Cross subjeeci
Figure 8: A. Projected latents from the source data used to train the reference VAF. Each trajec- tory corresponds to the stimulus conditioned mean. B. Inferred latents after aligning session from the same subject (above) and across subject(below) projected on the same space.
16


Reviews:
[{"review_id": "2NOEwOENoK", "reviewer": "Reviewer_MYCT", "summary": "Stable and Effective inference models are crucial for decoding neural recordings, yet the need to train new models for each dataset due to variability is computationally demanding and inefficient. This is thus an interesting scientific question. This study introduces a novel alignment method that applies learned dynamics to new data, enabling the reuse of pre-trained models and facilitating the sharing of generative models across different distribution settings. The method's effectiveness is demonstrated by using a seqVAE trained on monkey behavior datasets, underscoring the significance of low-dimensional neural representations and offering a new perspective on handling the neural variability between sessions.", "questions": "Please consider the things listed in the “Weaknesses” section.\nAlso please consider providing information regarding any potential future improvements.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper focuses on an insightful and scientific research question: alignment of neural recordings. Since the generalization ability of models is a great concern.\n2. The paper's words and figures are well-written and easy to follow.\n3. The proposed method is analytically tractable and is with good theoretical guarantees.", "weaknesses": "1. The baselines in the experimental part is too few, just the recently proposed SOTA method ERDiff [1]. There are many classical methods like [2] and [3].\n2. The spatio and temporal structure has already been noticed by the SOTA method ERDiff, and ERDiff also employs a generative model (score-based model) for alignment. Thus what's the new motivations and insights of your work?\n3. There should be more experiments and empirical results to support your method.\n\n[1] Extraction and Recovery of Spatio-Temporal Structure in Latent Dynamics Alignment with Diffusion Model.\n[2] Stabilizing brain-computer interfaces through alignment of latent dynamics.\n[3] Robust alignment of cross-session recordings of neural population activity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ns1Siyi5Mb", "reviewer": "Reviewer_Jh7W", "summary": "The authors present a novel algorithm for leveraging pre-trained seqVAEs for fitting neural recordings. This algorithm builds on the assumption that the neural data in question share the same neural dynamics as the ones used to train the pre-trained model. The two key components to this algorithms are: (1) learn a new observation model from latents to observations and (2) _implicitly_ learn an alignment function from the new observations to the old observations. The authors validated their algorithm on both synthetic data and a monkey reaching dataset.", "questions": "* I assume $q_\\phi(x_t, x_{t-1}|y_{1:T}) = q_\\phi(x_t|y_{1:T}) q_\\phi(x_{t-1}|y_{1:T})$ in equation 3? Might be worth clarifying.\n* Is the equation supposed to have $p_\\theta(x_t|x_{t-1:t-k})$ instead of $p_\\theta(x_t|x_{t-1})$?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- originality\n    To the best of my knowledge this is a new approach for leveraging pre-trained generative models for fitting new data.\n- quality & clarity\n    - The paper is well-motivated and clearly written, notwithstanding some typos here and there (e.g., missing \"In contrast, [our method] does not require...\" in last sentence of section 2.)\n    - The experiments are relevant and convincing.", "weaknesses": "* I am not sure I follow why Proposition 1 implies good alignment? If the space of alignment function that the authors are trying to learn is linear (i.e., $g_\\theta(w) = \\theta w$), isn't it always the case that any $\\theta_\\star$ will be some linear offset away from the optimal alignment? Is this offset $B$ supposed to be small somehow? \n* The choice of alignment function $g_\\theta$ appears crucial to this paper, but there is very little discussion/evaluation about different choices. Specifically, if $g_\\theta$ is a point-wise nonlinear function approximator $g_\\theta(w_t)$ the authors are implicitly assuming that the latent dynamics are not just qualitatively similar, but **identifcal** between the new and old observations. However, if we make $g_\\theta$ to flexible (e.g., another full bi-directional RNN), then we are effectively re-learning the encoder.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2QeswtTasR", "reviewer": "Reviewer_H5ee", "summary": "Here, the authors consider the problem posed by the existence of multiple data sets of neural time series from the same or similar tasks, asking whether it is possible to transfer learned latent dynamics from one to the other. They propose a model in which a sequential VAE is first trained on a large data set to establish a latent space ${x}_{1:T}$ (assuming a linear decoder). On a new data set, these dynamics are frozen, a new (linear) decoder is retrained, and the encoder is reused by learning a nonlinear mapping $g(w)$ from the new observations $w$ to the space of the training observations $y$. This requires only the trained encoder and does not require labeled examples. This method is compared with several other alignment proposals, where it produces both quantitative and qualitative improvements.", "questions": "- How flexible is this setup to the specific architecture chosen? It's mentioned that using an ELBO that measures log predictive probability several steps ahead is important to achive a good embedding, but it's not entirely clear to me why.\n- How complex are the learned dynamics in cases that work versus don't work? The monkey reach data typically have rotational dynamics. Do you see anything more complicated in other data?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The problem of aligning multiple data sets across animals and experimental sessions is a major one in neuroscience. An ability to reuse or amortize model training across these would be of significant benefit.\n- The method is theoretically well-motivated and fairly flexible. It doesn't appear to require a particular architecture (apart from the linear decoder, which is a limitation of data availability as much as anything).\n- The approach appears to produce real qualitative improvements in the learned embedding (Figures 3 and 4).", "weaknesses": "- The approach uses a fairly strong assumption that the latent dynamics really are shared across data sets, which all but implies a shared task setup. That is, it doesn't appear to be the case that a sufficiently large task-free data set in one mouse would facilitate embedding of mice performing a task-based behavior. It would be surprising if true, but this weakness should be acknowledged, since this limits the range of applicability.\n- Given the large literature on data alignment/domain adaptation both within and without neuroscience, it's a bit surprising that there are only two approaches compared in Table 1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "T1gwXCA0PL", "reviewer": "Reviewer_nsT1", "summary": "The authors propose an unsupervised alignment approach that can apply a pre-trained model on other neural recordings to the new dataset. They mathematically prove their approach is optimal up to scaler in a simplified case. They validate their methods on a simulation and a real dataset.", "questions": "1. how do you model the latent dynamics pθ(xt|xt−1)? Do you add any smooth constraints? In figure 2b, it seems that the results of the proposed approach are much smoother than ERDiff. \n2. what is FA + LDS method that you compare to in figure 2? Can you list a bit more details on how it's implemented? \n3. figure 2e, why does the mean k-step r^2 for Re-training decrease as the number of training samples increases? Intuitively the r^2 should increase as there are more training samples.\n4. what are the results showing in table 1? Are they results for the real neural dataset? I found the r^2 number in table 1 doesn't match the number in figure 4.\n5. in table 1, why the k-step MSE standard deviation is smaller than the stddev of MSE for ERDiff and the proposed method? \n6. in figure 4, why the proposed method performs better than ERDiff for poisson observations but not for the smoothed observations, especially Figure 4b (Top, left)?\n7. The proposed method assumes that prior dynamics p(xt | xt-1) is fixed, and the encoder q(x | y_embed) is fixed. I wonder if the authors have checked these two assumptions? For example, if you jointly train a model on two datasets and learn the latent dynamics and encoder, will they be similar to the results of training separate models on each dataset?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The authors propose a new unsupervised method to solve the transfer learning problem in neuroscience by re-using the temporal dynamics structure. They also have mathematical proof of the optimality of their method in a simplified case.", "weaknesses": "The experiment results are limited. The paper only shows results on 1 real dataset. There are some unclear parts in the paper, which I listed in the Questions section below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)