Flow-field inference from neural data using deep recurrent networks
Authors: Timothy Doyeon Kim
Thomas Zhihao Luo
Tankut Can
Kamesh Krishnamurthy
Jonathan W. Pillow
Carlos D Brody
Venue: ICML 2025 poster (2025)
Keywords: computational neuroscience
dynamical systems
neural population dynamics
interpretability
Primary area: applications->neuroscience_cognitive_science

Abstract: Neural computations underlying processes such as decision-making, working memory, and motor control are thought to emerge from neural population dynamics. But estimating these dynamics remains a significant challenge. Here we introduce Flow-field Inference from Neural Data using deep Recurrent networks (FINDR), an unsupervised deep learning method for inferring low-dimensional, nonlinear, stochastic dynamics underlying neural population activity. Using spike train data from frontal brain regions of rats performing an auditory decision-making task, we demonstrate that FINDR performs competitively with existing methods in capturing the heterogeneous responses of individual neurons. When trained to disentangle task-relevant and irrelevant activity, FINDR uncovers interpretable low-dimensional dynamics. These dynamics can be visualized as flow fields and attractors, enabling direct tests of attractor-based theories of neural computation. We suggest FINDR as a powerful method for revealing the low-dimensional task-relevant dynamics of neural populations and their associated computations.

Flow-field inference from neural data using deep recurrent networks
Timothy Doyeon Kim
Thomas Zhihao Luo't Jonathan W. Pillow
Tankut Can 34 Kamesh Krishnamurthy" 15 Carlos D. Brody 16
Abstract
monkeys, humans, and even nematodes (c.g., Kato al. (2015); Nieh et al. (2021); Churchland et al. (2012); Pan- darinath et al. (2015); Safaie et al. (2023)). One influential premise in systems neuroscience is that the low-dimensional dynamics on this manifold mediate the computations per- formed by neural populations (Vyas et al., 2020; Duncker & Sahani, 2021).
Neural computations underlying processes such as Docision-making working memory, and mo- tor control are thought to emerge from neural population dynamics. But estimating these dy- namics remains a significant challenge. Here we introduce Flow-field Inference from Neural Data using deep Recurrent networks (FINDR). an unsupervised deep learning method for infer- ring low-dimensional. nonlinear, stochastic dy- namics underlying neural population activity. Us- ing spike train data from frontal brain regions of rats performing an auditory decision-making task, we demonstrate that FINDR performs com- petitively with existing methods in capturing the heterogeneous responses of individual neurons. When trained to disentangle task-relevant and ir- relevant activity, FINDR uncovers interpretable low-dimensional dynamics. These dynamics can be visualized as flow fields and attractors, en- abling direct tests of attractor-based theories of neural computation. We suggest FINDR pow- erful method for revealing the low-dimensional task-relevant dynamics of neural populations and their associated computations.
Several methods- here referred to as "neural population dynamics inference methods"-han been developed to in- fer these dynamics directly from neural population activity measured from an animal performing a computational task. These methods typically make certain assumptions about the dynamies, either to facilitate inference or increase model capacity, For example, the dynamics are assumed to be autonomous (Duncker al., 2019), linear (Macke et al., 2011; Gao et al., 2016), switching linear (Linderman et al., 2017; Nassar et al., 2019; Zoltowski et al., 2020), determin- istic except at specific time points (Pandarinath et al., 2018; Kim et al., 2021; Keshtkaran et al., 2022), one-dimensional (Genkin et al., 2021), or high-dimensional (Pandarinath et al., 2018: Keshtkaran al., 2022).
While these methods capture neural population activity effectively. there are two key areas where improvements can be made. Here, we propose to make these improvements with a new method called FINDR (Flow-field Inference from Neural Data using deep Recurrent networks; Sec- tion 2).
1. Introduction
How do neurons work together in large populations to solve a task? Experimental evidence suggests that neural popula- tion activity lies on a low-dimensional manifold across mul- tiple brain regions in different species, including rodents,
1) We relax the assumptions on dynamics by using a gated multilayer perception (MLP) FINDR's dynamics model (Kim et al., 2023), enabling it to learn wide range of com- plex dynamical systems, including those that are difficult to learn using existing methods. In Section 3.1, we construct an example of a dynamical system that is difficult to learn with existing methods, and show that FINDR can accurately infer its dynamics from simulated spike trains generated by this system.
Princeton Neuroscience Institute, Princeton, NJ Preseent ad- dress: Allen Institute & University of Washington, Seattle, WA "School of Natural Sciences, Institute for Advanced Study, Prince- ton, NJ *Present address: Department of Physics, Emory Uni- versity, GA 'Joseph Henry Laboratories of Physics, Princeton University, Princeton, NJ 'Howard Hughes Medical Institute, Princeton University, Princeton, NJ. Correspondence Tim- othy Doycon Kim <tdkim @ princeton.edus: edu>,, Carlos D. Brody <brody @ prinectoonnncc.........
2) The complex response patterns observed in real neurons may be decomposed into task-relevant and -irrelevant com- ponents (Rigotti et al., 2013; Gallego et al., 2018; 8; Sani et al., 2021). However, many currently available methods do not explicitly take this into account. FINDR distinguishes be- tween the task-relevant and -irrelevant dynamics, and infers
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Flow-field inference from neural data using deep recurrent networks
them separately. This allows us to more easily interpret the inferred dynamics, and demonstrate a link between the dynamics and task computation. We show this by applying FINDR to spike trains from rat frontal cortical regions in- volved in decision-making (Sections 3.2- 3.3). We find that task-relevant dynamics in these regions are low-dimensional (Extended Data Figure 3), and even though we can capture the heterogeneous responses of individual neurons well us- ing the existing approaches when we use sufficient number of latent dimensions in these models, we show that promi- nent alternative deep learning-based method infers dynamics that are inconsistent across different training and test splits of the same dataset, whereas FINDR infers dynamics that are consistent and also low-dimensional (Section 3.3).
the latent space R equivalent to the distance and angle in the inverse-softplus rate space TR-V (see Section 2.2). Also, by learning task-irrelevant time-varying representation d,, we encourage the now-dimensional representation 2; to be more task-relevant, and therefore more interpretable. The observed population spike counts Ye at time bin t of width Al are modeled by y: ~ Poisson( (A)).).
To achieve 2), FINDR models the dynamies of the latent representation 2, stochastic differential equation (SDE) discretized with the Euler-Maruyan method:
At 21 = =--1+ J ((z)-1, u.) + Vit
T Ev.
(2)
Because FINDR can represent neural population activity in low dimensions, even low as two or three dimensions, we can explicitly visualize the flow field (or the velocity vector field) underlying neural population activity. We show that when we visualize the flow field formed by the frontal cortical neural population during sistently see two slow points, one associated with stimulus favoring leftward choice and the other favoring rightward choice (Section 3.3).
Here, T is a fixed time constant of the SDE, and u, E R is the external input to the system at t, typically a set of task variables that the experimenter has control over. The drift function μ. describes the "rules" used by & to evolve over time. Therefore, μ represents the dynamics used by the neural population while the animal is performing the task. We can visualize the dynamics by plotting the flow field (or the velocity vector field) of (Figure 1).
maaiiig, we con-
While we showcase our method with a neural population dataset in decision-making we expect this method to be applicable to a wide range of neuroscience datasets.
In Equation (2), we model potential noise that may be present in the dynamics, represented as É, ~ NO, E). Here, the noise covariance Σ is a diagonal matrix, with each element a learnable parameter in the model. Noise &, can be a crucial component in modeling neural computation. As an example, in perceptual decision-making, we observe that even when an animal is presented with identical stimuli, the animal's 's choice behavior can vary from trial to trial. If we are given spiking observations from neural population that represents the animal's choice, we can model variability in the animal's behavior with noise in the dynamics of the neural population. Poisson noise does not model noise in the dynamics and is not sufficient to capture this variability.
2. Methods
Neural population dynamics inference methods 1) compress the activity of a large population of neurons at time t to abstract -iw-dimesioonl representation, and 2) learn the "rules" of how this representation evolves over time. We call these rules the dynamics of the system (Figure la).
Given population spike train data from N neurons an animal performing an experimental task, FINDR achieves 1) by inferring their underlying neural population firing rates X E IR-y at time step 1 (= 3, ..., T, where T is the total number of steps taken in a trial) using a low- dimensional task-relevant latent representation 2, E RR- This map from z, to the firing rates de is given by
The drift function in Equation (2) is parameterized by
u(z, u) (G( 1- - u)z,
(3)
A softplus(Czz + de),
(L)
where a is the sigmoid function that acts element-wise, and F and G are multilayer perceptions (MLPs). Function parameterized this way is practically more expressive and trainable compared models without the gating network af (G(z, u)) (Kim al., 2023).
where C' E NXL is a loading matrix and de E R is some task iis--irrelevant time-varying bias. Here, the softplus nonlinearity prevents the firing rates from being negative. Instead of Equation (1), we could have used a more general mapping X = softplus(V,, , (2.)), where V is differen- tiable map with parameters K, and could, for example, be a deep neural network. However, for simplicity and for interpretability, we confine our map to be affine. Notably, if C' is oothoonl,,,,,,,, this makes the distance and angle in
FINDR achieves 1) and 2) by finding the optimal parame- ters for the networks F and G in the drift function μe, the diagonal noise covariance Σ, the loading matrix C, and the time-varying bias d, that best capture the observed ye. We will denote these parameters as θ (0r,0G: Σ, C.drr)
2

---

