Dynamical Modeling of Behaviorally Relevant Spatiotemporal Patterns in Neural Imaging Data

[1]: motivation for ConvRNNs

We thank the reviewer for this insightful question. There are two reasons we need to pass the ConvRNN1 states to ConvRNN2. First, ConvRNN2 aims to capture residual neural dynamics not explained by behaviorally relevant states, $X_k^{(1)}$. To ensure $X_k^{(2)}$ learns dynamics not included in $X_k^{(1)}$, we pass $X_k^{(1)}$ to ConvRNN2. Thus, this connection functions similarly to a residual connection which avoids re-learning the dynamics of ConvRNN1. Second, Sani et al. 2021 have proven that for linear dynamical systems, to get a general disentangled linear state-space model, one needs a link from $X_{k+1}^{(1)}$ to inform the recursion of $X_k^{(2)}$ (i.e. how $X_{k+1}^{(2)}$ is derived from $X_k^{(2)}$). Although SBIND uses nonlinear ConvRNNs, this linear theory provides another motivation for passing the behaviorally relevant states from ConvRNN1 to ConvRNN2.

[2]: Hyperparameter (HP) Search

Primarily, we swept the latent dimensions $n_x$ and patch size for all datasets, as these were found to be the most critical HPs for performance. Other major HPs (e.g., learning rate, sequence length, convolutional layers were fixed based on the best validation performance on the WFCI datasets. For the fUSI dataset these HPs were picked based on best behavior prediction in the validation set using one representative session and then fixed at these values across all sessions.

[3]: Applying to Electrophysiological (EP) Data

We thank the reviewer for the insightful suggestion to assess SBIND's robustness on EP data. While SBIND is primarily designed for the unique challenges of neural imaging, we agree this is a valuable comparison. We applied SBIND to an EP dataset (O'Doherty et al., 2020) containing smoothed spike counts from 42 channels and 4D behavioral kinematics (horizontal and vertical velocity and position) from two sessions of NHP reaching movements. We did not have access to the spatial coordinates of the electrodes, so we could not form an image based on true coordinates. Instead, for comparison purposes, we formed an arbitrary 6x7 2D pseudo-image from the 42 channels.

Despite this the pseudo-image input lacking true spatial structure, SBIND performed comparably to DPAD (new Table B.4), suggesting the robustness in SBIND's dynamical modeling and disentanglement approach. Finally, as the reviewer noted, we agree SBIND holds potential for other imaging modalities or modalities that naturally possess a grid structure, such as EEG recordings.

Table B.4: Comparison with DPAD on EP Recordings

---

[4]: Comments & Suggestions

We thank the reviewer for the helpful suggestions to improve clarity:

• ConvRNN naming: SBIND employs an attention-augmented ConvRNN architecture. While the self-attention (SA) module captures spatial dependencies within each time step, the architecture remains fundamentally recurrent because the state equation ensures that $X_{k+1}$ depends on $X_{k}$. Therefore, we used "ConvRNN" to describe the overall model class, with the understanding that it incorporates SA. We will clarify this in the manuscript and ensure "ConvRNN1" and "ConvRNN2" are used consistently.

• $n_y, n_x, n_z$ definitions: Please refer to Reviewer z1x2 response [1].

[5]: Dataset Size

WFCI 1 and WFCI 2 contain 49k, and 39k frames. While the fUSI dataset has ~5k frames per session, the availability of 13 sessions provides enough data overall for evaluation. To mitigate overfitting, we employed standard techniques including L2 weight decay, dropouts, and early stopping. Example loss curves are provided via this anonymous link, showing that overfitting was not an issue.

[6]: Main Contributing Components

Ablation studies (Sec 4.2; Tables 1, A.1, A.2; Fig 4) show SBIND's performance benefits stem from its ConvRNN architecture with integrated SA and use of GDL loss for more precise neural dynamical modeling. Also, learning behaviorally relevant dynamics directly from raw images leads to better behavior decoding.

[7]: $f_A$ function

The SA operates purely in the spatial domain within a single time step, k. At each time, the latent image $X_{k}$ is divided into spatial patches, and the multi-head SA computes the relationships between these patches at that specific time (not across different time steps). The temporal dependencies are handled by the outer recurrent structure of the ConvRNN, where $X_{k+1}$ is computed based on $X_{k}$. The effect of patch size (which influences the spatial scope of attention) was investigated in our ablation study (Figure A.6), showing that larger patch sizes result in better self-prediction. Also, the sequence length HP was chosen to balance sufficient temporal context with the computational limits.

