Flow-field inference from neural data using deep recurrent networks

We thank the reviewer for their valuable feedback. While we can’t update our submission now, we will revise it accordingly.

Identifiability: Our definition (Appendix A.2) differs from [1], which builds on Roeder et al., 2021 ([A]), in two ways:

1. 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.
2. FINDR is a generative model, while CEBRA is a discriminative model. [1] and [A] focus on conditions for linear identifiability in discriminative models.

We will move A.2-A.3 to the main text, and clarify differences from [1].

Fig. 3: We 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 (autoLFADS), using config as in L796. For SLDS and rSLDS, we optimized over discrete latents (L308).

Following 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.

We will clarify details on evidence-conditioned PSTH R^2, normalized log-likelihood, and held-out neuron training.

Fig. 4: Here, 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.

Unlike 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.

Benchmarking & Expressivity: As 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.

Regarding 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).

Fig. 2: Fig. 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.

FINDR’s goals: We 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.

Flow map confidence: Please see Reviewer PTUj’s Q1.

Relevant work: We 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.

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.

The 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…”.

Regarding 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.

Regarding Equation 3---thank you for catching this! This is a typo, and it should have been \sqrt(dt).

We thank the reviewer for the positive feedback. We will change Fig. 1 to include d in the revision.

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.

In the revision, we could show a “confidence heatmap” showing the probability of the single-trial trajectories visiting a particular region of the state space.

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.

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.

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.

We appreciate the reviewer’s thoughtful comments.

Necessity of gating: Kim 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.

Flow fields and interoperability: We 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.

Previous 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.

Task-irrelevant components and consistency in flow field: By 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.

Learning 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.

Limited accessibility: We will release our code with documentation and a tutorial notebook.

Model robustness: 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.

Identifiability: Please see response to Reviewer UnSA.

Additional comparisons with autonomous dynamics: While 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.

Questions

(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.

(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).

(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).

(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.

(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).

(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.

(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]].

(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.