Flow-field inference from neural data using deep recurrent networks
2.1. FINDR optimization
field μ02z, u). If [ is too high, the inferred latent trajectory žir becomes more irrelevant the observed neural activ- ity, but it becomes highly likely that it is generated from our inferred flow field pzz, 24. It has been observed that using B > 1 can help arrive at interpretable latent representations (Higgins et al., 2017: Burgess et al., 2018). We let 3 = 2 to put slightly more weight on the vector-field inference at the cost of less accurate reconstruction of neural activity. However, as we will show in Section 3.2, we find that under many conditions, FINDR outperforms existing methods in reconstructing neural population activity.
To obtain the optimal parameters O", we minimize the neg- ative logeleielhoodd of the spike trains given O:
L = - log pe (VI.T).
(4)
We approach this problem by first optimizing for di.r and then optimizing the rest of the parameters 0 = O) d1.T. To
obtain the optimal d, = [d df(") :...; d :...; d J. we fit a linear basis function model (Bishop, 2007) for each neuron n with
We minimize C by training a sequential variational autoen- coder (VAE), where we model pe (ya, z2z de) in the first term of Equation (6) as y, ~ Poisson (A/x) with X, given by Equation (I), We model (1-11 Ue) in the second term with Equations (2-3). Finally, model (=(2)|=(_1,e. with
D. (1)
d (1)
(5)
where ip- {pp1 PP2, .... (DD.) } is set of raised cosine basis functions (Pillow et al., 2005; 2008; Park et al., 2014). In this setup, we are approximating Equation (4) as a problem of minimizing the mean squared error (MSE) between the observed vi and d!") for all time steps i and all neurons n.
v = -1-11-1 + -1 Un ex) +
At
T Eo
This time-varying bias dl(T) is meant to capture fluctuations in an individual neuron n's firing rate within and across
v(z u, e) - ((G(,, u. e)) @ + f(z, u, e)
where F and G are MLPs (separate from F and G). and e is the output at every time step t of a bidirectional gated re- current unit (GRU) (Cho et al., 2014) that gets the sequence [u:: at & as its input. When we model the variational posterior q this way, the KL divergence term in Equation (6) has an analytical form. Note that in principle, having u in addition to e as input to is redundant. Nevertheless, we choose to have this parameterization is empirically superior to the model without u in terms of performance. Figure 1b presents a graphical overview of the FINDR model, and Figure Ic shows schematics of the model architecture
trials of the task that are not directly relevant to performing the task itself. See Appendix 1. for details.
After fitting this linear basis function model to obtain de, we proceed to optimize 0. Minimizing Equation (4) directly with respect to f can be computationally expensive. There- fore, we instead compute an approximate upper bound of C, which we denote as É (see Appendix A.1.2 for derivation; Kingma& & Welling (2014); Chung al. (2015); Krishnan et al. (2017)):
an extra imput to y because we find that
T
c =
log po (Yr (2,, di)
t=1
We train FINDR with mini-batch gradient descent with warm restarts (Loshchilov & Hutter, 2017), where the gra- dient of the loss E is taken with respect to both 0 and the variational parameters ø (which include parameters of the bidirectional GRU, F. and G) (see Section A. for de- tails). The gradient is obtained by backprogation through time (BPTT).
].
- DKL (90 ze 21
where 2, ~ 1,,- and et is some time- varying representation that summarizes the entire sequence [u1:r: V. q is the variational approximation of the pos- terior p(zz 1,,,,, uIr V,.r), with & being the variational parameters of q. 2 is a sample from the variational posterior q.
2.2. Identifiability and interpreting the learned latent representation
Minimizing Ć has two effects: optimizing neural activity reconstruction (due to the first term in Equation (6)) and inferring the low-dimensional flow field that generated the neural activity (due to the second term Equation (6)). The coefficient / in front of the second term is important hyperparameter that determines the trade-off between the accuracy of the reconstructed neural activity and the discov- ered flow field, When ,3 is too low, the reconstructed neural activity may be accurate, but it becomes unlikely that the inferred latent trajectory ह is generated from the flow
Following Wang et al. (2021), we define latent variable z as identifiable if x # zg implies that pe(y # d) * pe (o 22, d) in the first term of Equation (6). Since we use linear projection C for pe (3/ po d), this is satisfied if C
injective. Note that is identifiable only up to linear trans- formation because for any invertible A E IR L we can re-write Cz such that Cz = CAA-'Z 2 = C'z', where C' = CA and x' = A 11 z. Therefore, for all models with po(y) z, d) = softplus (C'z + d) as in Equation (1),
3

---

Flow-field inference from neural data using deep recurrent networks
a
dynamics inference in reduced dimensional latent space
spike trains
firing ratos
time
neural poculation state space
X
nouron yd
(1)
Eat
latent trajectory
nouron
y
(2)
nouron 3
(09
A
(3)
law-D latent
latent dynamics (fow field)
nouron N
y
A(
(15)
(NN
space
A ())
(N)
5
6
4
1
6
A
(5) ...
A
(5)
b
sequential VAE model
C
model architecture
spike trains
minimize loss for noural activity reconstruction
model-inferre frinng rates
encoder
decoder
Icentity
Iafant trajectory
firing
IIIII
latent trajectory
kofpllu
rates A
BIGRU
e,
A
Ae
1
encoder
decoder
spite loins
extomal inpurs
noiso
ying
minimize loss for flow-fiedd infarence'
Bo-1
nncural activity reconstruction
external inputs
minimize ya
latent dynamics (flo field)
Intent dynamics
flow-field inference
minimize weighted MSB/
Figure 1. graphical description of FINDR. a, FINDR infers the firing rates À underlying N-dimensional spike trains y and compresses A to T-dimensional representation = (L & N). Based on the trajectory formed by z over time ("latent trajectory" over f + tB + to + tD), FINDR infers the dynamics of z, more specifically the "flow field" showing the rules of how x moves in the latent space. b. FINDR is a sequential VAE that models the probability density functions in the loss 2 in Equation (6). FINDR's s encoder network models the variational posterior go (z 2t-1, ec). and the decoder network models p(y, de). An additional network models p(z) 22./2z-1.0 u)) to infer the flow field. FINDR minimizes loss for neural activity reconstruction (the first term in Equation (6)) and loss for flow-field inference (the second term in Equation (6)). c, The encoder network (in blue) processes data with bidirectional GRU and infers the single-trial latent trajectories using gated neural SDE discretized with the Fuler-Maruyama scheme (Equation (7)). The decoder network (in red) transforms the latent trajectories to firing rates (Equation (1)). FINDR reconstructs the observed neural activity using the firing rates from the www-immesiionll latent trajectories and the time-varying bias. At the same time, FINDR infers the most likely flow field that generated the inferred latent trajectories using a gated MLP (in green).
including FINDR, and those with rectifying nonlinearity other than the softplus, we perform singular value decom- position (SVD) on C' = USVT, where U E N is a semi-orthogonal matrix, S E R is a diagonal ma- trix with its entries populated by the singular values and L. is an orthogonal matrix. Then, we set C + U and = + SVV post-training. This makes the distance and angle in the new latent space R and the distance and angle in the inverse-softplus rate space N the same, i.e., ||xx||2 = = C Cz = =x = || for all z. We then rotate the latent space of z so that the axes are its prin- cipal components. The latents z after this post-training transformation (effectively equivalent to model with a semi-orthogonal C) are thus identifiable up an orthogo- nal transformation. We did not place soft or hard constraints on C' to be semi-orthogonal during training, as this wors- ened performance as reported in other contexts (Vorontsov et al., 2017). See Appendices A.2 A.2 for more details.
3. Experimental Results
3.1. FINDR can accurately infer approximately continuous attractors in synthetic neural populations
V e L
To examine the validity of FINDR, we generated simulated population spike trains from a known low-dimensional dy- namical system and checked whether FINDR can infer la- tent dynamics that are similar to the ground truth. The low-dimensional dynamical system we use is inspired by the "n- flip-flop task" (Sussillo & Barak, 2013). In this task, the system receives transient pulse inputs from dif- ferent channels and needs to memorize the value of the most recent pulse in each channel. It is known that dynamical system can use attractors to solve this task and that the at- tractor structure of the system reflects the statistics of the pulses (Kim et al., 2023). For example, in Figure 2a, we let
4

---

Flow-field inference from neural data using deep recurrent networks
the system memorize the pulse values from two channels, where the pulse value in channel 1, C1, and the value in channel 2, C22 are constrained to satisfy 1 S V + S 2. For the system to have robust memory of the pulses in the two channels, it should form a 2-dimensional continuous attractor that has the shape of a disk (Kim et al., 2023).
rectangular attractor, the width and length of the rectangle were roughly preserved in the inferred latent representation (Extended Data Figure 1c). We find similar results for 3 3-dimensional dynamical system with a continuous attractor of the rectangular prism shape (Extended Data Figure d g). These results suggest that FINDR accurately infers latent dynamics with attractors of different geometries and dimensionalities. while preserving distance in latent space.
We simulated our data from 500 different Poisson spiking neurons, with the task inselinteeevant dynamics being a constant bias with around 5 spikes/s of firing rates (Figure 2b). The latent trajectories zim trace, with some small noise, the optimal solution to the task, given external inputs Unv (Fig- ure 2a). Then we asked whether FINDR, given the neural population activity and external inputs Uaee can reconstruct a 2 2-dimensional disk attractor (Figure 2c-d) needed to solve this task. This task was constructed such that methods that assume autonomous dynamics (e.g., (Duncker et al., 2019)) or linear dynamics (e.g.. (Macke et al., 2011)) cannot dis- cover the disk attractor. Methods assuming switching linear dynamics (e.g., Linderman et al. (2017)) may have difficulty approximating the disk with a few interpretable discrete states because the latent space typically partitioned lin- early in these models. Methods learning high-dimensional dynamical system (e.g., Pandarinath et al. (2018)) may have difficulty inferring the true latent dimensionality (in this case L = 2), and may use dynamical features not present in the data (Sedler et al., 2023).
3.2. FINDR performs competitively against existing methods in capturing real neural population responses
To evaluate FINDR's performance relative to existing meth- ods in predicting the heterogenous responses of individual held-out neurons, we applied FINDR, switching linear dy- namical systems model (SLDS; Linderman et al. (2017)). re- current switching linear dynamical systems model (rSLDS; Linderman et al. (2017)), autoLFADS (Keshtkaran et al., 2022: Sedler & Pandarinath, 2023), and Gaussian Process Factor Analysis model (GPFA; Yu et al. (2008)) dataset comprising 67 choice-selective neurons, selected from larger population of 464 simultaneously recorded neurons from dorsomedial frontal cortex (dmFC) and medial pre- frontal cortex (mPFC) of a rat engaged in a decision- maaiing task across 448 trials (Luo et al., 2023). On each trial, the rat listens to two simultaneous, randomly timed auditory click trains played from loudspeakers on its left and right. At the end of the stimulus, it turns to the side that had the greater total number of clicks for water reward. The spike trains and the auditory click times given to the models were aligned to the stimulus onset (Figure 3a; for a more detailed description of the task and the selection criteria for neu- rons, see Luo al. (2023)). The spike trains and auditory click times were binned using bin width of 10ms, and the binned click times from the left and the right speakers were provided to the model as a 2-dimensional external input u.
To identify whether FINDR correctly captures the true la- tent dimensionality (L = 2) of the population spike trains, we trained multiple FINDR models, each assuming differ- ent latent dimensions (T. 1,2, 2,.,... 6), on 600 trials of the simulated population spike trains. For each of the FINDR models assuming different latent dimensions, we did grid search over the hyperparameters (see Section 1. for de- tails) and found the best-performing model by evaluating the normalized log-likelihood score (Pei et al., 2021) on 200 validation trials not used during training. We find that the loglliie,,oood, evaluated on 200 test trials (which are separate from the validation trials), saturates around L = 2 (Figure 2c). Consistent with this result, we also find that when we do principal component analysis (PCA) on the FINDR-inferred latent trajectories z, we see that two prin- cipal components (PCs) are sufficient to explain more than 99% of the variance in each model that assumes different L (= 2, ..., 6) (Figure 2d). Furthermore, when we project the flow field inferred from FINDR onto the first two PCs, we find an approximate disk attractor across FINDR models assuming L = 2,....., ..., 6 (Figure 2d).
We held out 13 neurons (about 20%) from this dataset, and partitioned the dataset into 5 different folds, each containing a subset of trials in random order. Following Pei et al. (2021), we held out 20% of the neurons by supplying the encoder with the remaining 80% of the held-in neurons. Then, the decoder in Equation 1 reconstructed all neurons. We used 3 of these folds for training, 1 fold for validation, and the remaining 1 fold for testing. We evaluated the 5-fold cross-vale validated likeliodl of held-out neural activity to measure model performance. This is a standard metric for evaluation known as co-smoothing (Macke et al., 2011; Pei et al., 2021). When we assess model performance on 5- fold cross-validated co-smoothing across latent dimensions from L = 1 to L = 6. we find that the log-likelihood for FINDR is consistently higher than existing models when L < 3 (Figure 3b; left). We also use the 5-fold cross- validated coefficient of determination (R²) to evaluate the
When we do similar analyses with simulated population spike trains generated from a 2-dimensional system with a continuous attractor that has the shape of a rectangle, we find that FINDR discovers these structures (Extended Data Figure 1). Furthermore, for the dynamical system with
5

---

Flow-field inference from neural data using deep recurrent networks
a
external Inputs Muse
b
C
latent dynamics Sino
Cinmit ~ (.)
durr
chamnel
Arrms wallppuueCunnnann dans)
normalized log sill-llllhhood
Firast Poiesion Mexaz)
1.2
(on test test data)
1.0
channel 2
1
neuron N
0.8
0-
0.6
1
5
6
0.0
0.5 time (9)
1.0
0.0
0.5
1.0
of latent dimensions (L)
(s)
assumed by FINDR
d
L =
L
4
variance explained
1.0
0.8
0.6
3
1
2
3
4
5
1
2
3
principal components (PC)
(PC 2)
2, (PC
Figure 2. FINDR infers approximately continuous attractors that reflect how the synthetic neural population stores external inputs. a, An example trial with external inputs Ure each channel shown black, and the state trajectory of dynamical system that maintains the previous pulse value shown in gray. b, An example trial of the simulated population spike trains. c. normalized log- lieliiklioood evaluated on the test data as function of latent dimensions (I) assumed by FINDR. d, Principal component analysis (PCA) on inferred latent trajectories across FINDR models assuming L. = 2, 6 show that the first two PCs are sufficient capture most of the variance latent trajectories. The inferred flow fields projected onto the first two PCs are consistent across FINDR models assuming 3, 6. Gray lines represent latent trajectorics.
match between observed and model-predicted peristimulus time histograms (PSTHs) of held-out individual neurons (Figure 3b; right). We included further details model evaluation metrics in Appendix A.4.
right panel Figure for I 6). We repeat this proce- dure for SLDS, rSLDS, autoLFADS, and GPFA across dif- ferent latent dimensions, and find that FINDR consistently outperforms competing models under this metric (Figure 3d, Extended Data Figure 2). These results together suggest that FINDR can predict held-out neural population activity from latent dynamics with fewer dimensions compared to existing methods.
Figure 3c shows example neurons' activity averaged across trials sorted by whether the stimulus favors a leftward or rightward choice (evidence-sign conditioned PSTH), and the held-out prediction of this conditioned PSTH from FINDR with 6. The R22 between the observed PSTH and FINDR's 5-fold cross-vall volsdaaldd prediction is computed for each held-out neuron, and shown a histogram, with mark- ers indicating where the example neurons fall within this distribution (Figure 3c; with neuron being the median of this distribution and what is shown as the blue circle in the
3.3. FINDR reveals dynamically consistent latent representations across cross-validation folds
That FINDR can capture neural population activity using latent representation with lower number of dimensions compared existing methods allows inspect the nature
6

---

Flow-field inference from neural data using deep recurrent networks
a
b
diference (mode) -taselin))
mPFC
dmFC
normalized log ikkllhood
0.05
median lann evddanse-iign conditoree PSTH Rt
0.8
more dlikks from the left or right?
0.03
0.2
FINDR
0.02
nose in center
-0.2
beseline model
rSLDS
0.00
-0.4
sudiory clicks
1-11-8--
FADS
2
3
4 5
6
5 of dimansions
response
GPFA
latant dimansions
C
FINDR goodness-of-fit
neuron (best)
neuron 2 (median)
neuron 3 (worst)
13-
held-out neurons
spikess
R² 0. 96
R 0.80
F? 0 0.52
0
evidence-sigg
time from stimulus onset (9)
right evidianoe leftt evidence
observed 95% CI FNNDR test data prediction
conditioned PSTH
conditioned PSTH R2
L L=2
==3
L
L= 5
L = 6
Other model
evidence-aige
0 -1
1 -1
-1
0
-1
1
-1
FINDR evidence-sign conditioned PSTH R
SLDS
aulaLFFADS
(SLDS
GPFA
Figure FINDR outperforms SLDS. iSLDS, autoLFADS, and GPFA reconstructing the neural population activity of rat performing perceptual decision-
decision-making. making. a, Neurons from dorsomedial frontal cortex (dmFC) and medial prefrontal cortex (mPFC) were recorded while the rat listened to stream click trains from left and right sides operant chamber, and oriented the side that had more clicks to receive water reward. b, 5-fold cross-validated normalized log-likelihood difference score (Pei et al. (2021); also known as "co-smoothing": left) 5-fold cross-validated median evidence-sign conditioned PSTH R² (right) on held-out neurons across different latent dimensions for FINDR, SLDS, rSLDS, FADoLFADS, and GPFA. For SLDS and rSLDS, consider only the best best-performing model among models assuming the number of discrete latent states = (1, 2,,,,,,44 4). For LFDDS corresponds to the factor dimension, not the size of the generator recurrent neural network (RNN). AutoLFADS with L fails train. The baseline model here is defined constant firing rate model with the constant being the mean of the observed neural activity. c, FINDR with L = captures the complex trial-averaged temporal profiles of individual neurons in mPFC and dmFC. The goodness-of-fit is measured using the R d, We perform an analysis similar to for other models and compare the R² obtained from these models and the R² from FINDR.
of the learned representation in intuitive and interpretable manner. When we compute the R2 between the observed and FINDR-reconse recdstrecont PSTHs for all neurons in the dataset in Figure 3a, we find that two latent dimensions are sufficient to describe the data (Extended Data Figure 3). Therefore, we focus on L for analyses of the learned dynamics.
average, given enough time (4-5s), the latent state falls into the slow point in the upper part of the state space when the stimulus favors rightward choice and that it falls into the slow point in the bottom part of the state space when the stimulus favors leftward choice (Extended Data Figure 4). Furthermore, we find that the learned representations were similar for the FINDR model with the second-best set of hyperparameters, and across different random seeds for network initialization and the order in which the mini- batches of the dataset were supplied to the model during training (Extended Data Figure 7). All of our analyses with FINDR in this Section involved learning the task-relevant and -irreeeaant components separately (Section 2.1). We find that we get less consistent and interpretable dynamical
In Figure 4b, we explicitly visualize the flow field inferred by FINDR with I, = 2, and evaluate consistency in manner similar in spirit to Genkin & Engel (2020). We find that FINDR learns dynamical representations that are consistent across 5 different folds (i.c., different training and test splits of the same data). Specifically, across all folds, we find that they consistently have slow points. We see that

---

Flow-field inference from neural data using deep recurrent networks
representations when we do not learn the task-irrelevant component (Extended Data Figure 5).
find that sensory inputs perturb dynamics roughly along PC in the latent space of FINDR, but this would be difficult to know using CEBRA.
Do we find similar slow-point structures across folds for alternative dynamical models? We evaluate this in autoL- FADS, which performs similarly to FINDR in predicting neural responses (Figure 4a). For each fold's autoLFADS, we ran the trained generator forward in time for 5s, starting from the initial conditions inferred from the encoder. We found that while the majority of autoLFADS states reached approximate steady-states by 5s, they did not form two clus- ters as would be expected from bistable attractors, both for folds and 2 (Extended Data Figure 8).
Similar to the analysis in Figure 2d, when we fit FINDR models that assume different latent dimensions (I. 2, 3, ..., 6), and project the inferred latent trajectories to the first two PCs, we also see consistency across dimensions (Extended Data Figure 6). This suggests that FINDR can discover consistent and interpretable task-relevant latent dynamics.
Our analyses suggest that FINDR representations are consistent across folds- not only topologically but also geometrically-and reveal dynamical consistency, specif- ically two slow points associated with left/right choices. Among all methods tested, only FINDR achieves both 1) strong performance on neural data, and 2) discovery of con- sistent low-l) dynamical representation, with interpretable slow points. While we find empirical evidence that FINDR can discover consistent representations, consistency is not theoretically guaranteed and should be verified empirically on new datasets.
To quantify consistency across folds, we sorted single-trial trajectories by evidence sign and computed the trial average of each group. Then, we calculated Pearson's 's [r] of these trajectories between fold 1 and 2 for each latent axis and took the average of |r| across the axes (we will refer to this metric as (Ir))) Note that the axes here are defined by the principal components of the latents. With this, FINDR folds were consistent by 0.99, while autoLFADS folds were consistent by 0.53. The consistency score for FAtD increased to 0.99 if we linearly transform autoLFADS fold 1 to match fold 2. However, this linear transform stretches and rotates the latent space, so the distance in latent space is no longer preserved in the neural space (ignoring soft- plus). This suggests that while autol. is topologically consistent, it is not geometrically consistent as in FINDR. To further quantify dynamical consistency, i.e., to quantify whether we consistently find two slow points across folds, evaluated whether the distribution of the fold-1 states in approximate steady-state match the distribution of the fold-2 states. We affine transformed the autoLFADS latent trajectories from 4-5s in fold to match those in fold 2, and when we computed (Irl). we found that they were not consistent ((rl)) = 0.22). This suggests that even when the autoLFADS factor trajectories across folds are topolog- ically consistent ((rl)> = 0.99), this does not guarantee that the underlying dynamics that generated the trajectories by LutoS are consistent. Using a similar procedure for FINDR, we found ([rl)- = 0.94, consistent with the visualization in Extended Data Figure 4.
4. Discussion
With recent advances that enable recording from a large population of neurons (Jun et al., 2017: Steinmetz et al., 2021; Stringer et al., 2019; Siegle et al., 2021), many meth- ods (e.g., Macke et al. (2011); Linderman et al. (2017); Pandarinath et al. (2018)) have been introduced to infer the dynamics underlying the recorded population activity (e-g.. Vinograd et al. (2024); Liu et al. (2024)). Benchmarks (Pei et al., 2021; Versteeg et al., 2025) have also been introduced to facilitate rigorous and standardized comparisons between these methods.
FINDR improves on the previous methods by enabling flexi- ble inference of stochastic, nonlinear latent dynamics, while also separating task-relevant and -irrelevant components. In population spike train dataset from frontal cortex of the rat during perceptual decision-making these two features allowed FINDR to capture the complex, heterogeneous neu- ronal responses, and discover dynamically consistent and interpretable latent representations that we can easily relate to the task computation.
We additionally performed analyses with a prominent latent representation learning method, CEBRA (Schneider al., 2023), and found that the parts of the state space traversed by the trajectories depend on evidence strength (Figure 4d). We also saw that the first two folds were topologically con- sistent by (Irl) = 0.99. However, CEBRA and related methods (c.g., Chen et al. (2025)) do not learn dynamical representations (i.e., Equation (2)), and, unlike FINDR, it is difficult to perform fixed-point analysis on the latents or evaluate dynamical consistency. Another key distinction between FINDR and CEBRA is that, for this dataset, we
FINDR is general dynamical inference framework that goes beyond the particular use case shown here (perceptual decision-makinn
and supports broader application to other types of neural computation, such as orientation tuning in visual cortex (Hubel & Wiesel, 1959) or rotational trajec- tories in motor cortex during movement (Churchland et al., 2012). The precise dynamics underlying orientation tuning motor control remain unclear (Khona & Fiete, 2022). FINDR provides new tool for investigating dynamics in
8

---

Flow-field inference from neural data using deep recurrent networks
b
fold1
fold 2
fold
fold 4
fold
0.6
10
pmo 054
FINDR
naurnA
(PC 2)
-0.6
2, (PC 1)
-.5
0.5
evidonce sign cood AR
autoL FADS
0.6
(FINDR 1-2 FADoLFADS (-22)
PC
overages sorted by evidence strength
-0.6
PC 1
CEBRA
PC 2
-0.6
PC 1
Figure 4. FINDR reveals slow points in the latent dynamics of frontal cortex during perceptual decision decision-making, consistently across different training and test splits of the same data. a, FINDR with L - 2 and autoLFADS with L = 20 are similar in how well they capture the heterogeneous responses of individual neurons, when evaluate their performance using the 5-fold cross-validated evidence-sign conditioned PSTH n22. b, To show how the decision process depends on the input, trials are sorted by their evidence strength, and we show the trial-averaged trajectories as colored lines. innee-infered flow fields consistently reveal slow points across different folds used for computing the PSTH R² in a. To focus on the portion of the state space that is relevant to computation, we show only the region occupied by at least 50 of 5000 simulated single-trial latent trajectories-genenaiiee gennraate from the learned dynamics in Equation (2) for outlined by the dotted line. Trial- avee-ae trajectories from autoLFADS, projected onto the first two PCs. d, Trial-averaged trajectories from CEBRA, projected onto the first two PCs.
these domains.
of representational similarity between dynamical systems (Gosztolai et al., 2025) into the framework presented here.
While we expect FINDR to be generally applicable to a broad range of neural population data, its performance may vary across datasets. As deep learning- based model, FINDR performs best with datasets that have a large num- ber of simultaneously recorded neurons and many repeated trials. While the exact neuron and trial count that give good performance may vary depending on the dynamics in the dataset and the firing rates, generally an increase in the number of neurons should make FINDR's estimate of each single-trial dynamical trajectory more accurate, while an in- crease in the number of trials should make FINDR's 's estimate of the flow field more accurate, because FINDR has more latent trajectories that traverse the latent space to infer the flow field from. Future work could incorporate uncertainty estimates of the inferred flow field, as demonstrated by the recent Gaussian process- process-based switching linear dynamical system model (Hu et al., 2024).
The low-dimensional dynamical representation of neural population activity discovered by FINDR can serve as a bridge between the more fine-grained, neuronal-level rep- resentation and the higher-level, algorithmic representation of task computation. Future work could strengthen these cross-level connections. For example, recent studies have started to examine the relationship between the low-rank connectivity structure of a recurrent neural network (RNN) and its low-dimensional dynamics when trained on tasks and neural data (Valente al., 2022; Pals al., 2024; Langdon & Engel, 2025). In Luo et al. (2023), the authors show that low- low-dimensional dynamical representations inferred from data can facilitate the design of parsimonious, algorithmic- level models of neural computation, such as their multi- mode drift-diffusion model (MMDDM). Integrating these levels of description into unified framework is an inter- esting future direction. These directions highlight how new neural population dynamics inference methods like FINDR are a promising approach for bridging the gap between neuronal-level mechanistic descriptions and the higher-level algorithmic descriptions of neural function.
Additional extensions may include inferring external inputs (Schimel et al., 2022), developing more expressive decoders (Gao et al., 2016; Versteeg et al., 2024; Abbaspourazad et al., 2024), supporting online learning of nonlinear dy- namics (Zhao& & Park, 2020), or integrating the learning
9

---

Flow-field inference from neural data using deep recurrent networks
Acknowledgements
Duncker, L.., Bohner, G., Boussard, J., and Sahani, M. Learning interpretable continuous-tinn models of latent stochastic dynamical systems. Proceedings of the 36th International Conference on Machine Learning. 2019.
We thank Jesse Kaminsky, Belle Liu, Chethan Pandarinath, Andrew Sedler, Chris Versteeg, and Iman Wahle for discus- sions. We thank Jonathan Halverson for help with using the Princeton HPC clusters. This work was supported by grants from NIH R01MH108358, the Simons Collaboration on the Global Brain (SCGB AWD543027), and a U19 NIHNINDS BRAIN Initiative Award BU19NSI04648)
Elfwing. S., Uchibe, E., and Doya, K. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 2018.
Gallego, J. A., Perich, M. G., Naufel. S. N., Ethier, C., Solla, S. A., and Miller, L. E. Cortical population activity within a preserved neural manifold underlies multiple motor behaviors. Nature Communicutions, 2018.
Impact Statement
This paper presents work whose goal is advance the field of Machine Learning and Computational Neuroscience. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Gao, Y., Archer, E. W., Paninski, L.,, and Cunningham, J. P. Linear dynamical neural population models through nonlinear embeddings. Advances in Neural Information Processing Systems, 2016.
References
Abbaspourazad, H., Erturk, E., Pesaran, B., and Shanechi, M. M. Dynamical flexible inference of nonlinear latent factors and structures in neural population activity. Nature Biomedical Engineering. 2024.
Genkin, M. and Engel, A. Moving beyond generaliza- tion to accurate interpretation of flexible models. Nature Machine Intelligence, 2020.
Genkin, M., Hughes, O., and Engel, T. A. Learning non- stationary langevin dynamics from stochastic observa- tions of latent trajectories. Nature Communications, 2021.
Bishop, C. M. Pattern Recognition and Machine Learning.
Springer, 2007.
Burgess, C. P., Higgins, I., Pal, A., Matthey, L., Watters, N., Desjardins, G., and Lerchner, A. Understanding disentan- gling in 8-VAE arXiv, 2018.
Gholami, A., Keutzer, K., and Biros, G. ANODE: Uncondi- tionally Accurate Memory-Efficient Gradients for Neural ODEs. International Joint Conferences on Artificial In- telligence, 2019.
Chen, C., Yang, Z., and Wang. X. Neural embeddings rank: Aligning 3d latent dynamies with movements. Advances in Neural Information Processing Systems, 2025.
Gosztolai, A., Peach, R. L.. Arnaudon, A., Barahona, M., and Vandergheynst, P. MARBLE: interpretable represen- tations of neural population dynamics using geometric deep learning. Nature Methods, 2025.
Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural Ordinary Differential Equations. Advances in Neural Information Processing Systems, 2018.
Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. beta-VAE: Learning basic visual concepts with a constrained varia- tional framework. International Conference Learning Representations, 2017.
Cho, K., van Merrienboer, B., Gulcehre, C., Bahdanau, D., Bougares, F. Schwenk, H., and Bengio, Y. Learning phrase representations using encoder-decoder for sta- tistical machine translation. arXiv, 2014.
Chung, J., Kastner, K., Dinh, L., Goel, K., Courville, A., and Bengio, Y. A recurrent latent variable model for sequen- tial data. Advances in Neural Information Processing Systems, 2015.
Hu, A., Zoltowski, D. M., Nair, A., Anderson, D., Duncker, L., and Linderman, S. W. Modeling latent neural dy- namics with gaussian process switching linear dynamical systems. Advances in Neural Information Processing Systems, 2024.
Churchland, M. M., Cunningham, J. P., Kaufman, M. T., Foster, J. D., Nuyujukian, P., Ryu, S. I.. and Shenoy, K. V. Neural population dynamics during reaching. Nature, 487:51-56, 2012.
Hubel, D. H. and Wiesel, T. N. Receptive fields of sin- gle neurones in the cat's striate cortex. The Journal of Physiology, 148(3):574-591, 1959.
Duncker, L. and Sahani, M. Dynamics on the manifold: Identifying computational dynamical activity from neural population recordings. Current Opinion in Neurobiology, 70:163-170, 2021.
Jun, J. J., Steinmetz, N. A., Siegle, H., et al. Fully inte- grated silicon probes for high-density recording of neural activity. Nature, 551 251:222-236, 2017.
10

---

Flow-field inference from neural data using deep recurrent networks
Kato, S., Kaplan, H. S., Schrödel, T., Skora, S., Lindsay, T. H., Yemini, E., Lockery, S., and Zimmer, M. Global brain dynamics embed the motor command sequence of caenorhabditis elegans. Cell, 163(3):656-66)) 2015.
Loshchilov, I. and Hutter, F. SGDR: Stochastic Gradient
Descent with Warm Restarts. arXiv, 2017.
Luo, T. Z., Kim, T. D., Gupta, D., Bondy, A. G., Kopec, C. D., Elliot, V. A., DePasquale, B., and Brody, C. D. Transitions in dynamical regime and neural mode underlie perceptual decision-making bioRxiv, 2023.
Keshtkaran, M. R., Sedler, A. R., Chowdhury, R. H., Tandon, R., Basrai, D., Nguyen, S. L., Sohn, H., Jazayeri, M., Miller, L. E., and Pandarinath, C. A large-scale neural network training framework for generalized estimation of single-trial population dynamics. Nature Methods, 19: 1572-1577, 2022.
Macke, J. H., Buesing, L., Cunningham, J. P., Yu, B. M., Shenoy, K. V., and Sahani, M. Empirical models of spiking in neural populations. Advances in Neural Infor- mation Processing Systems, 2011.
Khona, M. and Fiete, I. R. Attractor and integrator networks in the brain. Nature Reviews Neuroscience, 2022.
Nassar, J., Linderman, S. W., Bugallo, M., and Park, M. Tree-structured recurrent switching linear dynamical sys- tems for multi-scale modeling. International Conference on Learning Representations, 2019.
Kim, T. D., Luo, T. Z., Pillow, J. W., and Brody, C. D. Inferring latent dynamics underlying neural population activity via neural differential equations. Proceedings of the 38th International Conference Machine Learning, 2021.
Nich, E. H., Schottdorf, M., Freeman, N. W. Low, R. J., Lewallen, S., Koay, S. A., Pinto, L.,, Gauthier, J. L.., Brody, C. D., and Tank, D. W. Geometry of abstract learned knowledge in the hippocampus. Nature, 2021.
Kim, T. D., Can, T., and Krishnamurthy, K. Trriiaability. Expressivity and Interpretability in Gated Neural ODEs. Proceedings of the 40th International Conference on Ma- chine Learning, 2023.
Onken, D. and Ruthotto, I.. Discretize-optimize vs. optimize-discretic for time-series regression and con- tinuous normalizing flows. arXiv, 2020.
Pals, M., Sağtekin, A. E., Pei, F. C., Gloeckler, M., and Macke, J. H. Inferring stochastic low-rank recurrent neural networks from neural data. Advances in Neural Information Processing Systems, 2024.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. International Conference on Learning Represen- tations. 2014.
Klambauer, G., Unterthiner, T., Mayr, A., and Hochreiter, S. Self-normalizing neural networks. Advances in Neural Information Processing Systems, 2017.
Pandarinath, C.,, Gilja, V., Blabe, C. H., Nuyujukian, P., Sarma, A., Sorice, B. L., Eskandar, E. N., Hochberg, L. R., Henderson, J. M., and Shenoy, K. V. Neural popula- tion dynamics in human motor cortex during movements in people with als. eLife, 4:e07436, 2015.
Krishnan, R. G., Shalit, U., and Sontag, D. Structured inference networks for nonlinear state space models. Pro- ceedings of the Thirty-Fire fiirty-Firtt.ie AAAI Conference on Artificial Intelligence, 2017.
Pandarinath, C., O'Shea, D. J., Collins, J., et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature Methods, 15:805-815, 2018.
Langdon, C. and Engel, T. A. Latent circuit inference from heterogeneous neural responses during cognitive tasks. Nature Neuroscience, 2025.
Park, I. M., Meister, M. L., Huk, A. C., and Pillow, J. W. Encoding and decoding in parietal cortex during senso- rimotor decision- making.. Nature neuroscience, 17(10): 1395-1403, 2014.
Li, X., Wong, T.-K. L., Chen, R. T. Q.. and Duvenaud, D. Scalable gradients for stochastic differential equations. Proceedings of the 23rd International Conference on Ar- tificial Intelligence and Statistics, 2020.
Pei, F., Ye, J., Zoltowski, D., Wu, A., Chowdhury, R. H., Sohn, H., O'Doherty, E., Shenoy, K. V., Kaufman, M. T., Churchland, M., Jazayeri, M., Miller, L. E., Pillow, J., Park, I. M., Dyer, E. L., and Pandarinath, C. Neu- ral Latents Benchmark "21: Evaluating latent variable models of neural population activity. Advances in Neural Information Processing Systems, 2021.
Linderman, S., Johnson, M., Miller, A., Adams, R., Blei, D., and Paninski, L. Bayesian Learning and Inference in Re- current Switching Linear Dynamical Systems. Proceed- ings of the 20th International Conference Artificial Intelligence and Statistics, 2017.
Pillow, J. W., Paninski, L., Uzzell, V. J., Simoncelli, E. P., and Chichilnisky, E. J. Prediction and decoding of retinal ganglion cell responses with a probabilistic spiking model. Journal of Neuroscience, 25(47): 000-110113 2005.
Liu, M., Nair, A., Coria, N., Linderman, S. W., and An- derson, D. J. Encoding of female mating dynamics by a hypothalamic line attractor. Nature, 2024.

---

Flow-field inference from neural data using deep recurrent networks
Pillow, J. W., Shlens, J., Paninski, L., Sher, A., Litke, A. M., Chichilnisky, E. J., and Simoncelli, E. P. Spatio-temporal correlations and visual signaling in complete neuronal population. Nature, 454:995-9999 2008.
Valente, A., Pillow, J. W., and Ostojic, S. Extracting com- putational mechanisms from neural data using low-rank mas. Advances in Neural Information Processing Sys- tems, 2022.
Ramachandran, P., Zoph, B., and Le, Q. V. Searching for activation functions. arXiv, 2017.
Versteeg, C., Sedler, A. R., McCart, D., and Pandarinath, C. Expressive dynamics models with nonlinear injective readouts enable reliable recovery of latent features from neural activity. Proceedings tfe 2nd NeurlPS Workshop on Symmetry and Geometry in Neural Representations, 2024.
Rigotti, M., Barak, O., Warden, M. R., Wang, X.-J,, Daw, N. D., K., M. E., and Fusi, S. The importance of mixed selectivity in complex cognitive tasks. Nature, 497:585- 590, 2013.
Versteeg, C.,, McCart, J. D., Ostrow, M., Zoltowski, D. M., Washington, C. B., Driscoll, L., Codol, O., Michaels, J. A., Linderman, S. W., Sussillo, D., and Pandarinath, C. Computation- whrough-dynamiice benchmark: Simulated datasets and quality metrics for dynamical models of neural activity. bioRxiv, 2025.
Safaie, M., Chang, J. C., Park, J., Miller, L. E., Dudman, J. T., G., P. M., and Gallego, J. A. Preserved neural dynamics across animals performing similar behaviour. Nature, 2023.
Sani, O. G., Abbaspourazad, H., Wong, Y. T., Bijan, P., and Shanechi, M. M. Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification. Nature Neuroxcience, 2021.
Vinograd, A., Nair, A., Kim, J. H., Linderman, S. W., and Anderson, D. Causal evidence a line attractor encod- ing an affective state. Nature, 2024.
Schimel, M., Kao, C., Jensen, K. T., and Hennequin, G. ilqr-vae: control-based learning of input-driven dynamics with applications to neural data. International Conference on Learning Representations, 2022.
Vorontsov, E., Trabelsi, C., Kadoury, S., and Pal, C. On orthogonality and learning recurrent networks with long term dependencies. Proceedings of the 34th International Conference on Machine Learning, 2017.
Schneider, S., Lee, H., and Mathis, M. W. Learnable latent embeddings for joint behavioural and neural analysis. Nature, 2023.
Vyas, S., Golub, M. D., Sussillo, D., and Shenoy, K. V. Computation through neural population dynamics. An- nual Review of Neuroscience, 43:249-275. 2020.
Sedler, A. R. and Pandarinath, C. Ifads-torch: A modular and extensible implementation of latent factor analysis via dynamical systems. arXiv, 2023.
Wang, Y., Blei, D., and Cunningham, J. Posterior collapse and latent variable non-identifiabiiity. Advances in Neural Information Processing Systems, 2021.
Sedler, R.. Versteeg, C., and Pandarinath, Expressive architectures enhance interpretability of dynamics-based neural population models. Neurons, Behavior, Data anal- ysis and Theory, 2023.
Yu, B. M., Cunningham, J. P., Santhanam, G., Ryu, S.. Shenoy, K. V., and Sahani, M. Gaussian-process factor analysis for low-dimensione single-trial analysis of neu- ral population activity. Advances in Neural Information Processing Systems, 2008
Siegle, J. H., Jia, X., Durand, S., et al. Survey of spiking in the mouse visual system reveals functional hierarchy. Nature, 592:86-90 92, 2021.
Zhao, Y. and Park, I. M. Variational online learning of neural dynamics. Frontiers in Computational Neuroscience, 14: 71, 2020.
Steinmetz, N. A., Aydin, C., Lebedeva, A., et al. Neuropix- cls 2.0: A miniaturized high-density probe for stable, term- brain recordings. Science, 372(6539))cabf4588, 2021.
Zoltowski, D., Pillow, J., and Linderman,
general recurrent state space framework for modeling neural dy- namics during decision decision-making. Proceedings of the 37th International Conference on Machine Learning, 2020.
Stringer, C., Pachitariu, M., Steinmetz, N., Reddy, C. B., Carandini, M., and Harris, K. D. Spontaneous behaviors drive multidimensional, brainwide activity. Science, 364 (6437):255-25) 258, 2019.
Sussillo, D. and Barak, O. Opening the Black Box: Low- Dimensional Dynamics in High-Dimensional Recurrent Neural Networks. Neural Computation, 25(3):626-644 649, 2013.
12

