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

Authors: Tom George, Pierre Glaser, Kim Stachenfeld, Caswell Barry, Claudia Clopath
Primary area: applications to neuroscience & cognitive science
Keywords: neuroscience; place cells; grid cells; representations; neural data; hippocampus;

Abstract:
Neural activity in the brain is known to encode low-dimensional, time-evolving, behaviour-related variables. A long-standing goal of neural data analysis has been to identify these variables and their mapping to neural activity. A productive and canonical approach has been to simply visualise neural "tuning curves" as a function of behaviour. However, significant discrepancies between behaviour and the true latent variables -- such as an agent thinking of position Y whilst located at position X -- distort and blur the tuning curves, decreasing their interpretability. To address this, latent variable models propose to learn the latent variable from data; these are typically expensive, hard to tune, or scale poorly, complicating their adoption. Here we propose SIMPL (Scalable Iterative Maximization of Population-coded Latents), an EM-style algorithm which iteratively optimises latent variables and tuning curves. SIMPL is fast, scalable and exploits behaviour as an initial condition to further improve convergence and identifiability. It can accurately recover latent variables in spatial and non-spatial tasks. When applied to a large hippocampal dataset SIMPL converges on smaller, more numerous, and more uniformly sized place fields than those based on behaviour, suggesting the brain may encode space with greater resolution than previously thought.

Decision: accept (poster)
Reviewer scores: 6, 8, 5, 6

Reviews and discussion:
<thread>
creation date:
1729955646119

thread_type:
official_review

<thread_object>
creator: reviewer_5qja

<content>
summary:
The paper proposes a new method, SIMPL for estimating neural representation regarding behaviors. The method is based on the Expectation and Maximization method where during E-step, a model produces estimates of latent trajectories using Kalman smoothing and during M-step, the model uses KDE for fitting the intensity function.

strengths:
* The paper is well-written. Especially, the references in the introduction can be a good guide to readers who are not familiar with neuroscience, especially, the hippocampus and MEC although the authors simply listed 10 papers in the actually related work section.

* The figure is generally helpful in understanding the method and the results.

* The proposed method is a very simple E-M-based algorithm but it is more efficient and has better performance than CEBRA in the grid cell simulation with the RatsInABox simulator.

weaknesses:
1. 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).

2. The authors misuse big-O notation. For example, the authors mentioned O(1 hour, 200 neurons, 10^6 spikes) ~ O(1 min) in L 120 and O(hours) in L219. This causes confusion, especially to computer scientists who have been trained in big-O notation for measuring time or space complexity. If the authors want to use big-O notations, please change them to match the function of neurons/spikes and so on. However, I believe, just changing to plain text might be better.

3. Except for Section 4.2, there are no comparisons with existing methods or ablation studies. Although the results appear promising, it remains challenging to discern whether this is due to the simplicity of the task or the efficacy of the proposed method. Moreover, I am not sure whether CEBRA is the only method to compare with.
The authors should provide more details on the training and testing splits. L305–306 and 416–417 lack information: are the splits based on spatial segmentation (using specific parts of the box for training) or trial separation?

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

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

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

7. The writing can be improved a lot:

7.1 It is possible that the readers do not know what the “turning curve” is since it is used in neuroscience literature and ICLR is generally a machine learning community. It is good to define what it is in the introduction.

7.2 It would be better to move the Related Work section to Before Method (standard ML conference styles) or Discussion (some ICML styles). Currently, it disconnects methods and results.

7.3 Please choose between British English and American English. Currently, it is used together (e.g., In L477, optimization, behaviour).

7.4 Consider that readers may not be familiar with “tuning curve,” a term more common in neuroscience literature. Since ICLR has a general machine learning audience, defining it in the introduction would be helpful.

7.5 Moving the Related Work section to a position before the Method section (per standard ML conference styles) or to the Discussion (in line with some ICML styles) may improve flow, as it currently separates the methods and results sections.

7.6 Maintain consistency in English dialect, as the paper currently alternates between British and American English (e.g., "optimization" vs. "behaviour" in line 477).

7.7 In line 131, specify “Appendix A” instead of “Appendix,” even if there is only one section.
7.8 Ensure the font for v is consistent in lines 142 and 143.

7.9 Since space permits, it may be preferable to adjust the placement of Figure 5 so that it doesn’t span pages 7 and 8, avoiding the empty space near lines 399–405.

7.10 IN L816, 1, …, P => \{ 1, …, P \}.

7.11 Correct citation formatting:
* In line 116, “Tanni et al. (2022)” should use citep.
* In lines 486, 493, and 790, change citep to citet.
* In line 522, switch citet to citep.