References: Line 495

[1]: Comments & Suggestions

Thank you for these helpful suggestions for improving the clarity of our work.

• We will fix the typo on Line 87 and ensure all symbols are clearly defined upon first use. Specifically, $n_y$ is the number of neural images in $Y_k$ at time k, and H and W are the height and width. $n_x$ is the number of dimensions of the latent states, and $n_z$ represents the dimension of the behavior vector, $z_k$. Across all datasets, $n_y=1$, but we include $n_y$ for generality of the formulation.
• We agree with the suggestion regarding notation consistency and will amend Line 162 to clarify the RNN estimates $X_{k+1}$ given $Y_{1:k}$.
• Inference using the full sequence $Y_{1:K}$: Our model performs inference causally (prediction at time k+1 uses observations up to time k), given its importance for causal neuroscience investigations and real-time BCI applications. Using the full sequence would correspond to non-causal inference (i.e., smoothing), which was not the goal of this work but can be achieved using bidirectional RNNs in future studies.
• Finally, we have verified that the anonymous link to the code works; the previous issue may have been a temporary glitch.

---

[2]: Combined Loss

We adopted the two-stage approach because it allows us to explicitly disentangle the behaviorally relevant latent states $X^{(1)}$. The ConvRNN in Stage 1 is optimized specifically for behavior prediction, allowing the first latent state, $X^{(1)}$, to focus on capturing the behaviorally relevant dynamics. Stage 2 then focuses on modeling the residual neural dynamics, which are not predictable from $X^{(1)}$. This provides a clear separation/disentanglement of states.

While one-stage training with a mixed neural-behavioral loss is possible, doing so poses a challenge for clear disentanglement. Specifically, in this case, there is no guarantee that any state dimension will just focus on behaviorally relevant or just the residual dynamics. As such, states may actually be mixed up rather than disentangled. Doing so also may create practical challenges. First, the mixed loss optimization will heavily rely on tuning the relative weights of the neural and behavior loss terms, which can be difficult and sensitive to hyperparameters. Second, achieving disentanglement with a single combined loss may necessitate incorporating additional specialized loss terms or architectural constraints specifically designed for separation –e.g., KL divergence terms in TNDM (Hurwitz et al., 2021) that we now compare to and show that SBIND outperforms it (see Table B.2); finding the optimal balance for these specialized terms typically requires a potentially sensitive hyperparameter tuning process. In contrast, our two-phase approach avoids this sensitive hyperparameter tuning.

---

[3]: Multi-step Decoding & Inference Speed

We thank the reviewer for this important question. We will add these points to the manuscript.

Multi-step Prediction: SBIND's recurrent architecture inherently allows for multi-step prediction. We can do so with no additional optimization or retraining by performing recursive forecasting during inference as follows: instead of feeding the neural observation at the next time step into the neural encoder ($K^{(1)}$), we can feed the model's own one-step-ahead neural image prediction, $\hat{Y}_{k+1}$.

This allows the model to predict the subsequent state $X_{k+2|k}$, behavior $z_{k+2}$, and neural image $Y_{k+2}$ without having the observation at time k+1, and this same process can be iterated further into the future.

Inference Speed: A single inference step of the SBIND model takes 17.9 ms on an NVIDIA RTX 6000 GPU, which is faster than the sampling intervals of WFCI (~33-67 ms) and fUSI (~100-500 ms), suggesting potential feasibility for real-time BCI applications (Rabut et al, 2024; Mace et al, 2011).

---

[4]: Equation 2

This is another great question. Because $X_k$ and $Y_k$ have different spatial dimensions ($H' \times W'$ vs. $H \times W$), we must first encode $Y_k$ using $K(\cdot)$ to obtain a representation $K(Y_k)$ with the same spatial dimensions as $X_k$. Afterwards, there are two main options: either sum $K(Y_k)$ with the processed state $f_A(X_k)$ (as formulated in Eq. 2, or equivalently, Eq. A.1), or concatenate $K(Y_k)$ with the state $X_k$ before applying the recurrent function $f_A(\cdot)$ (as discussed in Appendix A.1.5). For the WFCI datasets reported, we utilized the concatenation approach (Eq. A.12), and it yielded better performance as seen in the new Table B.3.

