SIMPL: Scalable and hassle-free optimisation of neural representations from behaviour

"Lines 477–479 are ambiguous as to whether the authors discuss place cell remapping or another concept. If remapping is the intended topic, I recommend explicitly stating this and including a citation to aid readers from the machine learning community who may be unfamiliar with this concept."

We apologise for the confusion. We were not refering to hippocampal remapping in lines 477-479 but were discussing how the trajectory makes a discontinuous "jump" to another position in latent space. We will try and re-word in the camera-ready to avoid confusion.

"The writing can be improved a lot"

Thank you for this very thorough check! We will certainly fix all of these errors in time for the camera-ready version.

"What the arrow of epoch 1 -> 10 in Figure 3 is different from others? Does it have a specific meaning?"

This just refers to the fact that epochs 2--9 are collapsed and not shown to reduce clutter.

"What does “data now shown in L466 mean? Does it mean that the authors intentionally did not include the results since it is insignificant? If so, please include it in the appendix."

Yes, we did no show this result because it is insignificant and we wanted to keep only the most important results in the main figure to reduce clutter. We will add these statistics in for the camera ready.

We thank the review for their time and thorough review. We would be very happy to hear about, and address, about any remaining major concerns they may have regarding technical contributions, soundness and impact. We remain commited to addressing any concerns that may prevent the reviewer from recommending acceptance via an updated score.

We thank the review for their comments. In response we have added new benchmarks with alternative techniques and new analysis on their suggested macaque dataset. As such we believe we have satisfied all of their concerns. We now respond to their comments point-by-point:

"The author claims general but limited to specific regions of brain place cells and grid cells. I highly recommend authors compare with various datasets such as the macaque dataset and the mouse visual cortex datasets used in Schneider et al. (2023)."

Following this review we have now tested SIMPL on one of the Neural Latents Benchmark datasets (specifically the Area2_Bump data from somatosensory cortex for a macaque doing a centre-out reaching task collected by Chowdhury and Miller), the same dataset as used in the Schneider et al. (2023) paper. We find SIMPL works well (see the revised manuscript for a more detailed discussion of the results). In the results we have now added a new section and figure detailing our findings briefly summarised as follows: We tested 3 versions of SIMPL on the hand-reaching data.

• SIMPL2D(position): The latent is initialised with the monkeys x- and y-hand position.
• SIMPL2D(velocity): The latent is initialised with the monekys $v_x$- and $v_y$-hand velocity.
• SIMPL4D(postion&velocity): A 4D latent is initialised with $x$, $y$, $v_x$ and $v_y$. In all three models SIMPL optimises the latent variable (the test-log-likelihood improves), uncovering a smooth latent variable correlated to (but substantially different from) the behavioural initiates. Corresponding tuning curves revealed neurons with "hand-position-like" or "hand-velocity-like" selective receptive fields. The 4D version of SIMPL performed better than either 2D version, revealing disentangled latents with a higher overall log-likelihood than either of the 2D models. Our finding reveal:

1. SIMPL can be applied to the types of non-hippocampal datasets commonly used in the LVM literature.
2. Both position and velocity seperately do a good job explaining the latent dynamic but position and velocity combined is better.
3. In the optimised latent space neurons have localised receptive fields reminiscent of place cells in the hippocampus.
4. SIMPL can be extended beyond the 2-dimensional latent spaces initially tested.

"The authors misuse big-O notation."

We'll correct this for the camera-ready version.

"3. Except for Section 4.2, there are no comparisons with existing methods"

Thank you for this suggestion, which was echoed by reviewer gB66 (where you will find a more detailed discussion). In summary, we have added two more comparsion to well known techniques (piVAE and GPLVM) which are, arguably, more relevant. In both cases these models perform well on the synthetic grid cell dataset but still worse than SIMPL and with over 10x compute cost. We think these additional benchmarks substantially strengthen our claim that SIMPL outperforms the most relevant, popular alternatives.

"Although experiments with Tanni et al. (2022) represent the main result in this paper, the authors scarcely describe the task and the significance of each plot (particularly in Figure 6c), simply referring back to the original paper."