---

Flow-field inference from neural data using deep recurrent networks
A. Appendix
A.1. Model Specification
Dynamics in a neural population may be associated with not only the task that the animal performs but also other factors that are not relevant to the task. Therefore, in FINDR, we distinguish between the task-irrelevant dynamics d, and task-relevan low-dimensional dynamics 2, and a separate inference procedure for the two. These latent variables and task stimulus u, generate spike trains Y, from N recorded neurons. We assume that the total length of trial is T, and we bin T into T' bins, with each of the bins having the same width At. We assume that there are total of M trials in an experimental session. We denote t {1....,T) } representing the t-th time bin within a trial, and m E {1, M representing the m-th trial. In all tasks that we consider this work, T = second, with A/ = 0.01 second. In Section 3.3, we only analyzed the epoch trial from stimulus onset to right before movement initiation, which was variable across trials and less than or equal to second. In Section 3.2, for each trial, we looked at the first second from stimulus onset. Thus, if the stimulus duration was less than second on a given trial, we also considered spiking activity during movement. We fixed the lengths of all trials to be the same in this Section to allow comparisons existing models (datasets with variable trial lengths not yet supported in some model implementations consider).
A.1.1. INFERENCE OF TASK-IRRELEVANT DYNAMICS
To model task-irrelevant fluctuations in an individual neuron's 's firing rate across a total of M trials, we first compute each neuron n's average firing rate for each trial m with
# of spikes in neuron n. in trial
m
(8)
duration of trial m (in seconds)
Here we let the 12-th element of the N -idimensional vector Ym to be ym We partition M randomized trials into 5 equal subsets, where 3/5 of the M trials are used for training (Durai). 1/5 for validation (Dvat), and 1/5 for testing (Dwa). For each neuron n, we a linear basis function model (Bishop, 2007) with the following loss function:
De
min Pppens
( (peross,(1) 122 + across
I across, (rrooss(n)]pp) (n) 112
(9)
Dasin
where
D
( (n)
w
(nn) (n)) (m),
Pij bp
j=1
(10)
Pj
(n)
(m)
= cos H
(1) -1) π n 28max
(S(m) - dooss
acTONS [i]) + +1).
2
disp
[]] = linspacee (Smin, Smax, Decross [i].
Here,
n. if <
H(x)= = I,,
if - S
(11)
-T i
and ((m) is the time stamp of the onset of trial m (which lies between the beginning of the session, Smin = Os, and Smax 10, 000s, and rounded to the nearest second; we assume no session runs more than Smax). Here, linspace is a function that returns aecoss evenly spaced numbers over the interval (Smin Smax). The hyperparameters Dacross E 5,.,,.,, 10) and peross E (0.1, 0.1 0.1 0.1 0.1 are optimized by evaluating the loss D
To model iase-itrelevant fluctuations in individual neuron's firing rate within each trial, we again fit, for each neuron n. linear basis function model (Bishop, 2007) with the following loss function:
Deiitii
min vatin. (x)
d (n)
d. (ithin,(m))
112 + within
(m))
(12)
Lippiia m Dunis
13

---

Flow-field inference from neural data using deep recurrent networks
b
1.0
function values
function
0.2
0.0
5000 7500 time in session (5)
10000
1.00
time from stimulus onset (s)
Figure 5. Examples raised cosine basis functions (Pillow et al., 2005; 2008; Park et al., 2014) used in the inference of task irrelevant dynamics. a, opymmm (m) (m) when D De
b, PPj within,(m) (t) when Diithin
where
d
(wthin,(n))
within, (n) (m) (I);
jjj
oj
==1
(n) (1)
nT DDDiiiiin - 1) Carsinh(111885 (arsinh(0.12(( 1)) disp within (+]) b1))
+1). for l E (l, T),
(13)
cos H
disp wilhin [i] linspace(O, arsinh((1.885), Divithnn [i]]
Here, acmocsss.. the solution to Equation (9). The hyperparameters Dwithin {5, 6.., 10} and p p (0.1,0.12, 0.13, 1.10..1,.... 0.10) are optimized by evaluating the loss on D After obtaining the solution of Equation (12), d; dr (n) we set d d (rrossss + d. (x) Once we do this procedure for all neurons n, we get dmn,t, which is used in the task-relevant dynamics inference in Appendix A.1.2.
A. 1.2. INFERENCE OF RELEERE DYNAMICS
For simplicity, going forward we suppress m in our notation whenever we can. To maximize the log-likelihood of observing the population spike trains y given the task-related external inputs u, the task fask-irrelevant baseline inputs d and the model parameters need compute
log po (V JurT, d1.)) = log (/I
T
po (oryleee do)o u, )xz
(14)
[=]
Here, Po (z,, d,) is given by Equation (1). The term po(zz|z--1. u.) is given by Equation (2). We call Equation (2) the prior process. Here, we assume that the initial condition of the latent zg = O. Equation (14) does not have a closed-form solution, and computing this quantity can be computationally expensive. Therefore, we instead compute the evidence lower bound (ELBO) (Kingma & Welling, 2014):
log po U1.T, dr.r) 2 Eyg log po (V duar) DKL (qo(21rly),)) (Z
(15)
by introducing variational posterior qs. Here, DKL is the Kullback-Leible divergence between the prior pe and the variational posterior qø. We can further decompose the first term into:
T
Egg [log (1 dr:r)) = Egg
Σ
log pa pol.ue
(16)
14

---

Flow-field inference from neural data using deep recurrent networks
-1.0
-5.
0.8
0.5
stimulus
16-
12
17.5
12.5
splasss
10.0
5.0
2.5
100
200
300
400
100
200
300
400
1000
200
300
400
0.0 0.2
82- 0.6
0.4
trial
trial
trial
Figure 6. Estimated th closely matches the observed PSTH, and estimated d matches the true across-trial firing rate fluctuations.
a, 5-fold valsdavald PSTH reconstruction from FINDR (in bold line) plotted against observed PSTH (in shading; 95% confidence interval) for three example neurons. Histogram of the cross- ross-validated ft² between the observed and FINDR PSTHs (median=0.82) 5-fold vrass-validated firing rate prediction from FINDR plotted against the observed firing rate across trials for three example neurons. Histogram of the cross-validated R² between the observed and FINDR firing rates trials (median=0.05). For the majority of neurons, having a time-varying baseline across trials instead of a constant baseline improves goodness-of-fit.
The KL term also be decomposed into:
DKL (2o( Mir,, u1 (1or) (1:r)) = dz dz1:1g u1:r) log 40(21: Wi.rs urr)
po(2u:((u:r)
(17)
Bgg (DEL (20(z)l21-1 Marry ( u.)]]
1=1
Therefore, our objective becomes:
log logpo(W11. W
U1TT di.r)=
2
Eyg (ogpoly,l200 d.) DKI. (4o (ze 22-11 M1.r, u1:7) (p0 u.))
(18)
We specify do with Equation (7), and call Equation (7) the posterior process. The initial condition 2n again set to be here. We also let the diagonal matrix be the same in both the prior and posterior processes make the KL divergence finite (Li al., 2020). e, 2H encodes representations of y u.r and uır using bidirectional GRU (Cho et al., 2014):
GRU (f--1: Vey Uz) b, GRU,, (be+1, U. uz),
e, [f,; by],
(19)
where fg HaN and b, E R Hasu To train FINDR, we compute the gradient of
[(0 Σ log po (o 22, d.) 3KL. (ga(2, (e))l:==__ u.)))
(20)
with respect to 0 and & using backpropagation through time (BPTT), where we sampled žn once from (z ex) for t (1,..., T}. Note that we included 9 the KI. term in our loss [(0, d).
15

---

Flow-field inference from neural data using deep recurrent networks
With the included back into notation, Equation (20) be re-written
L)),0)
ΣΣ
Po m, Am,t dmm.))
(21)
Σ DRL (mm I:T. Mm.. lppo (2m, um,,))
where
ΣΣ 3DL
1:T. um.r:))lpp (=m.elmm.--1, um,)))
(22)
equal to:
E (v- 1 - (v-)), u),
(23)
m
where the arguments of the functions y and μ are (2m,-11, Em, Um,,) and (2m,t-1, Um,e). respectively.
1.1.3. MODEL ARCHITECTURE
For G and F in Equation (3), we use
G(z, u) WG,00 Va,ou + bb,0
F(z, u) WF,SSLL(WFOO + VVLu + bF,0) + bF
(24)
Σ = (diag(s)),
Similarly, for G and F in Equation (7), we use
G(z, u. e) = W + VV + UG 68,0 F(z,u, =
(25)
+ V7,0 + U + bp,o) + bp,1
where Wo.0, 8 RLXL Va,0, V2.0 E R.xxdi((u), Uc Uč.0 RL22HEN WFo, WF E RINNLL WF, W F. R Lx Hoss VyD, Vp.o E (Hinwxdim(u)), F. E H bG,o, bā E &RLbbro bi.u uu E IR Hoss, bri, E R,, and R- are trainable parameters. Here, 2HRNN hyperparameter that determines the size of the bidirectional GRU, and HENN is hyperparameter that determines the width of the hidden layer. In Equation (19), and by are trainable parameters. We let
(2 yy Em,t, d Poisson(AtA... At softplus dm,)));
(26)
where C E Nx is trainable. For all FNNs in this paper, we use the SiLU (a.k.a. swish) activation function (Ramachandran et al., 2017; Elfwing et al., 2018).
1.1.4. INITIALIZATION
We the initialization scheme in Kim et al. (2023) for the kernels that transform the hidden states in our gated MLPs. We use the orthogonal initializer for the kernels that transform the hidden states in the GRUs. We use the Lecun normal initializer (Klambauer et al., 2017) for the kernels that transform the input in both the GRUs and gated MLPs. Biases are i..dd normal with variance 10 10-6.
A.1. OPTIMIZATION
To train this model, we use the discrete adjoint sensitivity (i.e., standard backpropagation through time) to compute the gradient of Equation (21) with respect to {0, Φ). A few studies (Gholami al., 2019; Onken & Ruthotto, 2020) show that the discrete adjoint sensitivity produces more accurate gradients than the continuous adjoint sensitivity used in Chen et al. (2018). We train for total of 3000 epochs and minimize loss using mini-batch gradient descent with warm restart (Loshchilov & Hutter, 2017). The learning rate increases from 0 to linearly for 10 epochs every Deyele 2 D epochs, where goes from to lum-- After the 10 epochs, the learning rate decays in cosine manner, where at Doyete, the learning rate becomes 0. inn is determined by the minimum Σ Diml Daycle, which is greater than or equal to 3000. D is set be 200.
16