Table B.3: Comparison of SBIND Variants on WFCI1 Dataset.

References:

Rabut et al., Functional ultrasound imaging of human brain activity through transparent. Sci Transl Med. 2024.

Also see Lines 502, 462.

[1]: New Baselines

We thank the reviewer for their constructive feedback. We agree that comparing SBIND with more baselines strengthens our contribution. In response, we have added new baseline comparisons.

Comparison with STNDT:

STNDT uses a Transformer architecture for spatiotemporal modeling of neural population spiking activity. Unlike SBIND, which processes raw images, STNDT was originally designed for Poisson-distributed spiking data. We adapted STNDT to accept preprocessed imaging data extracted using LocaNMF, which is a widely used approach (Wang et al, 2024). We put a Gaussian prior on LocaNMF (instead of the original Poisson prior for spikes), and trained with STNDT's original objectives (masked reconstruction and contrastive loss). Following the original work, behavior decoding was performed via ridge regression on the learned latents. We swept the number of latents for STNDT. The results (Table B.1), show that SBIND significantly outperforms STNDT on WFCI1 dataset.

Table B.1: Comparison with STNDT.

Comparison with TNDM:

TNDM (Hurwitz et al., 2021) uses a sequential variational autoencoder (SVAE) to learn two sets of latent factors in spiking data with dimensionality $n_1$ and $n_2$ for behaviorally relevant and irrelevant dynamics, respectively. TNDM was originally designed for Poisson-distributed spiking data. To provide a meaningful comparison on our WFCI dataset, we adapted TNDM to accept preprocessed imaging data extracted using LocaNMF. We assumed a Gaussian distribution for these inputs and changed TNDM's neural reconstruction loss to MSE. We also swept $n_1$ and $n_2$ values for TNDM. As shown in the newly added Table B.2, SBIND achieves superior performance compared to this adapted TNDM in both neural prediction and behavior decoding, demonstrating the benefit of SBIND's architecture that is specifically designed for spatiotemporal image data.

Table B.2: Comparison with TNDM on WFCI1 Dataset.

The reviewer also mentions two other methods, but direct empirical comparison to these was not feasible as described below:

• PoYo (Azabou et al., 2023) is a foundation model for spiking activity, which builds an unsupervised, large-scale foundation model to generalize across different subjects and tasks. This is a different goal compared to SBIND’s goal of joint neural-behavioral modeling to disentangle behaviorally relevant neural dynamics. PoYo uses a tokenization scheme specific to spikes, whereas SBIND focuses on neural imaging modalities. Furthermore, the code for PoYo is not available.

• BeNeDiff: The code for BeNeDiff is also not publicly available. BeNeDiff employs SVAEs similar to TNDM (which we now compared to). A key distinction is that SBIND processes raw imaging data, while BeNeDiff uses LocaNMF-preprocessed data. SBIND captures local and global spatiotemporal dependencies in images, which may be lost after preprocessing with methods such as LocaNMF.

---

[2]: Neural Latent Benchmark (NLB) Comparison Request

We thank the reviewer for raising this point. Unfortunately, however, all NLB datasets consist exclusively of electrophysiological (EP) recordings. SBIND, in contrast, is specifically designed to model the spatiotemporal grid structure of neural imaging data (WFCI and fUSI). Our core contribution lies in developing a method tailored to leverage this structure from raw image sequences. Therefore, while the NLB is valuable for benchmarking models designed for EP data, it is not directly applicable for evaluating the specific contributions and application of SBIND.

---

[3]: Question on Raw vs Preprocessed Data

This is another great point about modeling raw versus preprocessed neural images. Methods like LocaNMF extract latent components roughly corresponding to brain regions. The adapted STNDT we compared against uses learnable positional embeddings to encode spatial relationships within LocaNMF components. However, as shown in our comparison (Table B.2), this method, even with mechanisms for spatial encoding on preprocessed data, yields inferior performance compared to SBIND.

In terms of computational efficiency, we acknowledge that training models on preprocessed data is generally faster than training on raw images. However, as shown in Tables 2 and 3, baselines using various preprocessing methods consistently perform worse than SBIND. Therefore, SBIND represents a trade-off that prioritizes higher predictive accuracy and the ability to learn directly from the neural images.