We will add a section to the discussion providing more details of this task. Regtarding panel 6c this shows violin plots of summary statistics for the tuning curves before and after SIMPL is applied, showing how they have changed. We'll clarify this for the camera ready.

"The authors compare SIMPL with CEBRA, a machine learning-based method running on CPUs. If they provide a runtime comparison showing that SIMPL on a CPU is faster than CEBRA on a GPU, it would further support their efficiency claims."

Thank you for this suggestion however we don't think this is an appropriate or fair comparison. Because the compute-heavy step in SIMPL (calculating the likelihood maps, see Appendix for why) is parallelizable, we would also expect huge speed up for SIMPL on a GPU. Since it is all written in JAX this should be readily acheivable. If the reviewer still thinks it worthwhile we can add SIMPL-CEBRA GPU-vs-GPU time comparison. On the other hand, we believe one of the core values of SIMPL is it's speed on a CPU. GPU usage is a barreir for some researcher and not everybody has access to to these compute resources so we would still like to focus attention on CPU compute times. We hope this argument makes sense.

"How much of the performance inferiority of CEBRA VS SIMPL comes from not carefully searching through the hyperparameter space? Is the comparison fair? Would increasing the iterations improve the performance of CEBRA?"

We don't think so. The relative underperformance of CEBRA is not because it hasn't converged but rather because it does not assume any dynamics on the underlying latent variable which would "smooth" (or "denoise") it. We made this comment in the original submission but will make it clearer in the camera-ready version. In CEBRA each spike bin is treated independently and is thus subject to its own irreducible noise. Conversely, SIMPL (and other dynamical LVMs) assumes the latent follows a linear dynamical systems thus employs Kalman smoothing which denoises the latent substantially.

"How could the author disentangle the possibility that this observation is a result of artifacts by their method versus reflecting true representation feature of the hippocampus?"

This is a an important question which we spent time addressing in the original submission (see paragraph 4 of the section on hippocampal dataset). In summary, we ran SIMPL on a control dataset of spikes sampled from behaviour and behaviour-fitted tuning curves. Results are shown in a grey shade on Fig. 6. For these spikes behaviour and behaviour-fitted tuning curves should be stable (since they are, by definition, exactly the generative model) and any changes reflect artifacts of the SIMPL algorithm. We observe no significant changes to the place fields for the control spikes suggesting the changes observed in the real neural data are not artifacts

Lastly, we thank the reviewer again for their thorough review and would appreciate if they would reconsider their score in light of our new additions to the manuscript which, we believe, have strengthened it substantially.

We thank the reviewer for their time take to thoroughly assess our paper. In response we have run new analyses as well has hyperparameter sweeps ot the manuscript. As a result we believe the paper is now in a much stronger position to be accepted. Our point-by-point response is as folllows:

"SIMPL has only two hyperparameters. Minimal hyperparameters make it practical for experimentalists."

Furthermore, we have now added a new figure to the Appendix (Fig. 8) where we sweep across these two parameters and show the model performs well across a large area of parameter space.

"Rigorous comparison with existing methods (CEBRA). It outperforms CEBRA (whose latent embedding was noisier and had larger final error) and is over 30 faster."

We have also added a comparison to GPLVM (and equivalent but Gaussian process based technique) and Pi-VAE as suggested by reviewer gB66.

"Limited Evaluation on Non-Spatial Tasks...[it] would benefit from evaluation on other tasks, such as the ones tested in the CEBRA paper"

Following this review we have now tested SIMPL on one of the Neural Latents Benchmark datasets (specifically the Area2_Bump data from somatosensory cortex for a macaque doing a centre-out reaching task collected by Chowdhury and Miller), the same dataset as used in the CEBRA paper. We find SIMPL works well (see the revised manuscript for a more detailed discussion of the results). In the results we have now added a new section and figure (Fig. 7) detailing our findings briefly summarised as follows: We tested 3 versions of SIMPL on the hand-reaching data.