7.12 It would also be beneficial to unify the reference formatting:
* Journals like Nature and Nature Neuroscience do not list volume numbers, while JMLR and most eLife papers do (except in lines 583–585). Additionally, both “elife” and “Elife” are used inconsistently, as are “Nature” and “nature.”
* Clarify “cite” with a red background in L608.
* While “ICLR” is abbreviated, other conferences like NeurIPS and ICML are listed in full; consider unifying this approach.

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

* 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730685572804

thread_type:
official_review

<thread_object>
creator: reviewer_gb66

<content>
summary:
The authors propose a new latent variable model for neuroscience that uses kernel density estimation to learn tuning curves from an unobserved latent variable for population spiking data. The model is fit using a closed form EM algorithm and is scalable to large datasets. The authors validate the model on simulated and real neural data.

strengths:
The validation on synthetic and real neural data is a nice addition. The simplicity of the approach has potentially some appeal to experimentalists who want to avoid some of the more bespoke or complex latent variable models for the neural spiking data.

weaknesses:
The authors provide a very brief overview a very broad field of latent variable models used in neuroscience to motivate their approach. They claim existing methods do not scale well, or have other shortcomings compared to SIMPL. To robustly demonstrate that this model indeed scales better than most or all of these existing approaches while retaining accurate latent and neural tuning identification would require far more comparisons, discussion and evaluation. 

In particular, the authors model is motivated in a similar way to the GPLVM, which is discussed but not compared to. As the authors point out GP based models often do have issues with scalability. However,  there have been many steps toward improving the scalability of these models in recent years -- note that 1 and 2 uses inducing points, a well-known approach to improve scalability in GP  models. These models however have the same 'tuning curve' interpretations that the authors posit SIMPL has, and so it would be important to directly compare to them for both scalability and neural tuning identification. 

1 "Manifold GPLVMs for discovering non-Euclidean latent structure in neural data" 
2. Learning interpretable continuous-time models of latent stochastic dynamical systems

There is also no mention of switching linear dynamical systems models, which often scale better than GP based methods due admitting an EM based approach and utilizing forward-backward passing (see e.g. 3, 4, but there are many others). 

3 Bayesian learning and inference in recurrent switching linear dynamical systems.
4. A general recurrent state space framework for modeling neural dynamics during decision-making.

Still there are others in this space that scale well and admit flexible non-linear latent characterizations. See for example 4, 5. and of course LFADS, as the authors discuss. The authors also mention Pi-VAE (and there are of course now many other VAE based models used in neuroscience), but they don't provide a principled comparison or discussion of these approaches. 

4 Inferring Latent Dynamics Underlying Neural Population Activity via Neural Differential Equations
5  Collapsed amortized variational inference for switching nonlinear dynamical systems.

The choice of CEBRA as the only benchmark is not well motivated. CEBRA is a fundamentally different model it uses behavioral information for contrastive learning, and it was designed to visualize a behaviorally-informed latent space. In this sense it isn't an unsupervised model used purely on spiking data (like SIMPL as well as the ones cited above, among others) and it does not permit a tuning curve interpretation -- which the authors emphasize is an important component of SIMPL and existing in many of these other LVMs. Comparison to the GPLVM, as well as switching LDS, or other nonlinear dynamical models would be far more appropriate then CEBRA. 

In it's current state, it is unclear exactly how much better SIMPL scales than any of the existing approaches, and how it's latent identification or tuning identification may compare to these many state-of-the-art methods.

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1730713980545

thread_type:
official_review

<thread_object>
creator: reviewer_kcef

<content>
summary:
The paper introduces a new method they call SIMPL, an EM-style algorithm, aiming to recover low-dimensional, time-evolving latent variables from high-dimensional neural activity in large neural datasets. Their  approach fits tuning curves to observed behaviour and iteratively refines these through a two-step process (EM like). The originality of SIMPL lies in its novel combination of expectation-maximization (EM) techniques with behavior-driven initialization, providing a straightforward yet effective approach to refine neural representations. It’s main advantage is that it is a scalable and fast approach compared to existing popular methods like CEBRA.

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

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

The capability to achieve results over 30 times faster than comparable methods while maintaining accuracy makes SIMPL particularly attractive for large-scale neuroscience research

Careful experiments on synthetic as well as real data.

weaknesses:
Limited Evaluation on Non-Spatial Tasks: While the paper demonstrates impressive results in the domain of spatial navigation, it would benefit from evaluation on other tasks, such as the ones tested in the CEBRA paper so that the two methods can be directly compared in more diverse settings.