---

Flow-field inference from neural data using deep recurrent networks
A. 1.6. HYPERPARAMETER GRID-SEARCH
For each of the 5-folds in single experimental session dataset, we do grid search on the parameters (1, HeNN, HRNN)) to identify the model that performs best when the objective is evaluated using the validation dataset. Here, nj {10- 2.0 10 10-1..25 10010-..000 10-00.75, 100000000 0.5. HFNN {30, 50, 100}. and HRNN {50, 100, 200). HENN is the num- ber of hidden units in FNNs Fg and Fg, where both networks had a single hidden layer. HRNN represents the number of units for both GRU,, and &R Thus the total number of units for the bidirectional GRU is 2/RNN.
A.1.7. FIXED HYPERPARAMETERS
We train FINDR for a total of 3000 epochs. For the first 100 epochs, we train only the first 30 time bins of the trials. Then for the next 200 epochs, we train only the first 50 time bins of the trials. Afterward. we fit all time bins in the trials. We set the coefficient of the lg regularization single hidden layer. We set the time constant 0.1s. We set 3 = 2. We set the number of trials in a mini-batch to be 25. We set the momentum in mini-batch gradient descent to be 0.9. We perform annealing to the KI. term in Equation (22). Specifically, the KL term is multiplied by ..iieeiaiinn
the weights of all model parameters to be 10-? We let Fg and Fg be FNN with
A.2. Identifiability of FINDR
We use the definition of latent identifiability in Wang et al. (2021) in the context of FINDR. FINDR is non-identifiable when pe lyy 2, d) pe( ( 2 d) for all possible z and ž. For the latent variable to be identifiable, the sufficient condition is for Po z, d) to be an injective function of z for some at the end of model training (Wang et al., 2021). In FINDR, when rank(()) = L, Equation (1) becomes injective. Therefore, we check at the end of model training whether the rank of the learned C' matches L, and consider only models that satisfy rank(C) L the end of training.
A.3. Post-Modeling Analysis
While we do not constrain C' in Equation (26) to be semi-orthogonal (i.e., C I) during the training procedure, a semi-orthogonal C may be desired because distance and angle in the latent space R are distance and angle in the inverse-softplus rate space R.V. More precisely, [x2||2 zT C Cz = =T ||2|| for all z. Having a semi-orthogonal, map C' would make the latent trajectories inferred by FINDR more interpretable. We do
and therefore distance-preserrvig not put either the soft or hard constraints on C' to be semi-orthogonal because having such constraints is known to worsen performance in other contexts (Vorontsov et al., 2017). We also find this to be true when we add soft constraints to enforce orthogonality.
Before we interpret the latent trajectories z and the inferred flow field (=(2, u). we perform singular value decomposition (SVD) on C USV where U e IR N is a semi-orthogonal matrix, IR Lx a diagonal matrix with its entries populated by the singular values and V R is orthogonal matrix. Then, we apply a transformation ğ = SV 2. We next perform principal component analysis (PCA) ă so that the first component of the transformed ž = U pa bpea corresponds to the first PC, and the L-th component of the transformed ž corresponds to the I-thh PC. This transformation by PCA is rigid. Therefore, the distance and angle in the space and axes given by & are still the distance and angle in the inverse sortpll rate space. The transformed flow field Po in the space of z is given by:
A = U SV
ž Az + bpen
(27)
Bo u) = (A (A
bpca), u).
In all of our analyses that show the latent trajectories and flow fields inferred by FINDR, we plot the transformed æ and Ao (ž, u). To project the flow field (ž, u) onto the first two PCs, we assume that the third and later components of ž are zero, and consider the first two components of श pe (2, u).
A.4. Model Evaluation Metries
In Section 3.2, the normalized log- difference score (NLL) is defined as:
NLL
E, [Poisson log- loglliielihood( Ae, y:) Poisson log-likelihood((1ngg (A:r; 9;)]
(28)
17

---

Flow-field inference from neural data using deep recurrent networks
similar the definition in Pei et al. (2021). Here, X1-r is the mean firing rate estimated from 91.r- We use another model evaluation metric called evidence- coneiiionon peristimulus time histogram (PSTH) R², which is defined for each neuron as:
SS
R² 1
SSres
=
Σ (PSTH PSTH + (PSTH PSTH L
pred,t)
)
(29)
SShot Σ
(PSTH obs, obs,t E, [PSTH )) (PSTH - E, (S.T :)²
For given neuron, we binned the spike train with a 10ms bin-width, and then convolved it with a causal Gaussian linear filter with standard deviation of 0.1s and a width of 0.3s. Then, PSTHs were computed by averaging this smoothed spike train across trials. Here, PSTH," represents the average firing rate computed from trials where there were more left clicks than right clicks, and PSTH represents the average firing rate computed from trials where there were more right clicks than left clicks. We computed the PSTHs from the observed and FINDR- predicted firing rates (* {obs, pred)) to compute the R².
A.5. Fitting FINDR to Synthetic Datasets with Autonomous Dynamics
We found that FINDR correctly captures autonomous limit cycle dynamics in synthetic dataset (Figure 7). To produce spike trains from the limit cycle dynamics (Figure 7a), we we used similar settings as the datasets in Section 3.1 and Extended Data Figure 1 (but with 500 trials, 80 neurons, and a constant bias of spikes/s).
a
b
True dynamics
Inferred dynamics
222
-1
PC1
Figure 7. FINDR correctly captures autonomous limit cycle dynamics from simulated spike train data. a, Ground truth dynamics used to generate the spike trains. b, Dynamics inferred from FINDR. The inferred dynamics captures limit cycle.
B. Software and Data
Our code is available as GitHub repository: http:://gghhha..oo//////////////
For SLDS and iSLDS, used code from httpp://ggtthub.ccm////ddeeemm abb.m/aaaaaaaammmaammmmmmmm For autoL
we used code from https: /tttss:/gitte https://github. http:////eeeeeee///aaaeeeeeeeee/ urations in configs/pbt.yaml. For GPFA, we used Elephant: httt::///githuub..mm//eeeta..// el elephant. For CEBRA, we used code from http://sss...a...cooaaaa/aaaaaar
with hyperparameter search config-
/ We fit a Euclidean-distance CEBRA-Time model using hyperparameters from http://aaannteototttttt//e/eaaaa natebooks/CCBRa_best___ps htw.:/eeees.mttttmmmmmmmmmmmm html come:mmesttttooooooeeeee..itiiii changes three hhw.ehnneeeetteeeeeeeeeeeeeeeee
max www.asee....teaatiiiioo.o..
C. Author Contributions
T.D.K. conceptualized the method. T.D.K. and T.Z.L. developed the inference method for task fasevinreeevant dynamics. T.D.K. developed the inference method for task-relevant dynamics. T.Z.L. collected data. T.D.K., T.C., and K.K. developed the gated multilayer perceptron (MLP) used this method. T.D.K. implemented the method as a software package. T.D.K. wrote the manuscript after discussions among all authors. J.P.P. and C.D.B. supervised the project.
18

---

Flow-field inference from neural data using deep recurrent networks
Extended Data Figure 1
inputs Masa
C
L =
L 3
L == 4
L = 5
L 6
latent dynamics
1.00
cherme
0.95
0.90
0.85
charne
2
2
4
3 5
# principal components
0.0
sme 16)
normalzzed log hg-ihellood
RC2
Oon bed (b66s
0.80
0.75
0 PC
0.65
(6)
L == 2
==3
L = 4
== 5
== 6
d
explainee
1.0
inputs More latent dynamics
0.9
3-
dainnl
0.8
-3-
0.7
2
2
4
3
5
2
5
2-
principal components
chennel2
g
noreallzed log Stllrood
-2
2.8
[on lost (att)
channel S
1.6
PC2
PC3
0
2.4
0.0
0.5 Smo (6)
1.0
2 3
5
-1
0
-1
0
1
-1
1
of oflinlntt dimeriion (1)
PC
PC
PC 3
Extended analysis related Figure 2. a. We generate transient pulse inputs from independent channels and let 2- dimensional system maintain the value of the most recent pulse in each channel. The pulse value in channel 1 (ci) satisfies -2 S and the pulse value in channel 2 (c2) satisfies -1 < c2 S 1. b-c, Analysis similar to Figure 2c-d. d. We generate transient pulse inputs from 3 independent channels and let 3-dimensional system maintain (with some noise) the value of the most recent pulse in each channel. The pulse value in channel (c1) satisfies -3 < < 3, the pulse value in channel 2 (c2) satisfies -2 < 02 I 2, and the pulse value in channel 3 (cs) satisfies -1 S c2 s 1. e-f, Analysis similar to Figure 2c-d. g, The range of values that the latent trajectory z takes along PC is larger than the range of values that z takes along PC and PC The range of values that takes along PC larger than the range of values that z takes along PC 3. This suggests that distance in this latent space is preserved and reflects the statistics of the pulses.
19

---

Flow-field inference from neural data using deep recurrent networks
Extended Data Figure 2
neuron
neuron
neuron
13
14
14
held-out neurons
FINDR
spkkess
4-
R²= ===00 0.96
R = 0.80
R² = 0.52
1.0
evidence-condflonne PSTH R²
time from atimulus onset (8)
13
14-
14
held-out neurons
18
SLDS
R² 0.47
R 0.53
R² = 0.21
0.0
0.5
1.0
evidence-conddlloned
PSTH R²
13
14
held-out neurons
rSLDS
0039.
0.19
0.0
0.5
1.0
0
eviddecce-oondilioned PSTH
held-out neurons
autoLFADS
4
Ra 0.80
R² 0.79
R² - 0.29
0.0
0.5
1.0
evidence-come condllioned
wwidence-condtliones
PSTH R²
13
14
held-out nourons
GPFA
V
R² = 0.71
R 0.70
R² - 0.17
0
1.0
evidencee-oo condiiioned PSTH R
observed 98% Ci prodiction
Extended analysis related to Figure 3c. Same Figure 3c. b, Analysis similar to Figure 3c, but for SLDS (L = 6). Analysis similar to Figure 3c, but for rSLDS (L 6). d, Analysis similar to Figure 3c, but for autoLFADS (L 6). Analysis similar to Figure 3c, but for GPFA (L- = 6).
20

---

Flow-field inference from neural data using deep recurrent networks
Extended Data Figure 3
a
C
R conditand
0.05
p-0
p=0.524
0.60
0.55
0.50
FEE FEET
-0.50 -0.25 0.00 0.25 0.50
of FINDR lateres
evidence conditicned FSTH
PST1AR AR
(FINDR (-2 FNNR (-1)
(FTER (ND FINER (+2)
d
p=0.549
p=0.001
p=0.327
neeions
-0.50 -0.25 0.00
0.25 0.50 PSTT AP
-0.50 -0.25
0.00 0.25
0:50 PIST AJY
-0.50 -0.25
0,00
0.25 0.50
PS
(FINER 2md NER 2x43)
(F 2nd FIRM 2m4))
(FINERE 2,4 TINERE (.45)
Extended analysis related to Figure 4. FINDR with L = is sufficient to describe the data. To show this, we split the dataset into 5 different folds, where each fold contains a subset of trials in random order. We train FINDR on 3 of the folds, validate 1 fold, and test its performance on the remaining fold. a, We compute the 5-fold cross-validated evidence-sign conditioned PSTH R² for all neurons in the dataset and take the mean. We see "elbow" at L = 2. b, We compute the evidence-sign conditioned PSTH It. for FINDR assuming L = 2 and L = 1. Then we take the difference between the πt obtained from FINDR with L 2 and R obtained from FINDR with L = 1 for each neuron. We find that FINDR with = 2 performs significantly better than FINDR with = (Wilcoxon signed-rank test, < 0.001). c, Analysis similar b, and = 3. e, Analysis similar but for FINDR with L 5 and 4. f. Analysis similar b, but for FINDR with and L 5.
but for FINDR with L 3 and 2. d, Analysis similar b, but for FINDR with
21

---

Flow-field inference from neural data using deep recurrent networks
Extended Data Figure 4
fold
fold
fold
fold 4
fold
t=1s
|=38
1=5s 1
Extended analysis related to Figure 4. While the maximum duration of the auditory stimulus is second in the data, and fit FINDR only the stimulus period, seconds, the trial-averaged trajectories reached a steady state either at point in the upper region of the state space or point in the bottom region the state space. What point the trajectories reached depended on the sign of the evidence (either left or right evidence), and suggested that these slow points might maintain the memory of the choice. Here, inside the dotted line represents part of the state space visited by single-trial trajectories, similar to Figure 4.
the model for more than 1 second. When the model was run for 3 - 5
Extended Data Figure 5
b
fold
fold 2
fold 3
fold 4
fold
B
Extended analysis related Figure When the time-varying task-irrelevant within- and across-trial dynamics are not learned, and instead d, was be a constant bias d for all and all trials, a. we find that evidence-sign conditioned PSTH R 's are significantly lower for this model compared to the FINDR model that learns the task-imelevant dynamics. b. We also find that the latent dynamics learned by FINDR are less consistent across folds. Here, inside the dotted line represents part of the state space visited by single-trial trajectories, similar to Figure 4.
22

---

Flow-field inference from neural data using deep recurrent networks
Extended Data Figure 6
a
fold 1
fold 2
fold 3
fold 4
fold 5
I=3
2
R2
1.44
1=5
100
-34
10
M
1:=6
16
Extended analysis related Figure 4. The latent trajectories discovered by FINDR with L> 2 are consistent across folds, and span mainly the first two PCs. As in Figure 4b, plot the trial-averaged trajectories sorted by their evidence strength, with red indicating leftward evidence and blue indicating rightward evidence. The trial-averaged trajectories are projected onto the first two PCs. b, The first two PCs of the latent trajectories explain 95% of the variance. c, The latent trajectories in a are mostly consistent across folds and across dimensions, with an "S"-shape to the trajectories. For those conditions where there was no prominent S-shape (e.g., fold 3 of 3, fold L = 4, and fold 3rd component was relatively higher than the other folds (as indicated by black dotted circles in c), and we could find the "S" "S"-shaped trajectories in the 3-dimensional PC space.
= 6), the variance explained by the
23

---

Flow-field inference from neural data using deep recurrent networks
Extended Data Figure 7
b
best model, seed=1
2nd best model, seed=1
best model, seed=2
4.2
2.72
1.6
261
6 6 12
Extended analysis related to Figure 4. a, We find that the second-best hyperparameter choice gives a representation consistent with the best choice. We also find a consistent representation when we train the model with a different initialization. b. The "confidence heatmap" showing the single-trial trajectories generated from the posterior. We expect that the dynamics inferred around regions traversed by more trajectories will be more accurate compared regions traversed less. c, To directly assess how well FINDR extrapolates from training data, trained FINDR to the same dataset used in Figure 4, but held out the last 0. of each trial (= 10 time steps). We then computed 5-fold cross-validated Rt for each time step between z from the full vs. the held-out models.
Extended Data Figure 8
b
a
fold
fold 2
3
state at 4s
state
1
2
PC
-1
-2
-2
0
-2
PC
PC 1
Extended analysis related to Figure 4. a, For fold initial conditions inferred from the encoder (448 initial conditions, because there were total of 448 trials in the dataset in Section 3.3). We found that the autoLFADS states at 4s and states at 5s are close to each other, suggesting they may have reached approximate steady-states (black line indicates displacement from state at 4s to state at 5s). However, we found that the states did not form two clusters as would be expected from bistable attractors. The states are visualized with the first two PCs. b. We performed the same analysis for fold 2.
the trained generator forward in time for 58, starting from the
24


Reviews:
[{"review_id": "gdkpyFqWw5", "reviewer": "Reviewer_geCp", "summary": "The paper introduces FINDR, a method for inference of low-dimensional stochastic dynamics from neural recordings. \nThe proposed approach estimates first the bias $d$ of the observation function (called task irrelevant component of the spiking activity) by  solving a regression problem for the average firing rates of each neuron. This aims to capture fluctuations in the spiking activity not attributed to the latent dynamics.\nThe authors first estimate an across-trial bias, by fitting a linear model with raised cosine basis functions to the trial-averaged firing rates, capturing slow fluctuations in each neuron’s baseline activity. They then model fast within-trial fluctuations using another linear model with raised cosine basis functions and sum both components to obtain the overall time-varying bias $d$ for each neuron.\nThe main inference method uses a sequential variational autoencoder with a semi-orthogonal loading matrix to infer a low-dimensional latent representation of neural population activity, and to learn underlying stochastic dynamics (flow field inference). The authors model the drift function with gated multilayer perceptrons The optimization uses backpropagation through time to compute the gradients.\n\n \nThe authors first demonstrate the method on synthetic data, showing that it estimates correctly the underlying dynamics, performing better than other models, when the latent dimension is correct. Then, on selected recorded data from rat PFC performing a decision-making task, the authors show that FINDR outperforms SLDS, rSLDS, autoLFADS, and GPFA in predicting held-out neuron activity. The authors also show that the inferred flow fields are consistent across cross-validation folds and identify attractors related to the decision-making task.\n\nOverall, the method demonstrates good performance. However, its technical implementation is somewhat involved. Nevertheless, the complexity of the approach is well suited for an application-focused paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s thoughtful comments.\n\n**Necessity of gating:**\nKim et al., ICML, 2023 ([B]) shows that gating increases expressivity and trainability of the dynamics function. Their Fig. 1 suggests gating is necessary for correctly inferring dynamics in our synthetic dataset.\n\n**Flow fields and interoperability:** \nWe respectfully disagree with the reviewer that flow fields can be visualized for all other baselines. While rSLDS allows flow field visualization, SLDS, autoLFADS, and GPFA do not. rSLDS failed to capture neural data (Fig. 3d), likely due to learning a constant d, and not a time-varying d, so we did not plot flow fields estimated using rSLDS.\n\nPrevious work has proposed that the model is interpretable if it discovers representations of computation that we can identify in the low-D dynamics ([B]; Duncker et al., ICML, 2019). Accordingly, one way to quantify interpretability in the flow field is in terms of the log-likelihood and R^2 metrics of how well neural data is captured in low-D (L<=3). Using these metrics, FINDR was the only model that captured neural data well in low-D.\n\n**Task-irrelevant components and consistency in flow field:**\nBy design, task-irrelevant components do *not* use the task-relevant inputs u to capture within- and across-trial baseline firing activity. In Ext. Data Fig. 5, colored trajectories represent trial-averages sorted by evidence strength (as in Fig. 4b), and inside the dotted line represents part of the state space *visited* by single-trial trajectories.\n\nLearning time-varying d improves performance (Ext. Data Fig. 5, last panel), and is necessary---when d is constant, the median evidence-sign conditioned PSTH R^2 is -0.009. We will include this result as the first panel, move the last panel next to this panel, and clarify the definition of the dotted line.\n\n**Limited accessibility:**\nWe will release our code with documentation and a tutorial notebook.\n\n**Model robustness:**\nWe find that the second best hyperparameter choice gives a representation consistent with the best choice. We also find a consistent representation when we train the model with a different initialization.\n\n**Identifiability:**\nPlease see response to Reviewer UnSA.\n\n**Additional comparisons with autonomous dynamics:**\nWhile we did not directly compare to Duncker et al. and Genkin & Engel, we will add analysis showing that FINDR correctly captures autonomous limit cycle dynamics in a synthetic dataset.\n\n**Questions**\n\n(1) \\tau=0.1s was a fixed hyperparameter (A.1.7), but the timescale of the dynamics is adaptive. That is, it depends on z and u through the gating function G. [B] analyzed expressivity as a function of \\tau (their Fig. 3), and showed that as \\tau increases, the drift function \\mu needs more parameters to fit data well. Given the number of our model parameters and trial durations (<1s), \\tau=0.1s was a reasonable choice.\n\n(2) Yes, external inputs u are given to the model. We will clarify in our revision that in Fig. 3 and 4, u_t was 2-D: [0;0] (no click), [0;1] (right click), [1;0] (left click).\n\n(3-1) If there is a constant shift in dynamics that is independent of u, this will be included as bias d and not as part of z. For the task-irrelevant component, because we use a linear basis function model, the solution we obtain is analytical and unique, given the regularization coefficients, which we optimize using the validation dataset (A.1.1).\n\n(3-2) We validated d_estimated against d_true using real datasets. Across-trial d_estimated closely matches true across-trial firing rate fluctuations, and within-trial d_estimated matches the observed PSTH (R^2=0.82). We will include relevant figures in the Appendix.\n\n(3-3) Our ablation study (Ext. Data Fig. 5, last panel) shows that without estimating time-varying d, the model fails to capture data (median R^2: -0.009).\n\n(4) Similar in spirit to the approach in Genkin & Engel, we split the data into five different folds to find features that are consistent across folds. We consistently find two approximate point attractors with the selection of \\beta=2 (Ext. Data Fig. 4). We will mention this, and cite references suggested by the reviewer.\n\n(5) We learned the parameters of the diagonal elements of \\Sigma (L99, 2nd column, Eq. (25)) directly via SGD. For generating the ground truth latents, we first generated the latents and added noise to each timestep with N(0, \\sigma^2=0.01). The inferred \\Sigma, after affine transformation to match ground truth latents, was found to be [[0.026, 0.001]; [0.001, 0.018]].\n\n(6) For Fig. 3, the inference of the task-irrelevant component doesn’t require GPUs, and takes <30 minutes. Task-relevant component (jax/flax-based) takes <1.5 hours per hyperparameter configuration on a single A100 GPU. It typically took total ~5 hours for FINDR to do grid search and complete training on our cluster. AutoLFADS (PyTorch) took ~6 hours. For SLDS, rSLDS and GPFA, no GPU was used, and typical runtime was <1 hour.", "claims_and_evidence": "- The main claims are that FINDR infers accurate dynamics, outperforms existing methods, and provides interpretable visualizations. The synthetic data experiments support the first claim by showing correct latent dimension recovery and attractor structure. \n- The gated parametrisation of the drift is more expressive and trainable compared to non-gated models. I am not sure that there is demonstrated evidence for this in the paper.\n- Consistency across folds and attractor visualization in flow fields support the third claim. However, such interpretable visualisation of flow-filed can be made for all other baseline methods used for comparison.\nThe interpretability of flow fields is demonstrated but might benefit from more quantitative measures. The separation of task-relevant/irrelevant components is shown to improve consistency, but the paper doesn't explore what exactly the task-irrelevant components capture. \n\n- The paper asserts that it can effectively disentangle task-relevant from task-irrelevant neural activity by estimating two components of the bias d, designed to capture within-trial and across-trial fluctuations that are not explained by the latent dynamics. However, the current version of the paper lacks clear evidence or numerical validation demonstrating the accuracy of this estimation method and the parameters of this part of the method are not rigorously analysed. In the Ext figure 5 the authors claim that by setting a single constant bias $d$ across trials the estimated flow fields are less consistent across splits, but to my eyes the projected trajectories in the five splits of that figure seem rather consistent, and the inconsistency seems to occur in the parts of the (projected) state space not visited by the system (see also my questions to the authors below for this part.)", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-motivated and suited for inferring neural population dynamics. \nThe authors first apply the method on synthetic data with known underlying dynamics, and demonstrate the performance in terms of normalised log-likelihood on test data.\nThe authors could probably use a metric that quantifies differences between ground truth and estimated flow field for this setting with known dynamics to directly show the accuracy of the flow field inference part.\n\nThe use of synthetic datasets, with known underlying attractor structures, provides a controlled environment to verify that the method can accurately recover latent dynamics. The application to real neural recordings using cross-validated metrics (normalized log-likelihood, R²) offers a framework for assessing performance on recorded data with unknown latent dynamics.\n\nThe authors could further report forward prediction metrics.", "theoretical_claims": "There are no theoretical claims in the paper.", "experimental_designs_or_analyses": "No", "relation_to_broader_scientific_literature": "The paper integrates and extends several strands of research in computational neuroscience and machine learning. First, it builds on previous approaches for modeling neural population activity via low-dimensional latent variables, such as GPFA, LFADS, and variants of switching linear dynamical systems. Most of available frameworks consider either latent deterministic or linear dynamics or autonomous dynamics. The work further builds on the literature of neural ODEs (and related differential approaches) by adapting these techniques specifically to spiking neural data.", "essential_references_not_discussed": "- Genkin, M., Hughes, O., & Engel, T. A. (2021). Learning non-stationary Langevin dynamics from stochastic observations of latent trajectories. Nature communications.\n- Genkin, M., & Engel, T. A. (2020). Moving beyond generalization to accurate interpretation of flexible models. Nature machine intelligence.\n- Schimel, M., Kao, T. C., Jensen, K. T., & Hennequin, G. (2022). iLQR-VAE: control-based learning of input-driven dynamics with applications to neural data. ICLR.\n- Zhao, Y., & Park, I. M. (2020). Variational online learning of neural dynamics. Frontiers in computational neuroscience.", "other_strengths_and_weaknesses": "Strengths:\n\n- Well-organized and clearly written paper.\n- Good performance on demonstrated datasets.\n- Consistency of estimated dynamics across splits.\n- Needs more evidence, but the bias estimation part is interesting, and to my knowledge the approach is novel.\n\n\nWeaknesses:\n\n-  The method is rather technically involved, with multiple steps and hyperparameter dependencies, potentially limiting accessibility.\n- Although a hyperparameter grid-search is described, additional explicit analyses demonstrating robustness or sensitivity across hyperparameter choices and initialisations is missing.\n- Also see my questions below.", "other_comments_or_suggestions": "The authors comment on the identifiability of the observation model in the appendix, but they don't discuss the identifiability of the latent dynamics (drift and diffusion of the SDE) when observed through spikes.\n\n\nI wonder how would the method compare in terms of accuracy and required amount of data when compared with the method off Dunker et al 2019, and Genkin and Engel 2020, when applied on a synthetic dataset with autonomous latent dynamics.", "questions_for_authors": "- In FINDR, the SDE’s time constant $\\tau$ (as used in Eqs. 2–3) sets the timescale of the intrinsic latent dynamics. Do you estimate $\\tau$ directly from the data, or is it fixed a priori? Have you observed any relationship between this timescale and factors such as the required number of recorded neurons/heir average firing rates for accurate inference of the latent dynamics?\n- Do you assume that the external inputs $u$ are known? If not, how do you estimate them? If you do assume that they are known, how do you set them up in the examples in Figure 3 and 4?\n- I have a few questions regarding the estimation of the bias term of the observation model, that the authors name \"task irrelevant activity\". The authors infer first this bias, before learning the dynamics and the other parameters.\n    - 1] How uniquely can you identify the bias term of each neuron, as opposed to attributing a shift of the dynamics part to this bias? \n    - 2] Does the proposed approach estimate the d values that were used when generating the spike trains in the simulated experiments, or are they just identified good enough to allow estimation of the latent dynamics? Can the authors provide a scatter plot with the d_estimated vs d_true values to demonstrate how well this part of the method works?\n  - 3] Moreover I think I would be useful to make an ablation study for this part of the estimation to test how well the inference method would work without estimating d, or without the individual across trial and within trial d components.\n\n- In the paper of Genkin,  & Engel mentioned above, the authors make a discussion regarding overfitting of latent stochastic dynamical models with increasing number of spiking data. Do you observe something similar in your framework requiring to adjust the parameter $/beta$, or does the selection of $\\beta=2$ suffice to overcome this issue?\n\n- In the model parameters to be inferred is the noise covariance matrix $\\Sigma$ of the latent process. How do you estimate $\\Sigma$ and how accurate is this estimation in the synthetic dataset?\n\n- How necessary is the gating function of the gated neural ODE for accurate inference of the latent dynamics? Would a nODE without $G$ fail capture the latent flow field?\n\n- What are the computational demands of each part of the proposed inference method, and how does is compare to the methods used as baselines?", "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": "CEJMqetvtp", "reviewer": "Reviewer_PTUj", "summary": "FINDR (Flow-field Inference from Neural Data using deep Recurrent networks) is an unsupervised method for discovering low-dimensional neural population dynamics. It uses a gated neural drift function, decomposing spiking activity into task-relevant latents and a time-varying bias capturing non-task effects. A sequential variational autoencoder ensures accurate data reconstruction and interpretable dynamics. FINDR offers an interpretable approach for modeling complex neural computations through low-dimensional flow-field visualizations and analysis.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback. We will change Fig. 1 to include d in the revision.\n\n**Q1:** We should have clarified that the colored trajectories in Fig. 4b represent trial-averages sorted by evidence strength, and inside the dotted line represents part of the state space *visited* by single-trial trajectories. We did not show part of the state space not visited by single-trial trajectories. \n\nIn the revision, we could show a “confidence heatmap” showing the probability of the single-trial trajectories visiting a particular region of the state space.\n\n**Q2:** By design, task-irrelevant components do not use the task-relevant inputs u to capture within- and across-trial baseline firing activity. We reasoned that by optimizing d first, we let the model predict neural activity as much as it can without task-relevant information. If this part is successful, then the dynamics capturing the residual of the task-irrelevant component from inputs u would be task-relevant.\n\n**Q3:** The Euclidean distance in the neural state space (ignoring softplus) is preserved in the latent space. For more details, please see Appendix A.3 and response to Reviewer UnSA on identifiability.\n\n**Q4:** We perform hyperparameter grid search to ensure that we analyze latent representations from a model with good training and performance (Appendix A.1.6). We found that the optimal size of the MLP, and other important hyperparameters depend on the dataset.", "claims_and_evidence": "The authors claimed that FINDR is able to compress the neural data into a abstract low-d latent space and learn the dynamics in the latent space, which have been supported by their experiments.", "methods_and_evaluation_criteria": "Yes, the cross-validated log-likelihood and PSTH-based R² are standard metrics for spike train models, helping verify predictive accuracy and interpretability.\nThe validity of their main method: gated drift dunction was verified in previous litratures, not included in the present work.", "theoretical_claims": "The main theoretical aspect is that the gating structure of the drift function improves the expressivity, stability, and trainability of the neural SDE, referencing prior work in gated neural ODEs, while there are no formal proofs in the text.\nOther claims seemed correct.", "experimental_designs_or_analyses": "The synthetic data experiments were designed to confirm whether FINDR could recover known continuous attractors, a standard check for dimensionality and topology in latent-variable models. The real-data analysis employed a five-fold cross-validation scheme, ensuring robust validation of held-out neurons. Both approaches appear well-executed and adhere to common practices (e.g., log-likelihood and PSTH R²). A minor limitation is that the real-data setup features only a two-choice decision-making task, which might not reflect more complex or multidimensional behaviors.", "relation_to_broader_scientific_literature": "FINDR extends classical state-space models (LDS, GPFA) and deep generative approaches (LFADS) by using a gated nonlinear architecture that separates task-irrelevant components. This design builds on continuous-attractor concepts and addresses mixed selectivity, yielding stable, interpretable latent spaces and contributing new insights into neural dynamics and computational neuroscience.", "essential_references_not_discussed": "not to my knowledge", "other_strengths_and_weaknesses": "- Strengths\n\n1.The paper is well organized, with a clear exposition of the model through equations and figures.\n\n2.It is novel in how it explicitly separates task-irrelevant features, an approach that helps maintain a stable latent space across trials.\n\n3.Both synthetic and real datasets are used to verify the model, enhancing confidence in its validity.\n\n- Weakness\n\nThe real-data experiments focus on a two-choice decision-making task, which may be relatively simple.", "other_comments_or_suggestions": "To my understanding, the decoder should take d as input to reconstruct the firing rates, but d is not included in the figure.", "questions_for_authors": "1.The flow field is learned from encoder-inferred trajectories reflecting real data; how accurate might the learned dynamics be in latent regions not sampled by the dataset?\n\n2.Why is it necessary to optimize d first, then the other parameters? How does the approach ensure that task-relevant and irrelevant features are distinctly captured by d vs. the latent variables?\n\n3.To what extent does FINDR preserve geometric relationships in the latent space? For instance, if an animal navigates a T-maze, would the model maintain the T-shaped structure in the low-dimensional representation?\n\n4.How large must the MLP be (e.g., number of layers and hidden units) for stable training and robust performance?", "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": "ZpKfOga2Hi", "reviewer": "Reviewer_Hy3r", "summary": "Authors introduce a new method for latent variable inference of neural data. The essence is a sequential variational autoencoder. The main innovation is a “prior” which encourages the latent variables to satisfy an ODE. Using this method, they show that low-D latents are recovered in synthetic examples. When compared with other methods that are also limited to low-D latents, the performance on real data is better for very low-D, and comparable for slightly higher dimensions.\n\n## After rebuttal\n\nAfter reading all the rebuttals and discussion with all reviewers, I am keeping my score. The topic is important, and many methods have been introduced in recent years. While the proposed method here is novel and seems promising, the limited comparison to other methods and benchmarks weakens the contribution.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback. We appreciate their suggestion to test FINDR on the Neural Latents Benchmark (NLB) to further support our claims. We agree that evaluating FINDR on multiple real datasets, including the public training and validation datasets available from NLB, could strengthen our findings. However, as the reviewer pointed out, we think that one of the main goals of our paper is to be able to discover interpretable dynamics (i.e., performing analyses similar to Fig. 4), rather than solely outperforming other models on neural activity prediction (though we should make sure that our model performs reasonably well on this). Since analyses similar to Fig. 4 for the models currently fit to datasets in NLB (for models where this is possible, like rSLDS) are not readily available, we think the dataset we use is as good a choice as datasets in NLB.\n\nThe winning model of the NLB challenge in 2021 was based on transformers, which does not give a dynamical interpretation. Similarly, autoLFADS (with or without the rank bottleneck) doesn’t enable plotting low-dimensional vector fields (even when L<=3). We suspect that transformers and autoLFADS with L>20 will outperform FINDR given that autoLFADS with L=20 performed similarly to FINDR with L=2 in terms of R^2 (Fig. 4a). However, given that FINDR performs reasonably well in predicting responses from held-out neurons (Fig. 3c), in addition to being able to provide low-dimensional flow fields, we think this is where the strength of FINDR lies. We will clarify in the main text that FINDR only outperforms existing methods in low dimensions, and revise the abstract so that it reads “we demonstrate that FINDR performs competitively against existing methods…”.\n\nRegarding A.2, thank you for pointing this out. We should have been more clear. The rank of `C` is computed with `numpy.linalg.matrix_rank`, which computes the SVD and counts the number of singular values that are greater than some value `tol` to compute the rank. By default, `numpy.linalg.matrix_rank` sets `tol = S.max(axis=-1, keepdims=True) * max(C.shape[-2:]) * numpy.finfo(S.dtype).eps`, where `S` is the vector of singular values. While, for example, the column for L=6 in Fig. 2d shows that variance explained is close to 1 after L=2, the singular values of `C` here were 96.303, 86.11196, 9.6449, 8.212957, 7.338667, 1.3648185, which were all greater than `tol` which, for this `C`, was 0.0057.\n\nRegarding Equation 3---thank you for catching this! This is a typo, and it should have been \\sqrt(dt).", "claims_and_evidence": "The main claims are partially supported. The method is able to recover low-D synthetic latents. But, the abstract claims that “FINDR outperforms existing methods in capturing the heterogeneous responses of individual neurons”. Existing methods (e.g. LFADS) were not tested without the limitation of low-D latents.", "methods_and_evaluation_criteria": "There are existing benchmarks (neural latents benchmark, also cited in this paper) that offer quantitative measures of prediction (without limiting latent dimensionality). These were not used in the present work.", "theoretical_claims": "Irrelevant", "experimental_designs_or_analyses": "The synthetic example is indeed a case where we expect a 2D continuous attractor, and the method recovers it successfully. The irrelevant dynamics (constant bias) are perhaps not challenging enough.\nAs mentioned above, it would be good to test the method on other benchmarks, such as neural latents.", "relation_to_broader_scientific_literature": "This is part of the latent inference line of work. The authors mention these other works. \nThe emphasis here is on a low-D latent space, which can improve interpretability.", "essential_references_not_discussed": "Two works that also emphasize the low-D latents.\nValente, Adrian, Jonathan W. Pillow, and Srdjan Ostojic. “Extracting Computational Mechanisms from Neural Data Using Low-Rank RNNs.” Advances in Neural Information Processing Systems 35 (December 6, 2022): 24072–86.\nPals, Matthijs, A. Erdem Sağtekin, Felix Pei, Manuel Gloeckler, and Jakob H. Macke. “Inferring Stochastic Low-Rank Recurrent Neural Networks from Neural Data.” arXiv, February 26, 2025. https://doi.org/10.48550/arXiv.2406.16749.", "other_strengths_and_weaknesses": "Main strength is a new approach to encourage interpretable latents, which is an important goal. The comparison to other techniques and other benchmarks is somewhat lacking, which weakens the paper.", "other_comments_or_suggestions": "none", "questions_for_authors": "In the identifiability section (A.2), you write that rank(C)=L. How does this fit with the results of figure 2D for L>2 ?\nEquation 3 – should this be sqrt(dt)?", "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": "vvMbmB3E7d", "reviewer": "Reviewer_UnSA", "summary": "This work sets out to infer the latent variables and their time dynamics from observed neural recordings.  To achieve this, the authors developed FINDR, essentially a recurrent neural network with multi-layer perceptrons defining the flow maps of the latent variables, whereas neural activities are defined as linear projections plus a soft plus (an analytically nicer form of relu) from the latent variables. The key contribution to the literature seems to be the use of a specific deep architecture, though similar approaches exist in the published literature.\n\n## After the rebuttal\n\nI believe that, following the promised revisions, this work now meets the criteria for acceptance at ICML. That said, I still feel that the clarity and overall impact of the paper could have been significantly improved had the authors chosen to present the full scope of the work in a single manuscript. The mention of a companion bioRxiv paper on scientific applications of FINDR (introduced only in the author response, see: \"we discuss scientific findings using FINDR in a separate bioRxiv paper\") clarified some of the unease I initially experienced while reviewing this submission. While this fact explains certain omissions that I felt would have benefited this work significantly (which is one big confound we can never account for in our reviews), it also highlights how a more integrated presentation could have elevated the contribution well beyond the acceptance threshold.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. While we can’t update our submission now, we will revise it accordingly.\n\n**Identifiability:**\nOur definition (Appendix A.2) differs from [1], which builds on Roeder et al., 2021 ([A]), in two ways:\n1. Following Wang et al., 2021, we define latent z as identifiable if z_1 != z_2 implies that p(y|z = z_1, \\theta) != p(y|z = z_2, \\theta), where y is the observed neural activity (this is Eq. (1) and L636). Since we use a linear projection C for p(y|z, \\theta), this is satisfied if C in Eq. (1) is injective. This is a weaker condition than [A], which requires both z **and** \\theta to be identifiable. For us, this is difficult as there is always some invertible A such that y = Cz = CAA^{-1}z. This was our motivation for performing SVD on C = USVt, and defining z_tilde = SVt z so that the distance in the space z_tilde is preserved in the space y (A.3). All FINDR analyses use z_tilde. FINDR’s latents (z_tilde) are thus identifiable up to an orthogonal transformation.\n2. FINDR is a generative model, while CEBRA is a discriminative model. [1] and [A] focus on conditions for linear identifiability in discriminative models.\n\nWe will move A.2-A.3 to the main text, and clarify differences from [1].\n\n**Fig. 3:**\nWe regret if there has been a misunderstanding. We *did* perform 5-fold CV to optimize hyperparameters of all models (paragraph starting in L332, 1st column). We used a hyperparameter-optimized LFADS (**auto**LFADS), using config as in L796. For SLDS and rSLDS, we optimized over discrete latents (L308).\n\nFollowing Pei et al., 2022, we held out 20% of neurons by supplying the encoder with 80% held-in neurons. Then, the decoder in Eq. (1) reconstructed *all* neurons. We trained FINDR on 3/5 of the trials on all neurons, validated it on 1/5 of the trials on all neurons via grid search (A.1.6), and tested performance on the remaining 1/5 on held-out neurons. We performed the same procedure for autoLFADS. In Pandarinath et al., 2018, they used a linear decoder with exponential nonlinearity, and *not* a nonlinear decoder.\n\nWe will clarify details on evidence-conditioned PSTH R^2, normalized log-likelihood, and held-out neuron training.\n\n**Fig. 4:**\nHere, axes represent PC 1 and PC 2 of latents. For autoLFADS, the latent trajectories projected onto these axes were not consistent, whereas for FINDR, they were. To further evaluate consistency across folds, we sorted single-trial trajectories by evidence sign and computed the trial average of each group. Then, we calculated Pearson’s |r| of these trajectories between fold 1 and 2 for each latent axis and took the average of |r| across the axes. With this, FINDR folds were consistent by 0.99. In contrast, for autoLFADS (after doing A.3 transformation, just like FINDR), this was 0.53. \n\nUnlike the transformation in A.3 which preserves distances in the latent space, linearly transforming autoLFADS fold 1 to match fold 2 increased |r| to 0.99, but, by doing linear transform, we stretch the latent space, so the distance in latent space is *not* preserved in the neural space (ignoring softplus). Without A.3, we wouldn’t be able to say e.g., the first latent dimension explains most of the variance of the task-relevant component of neural data.\n\n**Benchmarking & Expressivity:**\nAs the reviewer points out, [1] is static and is not a dynamical model in the sense that it does not learn representations like our Eq. (2), and we can’t perform fixed point analysis. This is similar for [3]. The model in [2] is a finite-dimensional LDS with a nonlinear decoder, meaning it can’t learn nonlinear dynamics like bistable attractors.\n\nRegarding explanation power (expressivity) of low-rank RNNs and benchmarking RNNs, Kim et al., 2023 defines a measure of practical expressivity and performs extensive analyses comparing RNNs of different architectures, including the one we use. Low-rank RNNs are a special case of our single-hidden-layer MLP without gating (Mastrogiuseppe & Ostojic, 2018).\n\n**Fig. 2:**\nFig. 2 *doesn’t* show bistable dynamics, but a disk attractor (L203, 1st column). A 1-D variant similar to this task would generate a line attractor. We find that FINDR can also recover limit cycles and bistable attractors in synthetic data.\n\n**FINDR’s goals:**\nWe will clarify that the two goals are field-wide, not just FINDR’s. We will make our contributions, including task-relevant and irrelevant dynamics, clearer with bullet points.\n\n**Flow map confidence:**\nPlease see Reviewer PTUj’s Q1.\n\n**Relevant work:**\nWe will cite Hu et al., Pals et al., and MARBLE as important concurrent work, and cite references mentioned by the reviewer. In particular, MARBLE estimates flow fields in neural space before embedding them in latent space, whereas FINDR estimates flows in latent space. Estimating flows in high-D space could be more sensitive to noise, and integrating state-space modeling (like the one here) with MARBLE could be an interesting future direction.", "claims_and_evidence": "I believe that some of the claims made in this work are not clearly supported by empirical or theoretical evidence. Specifically:\n\n- Identifiability: It is not clear what latent identifiability is, or how it is achieved. To my understanding, CEBRA ([1]) has solved a significant problem in this literature by proving linear identifiability in latent variables. This claim should be made more clear, as it is central to the latent variable identification literature.\n\n- Fig 3 experiments: The benchmarked models seem to be not optimized, rather some random parameters taken from demos are used? If true, model parameters should be optimized using cross-validation. The authors should also provide clear information on how projections from held-out neurons were trained for all methods. For instance, LFADS use nonlinear decoders. Did the authors train linear readouts? If so, this is also suboptimal. Overall, there is little to no information in the paper about the methodology behind Fig. 3.\n\n- Fig. 4: This figure does not support the claim that other methods do not find consistent representations across folds. For this claim, I believe there should be some form of statistical tests. Moreover, since latent variables are often only linearly identifiable, it is likely that the autoLFAD results might be quite consistent after linear transformations.", "methods_and_evaluation_criteria": "To my understanding, the field has moved away from LFADS. On the other hand, rSLDS are mainly used due to their interpretability, not computational power. Also, please see the latest implementation for rSLDS that seem to provide significant improvements compared to the traditional use (https://nips.cc/virtual/2024/poster/95587). \n\nIn general, I believe benchmarking should include CEBRA [1], or models derived from CEBRA-like architectures that also incorporate dynamical modeling of latent states (as is done in this work). See for instance, [2-3]. Additionally, I think at least one component of this paper should include a (low-rank) RNN benchmark. See for instance [4-5]. Before moving towards the deep recurrent network, one might expect whether shallow low-rank RNNs could have similar explanation power. For instance, [5] has shown a similar result as in Fig. 2 in this work.", "theoretical_claims": "NA.", "experimental_designs_or_analyses": "For Fig. 2, it may be more interesting to add other examples than the flip flop task. To solve flip flops, network simply generate bistable dynamics. Tasks such as sin-generation (limit cycle) and delayed addition/multiplication (line attractor) could bring additional breadth to this work and increase the appeal to the broader NeuroAI community.", "relation_to_broader_scientific_literature": "As noted earlier, the key contribution to the literature was somewhat unclear to me. Authors state that \"The goal of FINDR is to 1) compress the activity of a large population of neurons at time t to an abstract low-dimensional representation, and 2) learn the “rules” of how this representation evolves over time. \" \n\nAs stated, these are very broad statements that apply to many works in the field [1-5]. I was not able to identify these works as cited in the manuscript, and I think there are several other relevant ones that are cited within these works.", "essential_references_not_discussed": "Please see the end of this report. The current manuscript is missing some key references. This list is not exhaustive, but may be helpful.", "other_strengths_and_weaknesses": "As a big strength, I want to note that goal 2 stated by authors is very interesting! A fully static model like CEBRA cannot address this, though later variations did try to advance CEBRA in this regard. That being said, I always wondered the following question when reading papers with such claims: How we can trust the flow maps outside the regions that data is observed? In relevant literature, [6-7] have proven this for rSLDS by performing optogenetics manipulations. I am not sure if such drastic experiments are needed to support this claim, but as noted above, I do not believe Fig. 4 is sufficient either.", "other_comments_or_suggestions": "References:\n\n[1] Schneider, S., Lee, J. H., & Mathis, M. W. (2023). Learnable latent embeddings for joint behavioural and neural analysis. Nature, 617(7960), 360-368.\n\n[2] Abbaspourazad, H., Erturk, E., Pesaran, B., & Shanechi, M. M. (2024). Dynamical flexible inference of nonlinear latent factors and structures in neural population activity. Nature Biomedical Engineering, 8(1), 85-108.\n\n[3] Chen, C., Yang, Z., & Wang, X. (2025). Neural Embeddings Rank: Aligning 3D latent dynamics with movements. Advances in Neural Information Processing Systems, 37, 141461-141489.\n\n[4] Pals, M., Sağtekin, A. E., Pei, F. C., Gloeckler, M., & Macke, J. H. (2024, June). Inferring stochastic low-rank recurrent neural networks from neural data. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n\n[5] Valente, A., Pillow, J. W., & Ostojic, S. (2022). Extracting computational mechanisms from neural data using low-rank RNNs. Advances in Neural Information Processing Systems, 35, 24072-24086.\n\n[6] Vinograd, A., Nair, A., Kim, J. H., Linderman, S. W., & Anderson, D. J. (2024). Causal evidence of a line attractor encoding an affective state. Nature, 634(8035), 910-918.\n\n[7] Liu, M., Nair, A., Coria, N., Linderman, S. W., & Anderson, D. J. (2024). Encoding of female mating dynamics by a hypothalamic line attractor. Nature, 634(8035), 901-909.\n\nMore relevant work: https://www.nature.com/articles/s41593-020-00733-0, https://www.nature.com/articles/s41586-023-06714-0, https://www.nature.com/articles/s41467-018-06560-z, also see works from Durstewitz's and Ostojic's groups.", "questions_for_authors": "I have the following questions for the authors:\n\n1) Could you clarify how the model parameters for the benchmarks were chosen? Did you perform cross-validation to optimize them? How were the parameters for FINDR chosen?\n\n2) Could you address my comments in Claims And Evidence above: Specifically how held-out neurons were trained to connect to latent variables and what is meant here in terns of latent identifiability? \n\n3) How does your method compare to CEBRA ([1], published a while back) and MARBLE (https://www.nature.com/articles/s41592-024-02582-2) which just came out but has been on BiorXiv for a while? To clarify, for MARBLE, I am not asking for comparisons since it can be considered concurrent work, though should be cited as such, I am just asking for clarification.\n\nI believe this work can be a good fit for ICML, but as it stands, it requires substantial revisions. I am not sure if the limited interactions of the conference format allow such nuanced discussions. Hence, I recommend the authors address above concerns and resubmit to the next conference cycle, though if they choose to do a rebuttal, I remain optimistic and ready to change my evaluation if substantial evidence is presented.", "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": {"ethical_review_concerns": "NA"}}]

Author rebuttal:

Meta-review: This study develops a new method for latent variable inference from neural data. The essence is a sequential variational autoencoder. The main innovation is a type of prior which encourages the latent variables to satisfy an ODE. Using this method, the study shows that low-dimensional latent variables are recovered in synthetic examples. When compared with other methods that are also limited to low-dimensional latent variables, the performance on real data is better for very low-D, and comparable for slightly higher dimensions.
The reviewers appreciated that the paper makes a case for being accepted in ICML but also raised suggestions both of citation of literature and of additional analyses and results to be added, which are expected to be implement in a revision.
Decision: Accept (poster)