• SIMPL2D(position): The latent is initialised with the monkeys x- and y-hand position.
• SIMPL2D(velocity): The latent is initialised with the monekys $v_x$- and $v_y$-hand velocity.
• SIMPL4D(postion&velocity): A 4D latent is initialised with $x$, $y$, $v_x$ and $v_y$. In all three models SIMPL optimises the latent variable (the test-log-likelihood improves), uncovering a smooth latent variable correlated to (but substantially different from) the behavioural initiates. Corresponding tuning curves revealed neurons with "hand-position-like" or "hand-velocity-like" selective receptive fields. The 4D version of SIMPL performed better than either 2D version, revealing disentangled latents with a higher overall log-likelihood than either of the 2D models. Our finding reveal:

1. SIMPL can be applied to the types of non-hippocampal datasets commonly used in the LVM literature.
2. Both position and velocity seperately do a good job explaining the latent dynamic but position and velocity combined is better.
3. In the optimised latent space neurons have localised receptive fields reminiscent of place cells in the hippocampus.
4. SIMPL can be extended beyond the 2-dimensional latent spaces initially tested.

"Scalability Concerns for High-Dimensional Latents...What alternatives might be suitable when dealing with truly high-dimensional data?"

It is true that we remain cautious about SIMPL's performance in very high-dimensional latent spaces. This is primarily because, as a function approximation technique, KDE suffers from the "curse of dimensionality". We believe that neural-network based approaches might have the potential to perform better in such high-dimensional scenarios, a point we made in the orignal submission but will now further clarify. However, these models are harder to optimize, often requiring tedious hyperparameter tuning, and may put off prospective scientists.

"How does the method perform with smaller neural populations?...How much data is required to optimize SIMPL?"

We thank the reviewer for bringing this important point to our attention. These are important questions. We repeated our synthetic grid cell experiment using decreasing amounts of data (fewer neurons and shorter duration), the results are shown in a new subpanel of Fig. 3 (panel e). In summary we find that the size of our original dataset (225 neurons, 60 minutes) was much larger than actually required. SIMPL still performs well (and runs much faster) with only 50 neurons and 10 minutes of data. Of the two, we found that number of neurons is more important than the duration, with performance dropping off sharply when $\leq$ 20 neurons are used. This is easily achieved by most modern neural datasets and performance of SIMPL in lower data regimes is empirically confirmed in our new analysis on a somatosensory dataset which has only 65 neurons and 37 minutes.

As the discussion period is coming to a close we would like to ask whether the reviewer has had the time to consider our comments and, in light of these, reconsider their score. To summarise; along with a point-by-point discussion to each of the reviewers questions we believe we have handled their most important concerns by making the following changes to the paper:

• Adding a new analysis/figure where we analyse a non-spatial motor-task dataset.
• Adding a new subpanel to fFig. 3 where we test performance against dataset size, and show SIMPL works well with significantly smaller datasets.

Full details can be found in our original response above. Please would the reviewer let us know if they have any remaining major concerns.

• GPLVM: We find that GPLVM performs quite well on our synthetic grid cell dataset but terminates with a higher error than both SIMPL and CEBRA, probably due to its misspecified emission model meaning it does not properly account for the Poisson nature of the spiking data. Furthermore, to give GPLVM the best chance and put it on level footing with CEBRA and SIMPL we initialised the its latent with behaviour. Note that in order to make use of all datapoint we had to restrict ourselves to using 1000 inducing points. Total compute time was just under 12 minutes.
• Pi-VAE: Performed well, beating both CEBRA and GPLVM but finishing with an overall error (8.38 cm) about twice that of SIMPL and a compute time about 15.8x higher (we used a CPU for both). Like CEBRA, but less so, it's latent appears noisy which, also like CEBRA, we expect is due to the fact there is no explicit temporal dynamics/smoothing.

We agree with the reviewer that these two methods are both more relevant that CEBRA which we initally compared to because of its popularity and ease of access. We will leave CEBRA there for completeness but we are happy to move it to the Appendix if the reviewer thinks this would help the clarity of the paper. Another viable contender could be a technique such as PfLDS (Gao, 2016). This satisfies the first three of our desiderata, but, to our knowledge, is not identifiable in the sense that there is no meaningful way to give behaviour as an input.