Scalability Concerns for High-Dimensional Latents: The authors mention that kernel density estimation (KDE) may not scale well to high-dimensional latent spaces. This potential limitation needs more elaboration. What alternatives might be suitable when dealing with truly high-dimensional data.

questions:
How does the method perform with smaller neural populations? Can SIMPL perform well when there are limited neurons recorded. Some downsampling analysis on neural data can be helpful to check if the method is sensitive to neuron number.

How much data is required to optimize SIMPL? How does the mother deal with short recordings? For example, would a short half hour recording be sufficient to fit the model?

The author mentioned that they trained CEBRA on the synthetic grid cell data using out-of-the-box hyperparameters training for the default 10000 iterations. 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?


The author mentioned that SIMPL finds a modified latent space with smaller, more numerous, and more
uniformly-sized place fields, suggesting the brain may encode space with greater resolution than previously thought. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731167462679

thread_type:
official_review

<thread_object>
creator: reviewer_6vn3

<content>
summary:
The paper introduces SIMPL, an EM-style algorithm for refining latent variables and tuning curves from spiking neural data. The key idea is using behavior as initialization and combining kernel density estimation with Kalman smoothing in an iterative optimization framework. The authors validate their approach on synthetic datasets and real hippocampal recordings, showing improvements in both computational efficiency and biological interpretability compared to CEBRA (a recent popular deep learning approach for learning latent embeddings from neural and behavioural data)

strengths:
- The paper is a fresh departure from current latent variable models relying on deep neural networks, with a number of practical advantages (e.g., fast optimization, minimal hyperparameters (only 2), simple implementation based on standard components (Kernel Density Estimation and Kalman smoothing))
- The paper reveals scientific insights through the use of the proposed method (finer structure in place field representations, new interpretation of place field size distribution, relationship between behavioral uncertainty and neural encoding)
- The authors do a good job presenting the approach clearly and acknowledging the limitations of the proposed method.

weaknesses:
The proposed approach relies on strong assumptions that seem to limit the application of the method for more complex cases beyond the datasets used here (e.g., settings where behavioral recordings are missing, higher dimensional latent states, latent spaces not dominated by behavior, non-smooth latent trajectories). It would be great to demonstrate the applicability of the approach on commonly used datasets in the current literature on LVMs in neuroscience (e.g., Neural Latents Benchmark).

questions:
- The authors mention improved identifiability of the proposed approach. Is that solely because of the behavioral initialization?
- 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?
- Can the authors provide the intuition for choosing the kernel bandwidth in the KDE step? This seems critical for the method's performance.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731511951492

thread_type:
public_comment

<thread_object>
creator: ~cole_lincoln_hurwitz1

<content>
title:
Missing citations for neural-behavioral modeling

comment:
There are a bunch of methods for neural-behavioral modeling which are not cited in this paper. It would be good to include these references as building latent variable models that exploit behavior is a very active field of study.

- Sani, Omid G., et al. 2021: This work proposes a linear dynamics approach for modeling neural-behavioral data.

- Hurwitz et al. 2021: This work proposes a sequential VAE for modeling neural-behavioral data.

- Gondur et al. 2024: This work proposes a multi-modal gaussian process variational autoencoder for neural-behavioral data

- Sani et al. 2024: This work proposes an RNN-based architecture for neural-behavioral data.

Rabia Gondur, Usama Bin Sikandar, Evan Schaffer, Mikio Christian Aoi, and Stephen L Keeley. Multi-modal gaussian process variational autoencoders for neural and behavioral data. In International Conference on Learning Representations, 2024.

Sani, Omid G., et al. "Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification." Nature Neuroscience 24.1 (2021): 140-149.

Cole Hurwitz, Akash Srivastava, Kai Xu, Justin Jude, Matthew Perich, Lee Miller, and Matthias Hennig. Targeted neural dynamical modeling. Advances in Neural Information Processing Systems, 34:29379–29392, 2021.

Omid G Sani, Hamidreza Abbaspourazad, Yan T Wong, Bijan Pesaran, and Maryam M Shanechi. Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification. Nature Neuroscience, 24(1):140–149, 2021.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732131182068

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of the reviews and our response

comment:
_Nb: PLEASE RE-READ this has been edited to reflect further discussion with reviewer gB66_

Based on the reviews and subsequent follow-up discussion, we have submitted a revised manuscript which has undergone significant improvements. 