[1]: Neuroscientific meaning

We thank the reviewer for raising this key point regarding the scientific utility of neural imaging data compared to electrophysiology (EP) data, which we will now clarify in the manuscript. Brain function relies on diverse spatial and temporal scales from single neurons to large-scale networks of neuronal populations. To understand how the brain generates behavior, we need to model neural activity across all these scales. While EP recordings measure the small-scale spiking activity of a group of single neurons, they do not measure large-scale networks that are thought to be a key basis for cognition and complex behavior (Cardin et al, 2020). Neural imaging can measure these large-scale networks by providing large-scale spatial coverage across cortical areas or even whole-brain access (Macé et al., 2018). As such, neural imaging data are complementary to EP data and play an increasingly central role in modern neuroscience by allowing the study of large-scale network dynamics, functional connectivity, and mesoscale neural analyses, which are all crucial for a full understanding of how the brain generates behavior and task performance (Musall et al., 2019; Nietz et al., 2023). Additionally, functional ultrasound imaging (fUSI) may offer a less invasive approach compared to EP for developing brain-computer interfaces (BCIs) as recently demonstrated (Griggs et al. 2024; Rabut et al., 2024).

Finally, we agree with the reviewer that neural imaging and EP data have distinct statistical characteristics. Precisely for this reason, we developed SBIND to directly learn such specific spatiotemporal structure from raw neural images, to enable extracting more precise dynamical neural information. We recognize that our initial motivation for focusing on these imaging modalities could have been more clearly articulated and appreciate the reviewer’s insight.

[2]: Neural-behavioral models

First, as explained above, neural imaging data play a critical role in understanding large-scale network dynamics that are key to complex behavior and cognition. Furthermore, we demonstrated SBIND's application on two fundamentally distinct neural imaging modalities that are increasingly employed: an optical modality based on widefield calcium imaging (WFCI), and an acoustic modality based on fUSI.

Second, while we agree that the general concept of joint neural-behavioral modeling has been explored, particularly for EP data, we emphasize that SBIND's novelty and contribution lie in doing so for neural imaging modalities by developing a novel architectural design tailored for the unique challenges of neural imaging data. These challenges, as highlighted in the introduction and also noted by the reviewer, include high dimensionality, complex spatiotemporal dependencies, and prevalence of behaviorally irrelevant dynamics. Our model differs significantly from prior joint neural-behavioral models that typically operate on lower-dimensional EP recordings that have a much smaller spatial scale than neural imaging (Sani et al., 2024), or on preprocessed time series extracted from neural imaging data (Wang et al., 2024).

Finally, we demonstrated that SBIND outperforms SOTA methods, such as DPAD and CEBRA. We now also add a new baseline, STNDT (Table B.2 in our response to Reviewer 9wsc) to our results, showing SBIND’s superior performance in both behavior decoding and neural prediction compared with all baselines. Furthermore, as Reviewer z1x2 pointed out, "methods that can be applied to fUSI are relatively sparse", and to our knowledge, SBIND is the first dynamical latent state model for fUSI, whose successful application to fUSI underscores its potential to facilitate the design of non-invasive BCIs.

[3]: References

We thank the reviewer for bringing these relevant papers to our attention and will incorporate a discussion of them into our revised manuscript. While these represent important advancements, they differ significantly from SBIND. They all address electrophysiology data and focus on distinct goals such as cross-session latent alignment (Wang et al., 2023), modeling inter-regional communication (Li et al., 2024), capturing spatial and temporal dependencies in population activity (Le et al., 2022), or recovering interpretable inter-neuron connectivity (Zhang et al., 2024). Crucially, none are designed for joint neural-behavioral modeling, and thus they do not disentangle behaviorally relevant dynamics, nor do they employ image-specific architectural priors like SBIND's ConvRNNs with integrated self-attention. SBIND's focus thus remains distinct in leveraging the spatiotemporal structure of imaging data for robust, disentangled neural-behavioral modeling.

References:

Cardin et al., Shining a Wide Light on Large-Scale Neural Dynamics, Neuron 2020

Macé et al., Whole-Brain Functional Ultrasound Imaging Reveals Brain Modules, Neuron 2018

Also see Lines 452, 524, 529, 542, 559