Finally, related to benchmarking, we would also like to draw the reviewers attention to a new analysis and figure (Fig. 7), where we apply SIMPL to a hand-reaching macaque dataset from somatosensory cortex. This dataset on which SIMPL performs well, from the Neural Latents Benchmark suite, is commonly used in the LVM literature further supporting our claim that SIMPL is performant in comparison to similar LVM methods.

We hope that our response and additional comparisons cleared the reviewer's concerns on this topic. We are happy to further clarify the position of SIMPL in the spectrum of data analysis methods by incorporating elements of this response in the Related Work section for a camera-ready version. We will be happy to address any remaining concerns the reviewer may have. Otherwise, we would be appreciate if the reviewer could adapt their score and recommend acceptance of our submission. Once again, we thank the reviewer for their time and insightful comments which, in our opinion, have led to substantial improvements in the manuscript.

We thank the reviewer for their very thorough review and for pointing us to this additional literature on latent variable modelling. We will cite these works in the camera-ready version of the paper. We provide below some important clarifications regarding which methods constitute relevant alternatives to SIMPL, as well as describing additional benchmarks we have run in response to the reviewers valid concerns.

From a modelling standpoint SIMPL has four key properties which we view as desiderata for comparable techniques:

1. Complex, non-linear tuning curves
2. Time dependence of the latent variable.
3. Poisson emission probabilities
4. Identifiability from behaviour

By identifiability we mean whether there are any reassurances (theoretical or empirical) that, given behaviour, the model can recover the true latent up to some affine transformation. Many models simply do not admit any way to consume behavioural information.

We see desiderata 1 as non-negotiable since we are primarily focused on modelling cells with complex tuning curves (e.g. grid cells). Any model with a substantially restrictive intensity functions (specifically linear-type, of the form $y_t \sim \textrm{NoiseModel}(f(\mathbf{M}\mathbf{z}_t+\mathbf{c}))$ where $f$ is a simple function e.g. identity, softplus, exponential etc.) can never interpretably account the sorts of datasets (e.g. grid cells) we consider here. With that in mind we carefully checked all the references that the reviewer pointed us to. None of them satisify all these desiderata and, in particular, five of them (2, 3, 4, 5 and 7 as numbered in the table below), including LFADs, have linear-type tuning curves and, as such, we consider them mis-specified. Unless we have misunderstood, or the reviewer has a strong counter-reason, these method simply could not model the types of tuning curves we are interested in optimising and therefore we do not see them as valid comparisons. We will make this point much clearer in an updated manuscript.

The following table summarises all papers mentioned by the reviewer and their status with respect to our four desiderata:

To be clear, the fact none of these references satisfy all four desiderata does not mean that all these methods aren't useful in cases where SIMPL could be used however it means that we believe, independent of explicit comparisons, that SIMPL still represents a valuable contribution to the field.

More broadly, the closest alternative to SIMPL is possibly Poisson-GPLVM (Wu et al. 2017). However, unfortunately, the current P-GPLVM algorithm doesn't make use of inducing points, resulting in cubic complexity. We tried running this on 1000 datapoints and got a compute time of ~2 mins, already exceeding the runtime of SIMPL on all 36000 datapoints by a factor of ~3. We would be happy see future work directed at improving the scalability of P-GPLVM however, in this work, we have taken a different approach by defining a model that bypasses GPs altogether.

Nonetheless, we take seriously the reviewers suggestion that we need to improve our comparisons of SIMPL to alternatives. For this reason we have added comparisons to two other methods which only lack one of the desiderata listed above. Pi-VAE which imposes time independence on the latent (i.e. only fails desiderata 2) and GPLVM which assumes Gaussian emissions (desiderata 3). Both come with well maintained code bases. The GPLVM algorithm can exploit inducing points to make it scalable. Results are shown in the updated version of the paper. In summary:

(continued in next comment...)