**To summarise _strengths_**, the reviewers appreciated the novelty of our paper stating that it is "a fresh departure from current latent variable models relying on deep neural networks" (6vN3) as well as its computation efficiency and scalability (Kcef, 5Qja, gB66) and that "the simplicity of the approach has potentially some appeal to experimentalists who want to avoid more complex latent variable models" (gB66). They noted that it reveals new scientific insights (6vN3) and that we made "careful experiments on synthetic as well as real datasets"(Kcef, gB66) as well as "rigorous comparisons with existing methods"(Kcef) . They also commented that the paper was well presented (5Qja, 6vN3), the figures were clear (5Qja) and that we "acknowledged the limitation of the proposed method" well (6vN3)  

**To summarise _weaknesses_**, there were two primary concerns. Firstly, three of the reviewers (5Qja, 6vN3, Kcef) wanted to see SIMPL tested on additional non-spatial datasets, all pointing towards the same Macaque hand-reaching dataset used in previous LVM studies (we have now done this, consequently the reviewers have raised their scores to accept/strong accept-status). Secondly, reviewer gB66's primary concern was that we have not made sufficient, nor relevant comparisons to alternative techniques (we have discussed this in depth, added three more methods to our benchmarks, rewritten the related work and added a technique-comparison table to the appendix). The reviewers also had other more minor concerns about whether SIMPL performance depends strongly on the choice of hyperparameters (6vN3, we tested this and it did not), dataset size (Kcef, likewise), or continuity of the latent variable (6vN3, likewise). To a large extent, we are in agreement with the reviewers regarding all of these concerns. Consequently we have run new benchmarks, performed new experiments/hyperparameter sweeps, run SIMPL on the Macaque dataset and rewritten the related work section, all of which leave the manuscript in a much stronger position. Our most substantial changes are summarised in the following table: 

**Summary of most notable changes** 

| Addition | Comment | Relevant reviewers | 
|-------|-------|------|
| **Macaque dataset analysis** |SIMPL was applied to a somatosensory macaque dataset and performed well demonstrating its applicability to non-spatial and higher-dimensional latent spaces. | 6vN3, Kcef, 5Qja |
| **Additional benchmarks**| We have now benchmarked SIMPL against three additional, more relevant methods (pi-VAE, GPLVM and GPDM as well as the original CEBRA). In all cases SIMPL performs best and computes fastest. These comparisons cover the range of relevant LVM features (see extended discussion below), confirming SIMPL's superiority. | gB66, 5Qja |
| **Hyperparameter sweep** | A new figure has been adding confirming the robustness of SIMPL to its two main hyper parameters ($v$ and $\sigma$). | 6vN3 | 
| **Dataset size sweep** | We ran SIMPL on datasets of increasing smaller and smaller size (both neuron count and duration) to see how it performs in the low-data limit. | KceF | 
| **Discontinuous latent experiment** | We ran a new experiment where the latent is non-continuous and confirmed SIMPL can accurately recapitulate it (despite having an explicit smoothness prior) | 6vN3 |
| **Discussion of relevant alternatives** | Additional discussion of what models constitute relevant alternatives and why. | gB66 |
| **Rewritten Related Work** | To better explain the field of LVMs and frame SIMPL in relation to alternative techniques specifically clarifying a minimal set of unique features SIMPL and only SIMPL satisfies. | gB66 | 
| **Added missing citations** | All citations mentioned on this forum have not been added to the manuscript | gB66, C. Hurwitz (public commenter) | 

All new figures have been added to a revised manuscript in order that the reviewers can assess them. In this manuscript new text is shown in blue.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a latent variable model (LVM) for neural and behavioural representations. The model is able to learn low-dimensional latent variables from high-dimensional neural activity. The learning algorithm is based on the Expectation-Maximization framework and it is more scalable than competitive LVMs and can be applied to large datasets. 

The reviewers appreciate that this method is simple (only two hyperparameters), yet intuitive and effective. In particular, scalability gains are high, with results showing a 30x boost over comparable methods without losing accuracy. Furthermore, the reviewers find it interesting that the paper reveals scientific insights, such as the finer structure in place field representations and the relationship between behavioral uncertainty and neural encoding. 

There has been extensive discussion about adding clarifications, and the revised version has largely addressed those. With the addition of new experiments during the rebuttal period, the overall experimental evaluation section is stronger. 

However, there is one concern remaining: the comparisons are done in a simulated setting (figure 4) rather than showing the comparison in the real-world setting. This doesn’t seem to be a major concern about the convincingness of the method per se, but it is rather a concern about whether these evaluations support claims like _“when it is applied to a large rodent hippocampal dataset, SIMPL efficiently finds a modified latent space with smaller, more numerous, …”_. The wording of these claims should be easy to modify appropriately for the camera-ready version. 

Overall, I recommend the paper for acceptance.