As the discussion period is coming to a close we would like to ask whether the reviewer has had the time to consider our comments and, in light of these, reconsider their score. To summarise; the reviewer's major concern was that we had not benchmarked SIMPL against sufficient nor relevant existing methods. Largely we were in agreement with the reviewer and, as such, responded with:

• Two new benchmarks against more relevant LVM techniques mentioned by the reviewer: pi-VAE and GPLVM-with-inducing-point.
  • These techniques work well on our dataset but still significantly underperform SIMPL in terms of both absolute final error and compute time.
  • These have been added to Fig. 5.
  • Both our new benchmarks (assuming a fixed number of inducing points for GPLVM) have linear time complexity, matching SIMPL. But our experiments suggest they are both slower by approximately a factor of 10x.
• A thorough review of the literature flagged by the reviewer
• Substantial discussion on what methods constitute relevant benchmarks and why.

Full details can be found in our original response above. We thank the reviewer again for their suggestions and ask if they would let us know if they have any remaining major concerns.

We thank the reviewer for their time taken to assess our paper. We were especially happy to read, and would like to reiterate, their point about SIMPL being a "fresh departure" from previous techniques. We have made a number of changes to the manuscript in response to their review, several new figures and sub-panels as well as an entirely new analysis on a non-hippocampal dataset from the Neural Latents Benchmarks suite. We hope that these changes address their concerns.

"Higher dimensional latent states...and datasets in the current literature on LVMs"

Following this review we have now tested SIMPL on one of the Neural Latents Benchmark datasets (specifically the Area2_Bump data from somatosensory cortex for a macaque doing a centre-out reaching task collected by Chowdhury and Miller), the same dataset as used in the CEBRA paper. We find SIMPL works well (see the revised manuscript for a more detailed discussion of the results). In the results we have now added a new section and figure (Fig. 7) detailing our findings briefly summarised as follows: We tested 3 versions of SIMPL on the hand-reaching data.

• SIMPL2D(position): The latent is initialised with the monkeys x- and y-hand position.
• SIMPL2D(velocity): The latent is initialised with the monekys $v_x$- and $v_y$-hand velocity.
• SIMPL4D(postion&velocity): A 4D latent is initialised with $x$, $y$, $v_x$ and $v_y$.

In all three models SIMPL optimises the latent variable (the test-log-likelihood improves), uncovering a smooth latent variable correlated to (but substantially different from) the behavioural initiates. Corresponding tuning curves revealed neurons with "hand-position-like" or "hand-velocity-like" selective receptive fields. The 4D version of SIMPL performed better than either 2D version, revealing disentangled latents with a higher overall log-likelihood than either of the 2D models. Our finding reveal:

1. SIMPL can be applied to the types of non-hippocampal datasets commonly used in the LVM literature.
2. Both position and velocity seperately do a good job explaining the latent dynamic but position and velocity combined is better.
3. In the optimised latent space neurons have localised receptive fields reminiscent of place cells in the hippocampus.
4. SIMPL can be extended beyond the 2-dimensional latent spaces initially tested.

"Settings where the behavioural recordings are missing or latent spaces not dominated by behaviour"

We have already shown in Fig. 4 that SIMPL still works well without behaviour. The "catch", in such instances, is that you lose the identifiability guarantees that come when a relevant behavioural variable is used for initialisation. Because of this we acknowledge and agree with the reviewer that the most likely use-case for SIMPL will be to discover/refine tuning curves in neural data which is well characterized by a behavioural variable already available to the user.

Although this is, on the face of it, a limitation, it reflects a very conscious choice made to bypass a fundamental issue in latent variable modelling. Latent variable modelling is fundamentally hard and an entirely general purpose technique which is (i) totally assumption free and (ii) computationally cheap and hassle-free across all datasets simply doesn't (perhaps never will) exist. We contest that many popular techniques have actually under relied on behaviour which has resulted in methods which are overly complex (putting off non-theoreticians) or compute heavy. Our observation is that in a majority of neuroscience experiments the latent space is closely related to a behaviour which is (or could be easily) measured concurrently - in such a majority of cases SIMPL strikes a good balance between simplicity, interpretability and computational efficiency. We will add a discussion of this point in the revised manuscript.

"Non-smooth latent trajectories"

We agree with the reviewer that it is important to test SIMPL when the latent variable is non-smooth. Currently the amount of smoothness in the latent trajectory is controllable via the velocity hyperparameter and can be set arbitrarily close to zero (at which point the decoding procedure amounts to pure maximum likelihood estimation, capable of modelling non-smooth latents).

To further alleviate concerns over whether SIMPL can model non-smooth latents, we generated a synthetic "replay" dataset. In this dataset the latent and behaviour are identical except for regular, brief instances where the latent siccontinuously jumps to a new location in the environment then jumps back. With the same parameters as before SIMPL is able to accurately recover the latent, correctly recapitulating the discontinuous jumps. In summary, although SIMPL is biased towards smooth latents this is just a prior and, with reasonable parameter settings, non-continuous latents can be modelled as well. These new results are summarised in a new figure (Fig. 9) in the appendix.

"The authors mention improved identifiability of the proposed approach. Is that solely because of the behavioral initialization?"

Yes, we find that the initialization near behaviour results in convergence on a latent space which is the same as the ground truth (at least, in our sythetic experiments, which are the only examples where true ground truth is knowable). Results in favour of this hypothesis include that fact that when behavioural intialisation is removed SIMPL learns a warped/fragmented latent space. As well as the fact that CEBRA, which was not initialised with behaviour, did not find such a good latent and it's grid fields appear slightly warped relative to the ground truth.

"The authors say: "we see SIMPL as a specific instance of a broader class of latent optimization algorithms, ..." and then move on to the idea of replacing KDE with neural networks. Can the authors elaborate on that?"

This point makes clear that there could, in theory, be many ways of fitting tuning curves. We have chosen KDE because it is simple, interpretable and fast but a neural network (e.g. trained to output firing rates given a latent) could be used instead and would come with its own advantages and disadvantages which we have not studied here. We have not tested this but it is an interesting avenue for future work. We will rewrite this section to make this clearer.

"Can the authors provide the intuition for choosing the kernel bandwidth in the KDE step? This seems critical for the method's performance."

There are some general heuristics for choosing the kernel bandwidth, for exampled see Silverman's rule of thumb. Roughly, the goal is to choose the smallest bandwidth allowed by the data without over- or under-fitting: too small and individual spikes will be resolved, too big and high-frequency structure in the receptive fields will be smoothed. In practice SIMPL is not particulary sensitive to this parameter as shown in our new hyperparameter sweep performed in respjse to the reviewers concerns (see appendix Fig. 8) where we find that performance is good across an order of magnitude of kernel bandwidths (between 0.5 cm and 5 cm) and an order of magnitude of speed priors (between 0.1 ms$^{-1}$ and 1 ms$^{-1}$). These ranges will, of course, be dataset dependent, but they do suggest that the range of appropriate hyperparameters might not be very sharp. In summary, our new hyperparameter sweep shows SIMPL has only soft dependence on the hyperparameter.

On the basis of our adjustments and additional analysis on the Neural Latents Benchmark dataset, we believe that the paper is now in a much stronger position to be accepted. Please would the reviewer let us know if they have any further concerns or questions will might prevent them recommending our paper for acceptance.

As the discussion period is coming to a close we would like to ask if the reviewer has had the time to consider our comments and, in light of these, reconsider their score. To summarise; along with a point-by-point discussion we believe we have handled their most important concerns by:

• Adding a new analysis/figure where we analyse a non-spatial motor-task dataset.
• Adding a new experiment/figure where we show SIMPL can account for discontinuous latent trajectories.
• Performing a sweep across the kernel bandwidth and velocity prior hyperparameters showing SIMPL performance is not sharply dependent on their values.

Full details can be found in our original response above. Please would the reviewer let us know if they have any remaining major concerns.

We'll take a look at these references and filter them into the manuscript for the camera-ready version.