BaRISTA: Brain Scale Informed Spatiotemporal Representation of Human Intracranial Neural Activity

We thank the reviewer for their feedback and we reply to the comments inline below.

The reproduced results for PopT+BrainBERT underperform the original paper [1] -- specifically, 0.80 versus 0.90 on sentence onset detection, and 0.78 versus 0.93 on speech/non-speech classification.

Yes, this is a correct observation. First, we clarify that the reason for this is that we did not use the same data segments used by the authors of Brainbert/PopT [1,2] for our analysis. As such, the reproduction results would not necessarily exactly match. Specifically, those authors generated 5-second overlapping segments which they then partitioned into training/valid/test splits. Here we generated non-overlapping segments to minimize any possible information leakage between segments for a rigorous evaluation of spatial scales.

However, to address the reviewer’s concerns we now demonstrate that our pretrained PopT model can achieve the same performance reported in the original work using our evaluation framework. To do so, we used the publicly available code on GitHub to generate the downstream data splits used in [2] and evaluated our pretrained PopT model on these data splits. In this setup, we achieved 0.883 ± 0.008 on sentence onset and 0.925 ± 0.01 on speech, which are almost the same as the numbers reported in the original manuscript. We note that the small discrepancies can be due to differences in the evaluation framework, such as seeding and the channels used. In [2] the authors only used 1 finetuning seed (to the best of our understanding) whereas we use 5 to achieve a more robust estimate of model performance. Further, here we use all channels for each session rather than selecting a subset, which is done in [2] (the authors used 90 of the channels).

The above analysis shows that our pretrained PopT model is valid and we are able to use it to reproduce the original reported results.

A minor implementation difference exists in trial length -- 3 seconds in this work versus the original paper's 5-second segments.

Although this is a valid observation, we clarify that the 3-second vs 5-second distinction is not a limitation of our approach. We chose to use the 3-second interval for two reasons. First, as we noted in the previous response, we chose to generate non-overlapping data segments (distinct from the original work [1,2]). As a result, we would have had less (pre)training data had we used the same 5-second intervals as the original work. Thus we opted to use 3-second intervals. Second, we based our choice of 3-seconds on intervals used in prior works considering language processing [3-5]. For example, the technical paper associated with the Brain Treebank dataset [3] used intervals of 4 seconds for their analysis.

I note that the supplementary code for PopT [1] on OpenReview implements a lobe-based position embedding approach, which is not described in the PopT manuscript. Given the concerns raised in Major Weakness 1, the contribution of this embedding technique to model performance remains unclear.

We thank the reviewer for directing us to the supplementary code on OpenReview. First, we clarify that we have addressed weakness 1, as we have now shown reproduction of the original numbers (described above). Second, we have integrated the OpenReview implementation of region-based positional embedding into a local copy of the publicly available PopT codebase (from GitHub), and have pretrained a new PopT with region-based encoding. Using the default pretraining settings, the resulting pretrained model underperformed the channel-level encoding version of PopT that was reported in the published paper. We hypothesize that there can be several reasons for this:

1. The pretraining configurations provided by the default codebase may not be tuned for using parcel-based spatial encoding and thus may require more extensive modification and experimentation, which are beyond the scope of this work. Indeed, the original PopT paper does not report any results with parcel-based encoding, suggesting that they may have seen no improvement with the larger spatial scale.

2. It could be that the specific discriminative pretraining task used to train PopT benefits from channel-level information more than parcel-level information. For example, the model may have a harder time performing the channel-wise discrimination task if it has the parcel-level encoding rather than a coordinate-level encoding for the channel – as the individual channel’s identity may be critical for the task.

We will revise our manuscript to clarify that our results show the benefit of flexible spatial scales in masked reconstruction pretraining. The impact of spatial encoding in other pretraining approaches (e.g., discriminative tasks such as PopT's) remains an open question and is an interesting direction for future research. Overall, our model enables flexibility in choosing spatial scales for neurofoundation modeling with masked pretraining, shows the importance of considering such scales for strong model performance, and provides the capability for future work to study spatial scales in various downstream motor, cognitive, and sensory tasks.

** References:**

[1] Christopher Wang et al. BrainBERT: Self-supervised representation learning for intracranial recordings

[2] Geeling Chau et al. Population Transformer: Learning Population-level Representations of Neural Activity.

[3] Christopher Wang et al. Brain Treebank: Large-scale intracranial recordings from naturalistic language stimuli.

[4] Ariel Goldstein et al. A unified acoustic-to-speech-to-language embedding space captures the neural basis of natural language processing in everyday conversations.

[5] Ariel Goldstein et al. Shared computational principles for language processing in humans and deep language models.

We thank the reviewer for their feedback, questions, and comments. We reply inline below.

Limited evaluation

This is an important point. There are two reasons for not evaluating on those tasks. First, we clarify that the primary focus of our work was to assess the impact of different spatial scales during pretraining on downstream model performance. As such, we prioritized performing comparisons of our pretraining framework under different spatial encoding/masking configurations, rather than demonstrating our model’s utility/generalizability on a variety of downstream tasks. Second, we did not have sufficient labeled data for a per-subject pitch and volume discrimination task to enable robust model training. This is because we generated our downstream data segments differently than how the authors of PopT generated their data segments: as noted in our manuscript, we controlled our data splits to use non-overlapping segments to avoid any potential information leakage between training and test splits (especially since neural data is highly correlated over time), therefore enabling a rigorous evaluation of spatial scales. In contrast, the original segments in PopT did contain overlap. Because of our stringent no overlap requirement, combined with the quartile label generation (as used in [1,2] for the pitch/volume tasks), we did not have sufficient training samples left over to reliably train the model. Indeed, we observed chance performance across all models (PopT, Brant, our variants) for pitch/volume discrimination due to the small training set size. That is why we decided against evaluating on pitch/volume and instead focused on the speech/non-speech and sentence onset tasks.

However, to address the reviewer’s concern about limited evaluation we decided to modify our segmentation scheme to generate more labels, while still addressing the information leakage concern. To do so, we modified our scheme to allow overlapping segments but selected train/valid/test splits across time (i.e., the beginning of the recording session for training vs. the end for testing), instead of randomly shuffling the segments and then selecting splits – as is done in PopT’s segment generation. Based on Fig. 9 of the Brain Treebank technical paper [3], we decided to decode volume which showed high SNR across multiple brain regions. We present our results (average AUC across 7 test sessions, 5 seeds each) using PopT and our pretrained models from Tab. 1 (channels/channels and parcels/channels). As we can see, similar trends on the effect of spatial scales hold for this downstream task as well, thus expanding our model’s evaluation.

"high-frequency" information of sEEG signals; is your model better at capturing high-frequency information compared to Brant and PopT; mention limitation in the discussion

We thank the reviewer for their excellent suggestion of an interesting analysis. We looked into our method’s ability to capture high-frequency information and indeed observed that the majority of the spectral power in the reconstructed signal was in the low-frequency range (approximately <=25Hz on average). As per the reviewer’s recommendation we performed the following analysis: we reconstructed 1162 3-second segments across the 7 test sessions which we low-pass filtered for the low-frequency (<40Hz) range and band-pass filtered for the higher frequency (40-150Hz) range. We then computed the reconstruction error on the filtered signals and compared the performance between the low and high-frequency ranges. Below we present the results of our analysis for 3 encoding/masking pairs (channels/channels, parcels/parcels, lobes/lobes). We computed the normalized MSE (i.e., MSE normalized by the variance of the target signal), which we found was a better metric for comparing the two regimes since the high-frequency filtered signal had lower amplitude than the low-frequency filtered signal (due to the 1/f relationship noted by the reviewer). As the reviewer hypothesized, the reconstruction error for the high-frequency range is higher than for the low-frequency range.

We also compared against Brant by finetuning Brant for the channel reconstruction task, but model convergence was very slow and the performance was poor. We did not compare against PopT because we do not expect PopT to have the ability to faithfully reconstruct the channel activity. This is because PopT’s default temporal encoder (pretrained Brainbert) summarizes the input segments by computing an average temporal representation using the middle ~500ms (center 10 samples) – regardless of the input segment length. As such, we don’t believe the model has sufficient temporal information to reconstruct the channel activity for the full 3-second interval. We will revise the discussion section to explicitly make note of the limitation of high-frequency reconstruction and list it as a direction for future work. We will also include a more comprehensive supplementary section with the analysis discussed above.

justify why you did not parameter match your model with your baselines

This is also a good point. First we clarify that our primary goal in this work, as stated in the first response, was to investigate the impact of spatial scales on model performance. As such, we optimized our model to be lightweight enough to permit multiple pretrainings in a reasonable amount of time (to test the various encoding/masking combinations), while still maintaining sufficient modeling capacity to perform well on downstream tasks. This is why we did not parameter match the larger models. Indeed, the primary reason for including the comparisons against the baseline models in Tab. 1 was to demonstrate our proposed model’s & pretraining framework’s validity (by showing that we are able to achieve results comparable to prior published works when we used channel-level encoding).

The reviewer’s suggestion to train tiny versions of Brant and PopT is interesting. We do not have access to the original dataset the larger Brant model was pretrained on and thus don’t think it would be suitable to train smaller versions of Brant without this dataset. Moreover, Brant's pretraining code is not publicly released for us to pretrain a tiny version on a different dataset. However, having access to PopT's codebase (on GitHub) and the Brain Treebank dataset (used to pretrain PopT), we did try pretraining a smaller version of PopT. We identified a model configuration that allowed for ~1M (1790594) parameters by changing the hidden dimension from the default 512 to 64. However, pretraining using the default pretraining configurations provided by the authors was unsuccessful; we evaluated downstream performance on one task for a single test session and found the performance was close to chance. We also tried a ~3M (3772930) parameter model (hidden dimension 128), but pretraining was again unsuccessful with their default settings. Pretraining a tiny PopT may require more substantial changes to the pretraining configuration and/or the model configuration itself (i.e., number of layers, etc.) than is within the scope of this work. Finally, the larger Brant and PopT models are likely better than their tiny versions on the downstream tasks, as is generally the case with foundation models. Thus, we used the best possible performance from our baselines as a comparison point for our model.

explain why you chose to not show Table 7 of the supplement in place of Table 2 in main; consider replacing Table 2 with Table 7

We thank the reviewer for their question/comment. The primary reason for showing the best pretraining seed in Tables 1 and 2 was to maintain consistency with common practice and for fairer comparison with our baselines. In Table 1, we chose a single model to fairly compare against our baseline models which reflected just one pretraining seed (i.e., a single model); we used validation set decoding to select the model. Then, for the sake of consistency, we used the same criteria to select the best performing model within each category for Table 2. Finally, for completeness we presented the most robust estimate of our model’s performance for each encoding/masking combination in Table 7. Thus, we ended up choosing the simplified Table 2 to present in the main manuscript largely for ease of exposition/explanation and to maintain consistency with the criteria used for Table 1. Our results and conclusions are consistent in both Tables 2 and 7. We will additionally provide Table 7 in the main text of our manuscript following the reviewer’s suggestion.

Replace $\mathcal{R}$ with $\mathbb{R}$; include captions of tables (or at least of table 1); typo in lines 10-11 of supplement

We thank the reviewer for their recommendations and for catching the typo. We will make the suggested changes to improve readability of the text and to be more aligned with standard notation.

Supplement E and F are really strong; move them to the main text

We thank the reviewer for describing our results to be strong and for the suggestion. Upon manuscript acceptance we will move supplements E and F into the main manuscript.

References

[1] C. Wang et al. BrainBERT: Self-supervised representation learning for intracranial recordings.

[2] G. Chau et al. Population Transformer: Learning Population-level Representations of Neural Activity.

[3] C. Wang et al. Brain Treebank: Large-scale intracranial recordings from naturalistic language stimuli.

We thank the reviewer for their recommendations, questions, and comments. We address these below inline.

Do spatial scales vary given the classification task/experimental stimuli? E.g. smaller spatial scale needed to encode visual features in environments vs larger spatial scale for more cognitively demanding tasks such as the language tasks here? …it would be interesting to test if the same performance gains could be achieved using a second dataset in which people engage in a far less cognitively demanding task (such as viewing images of scenes). I.e. Are the benefits of varying spatial encoding less advantageous in situations where shorter, less flexible temporal and spatial scales are needed to detect changes in one's environment?

This is a great question and will be an interesting hypothesis to test in follow-on work. Based on the analyses performed here (i.e., the language-related decoding task and the channel reconstruction task), we hypothesize that the reviewer’s expectation may indeed be the case and that different applications will benefit from different spatial encodings – which can be assessed by incorporating and testing each scale using our model and its flexible spatial encoding capability. For example, we observed that models pretrained using smaller encoding spatial scales were able to perform channel-level neural reconstruction more accurately than models pretrained with larger encoding spatial scales. On the other hand, larger spatial scales helped more with language-related classification tasks.

To address the reviewer’s question, we decided to try decoding one of the visual features (global optical flow magnitude) provided by the Brain Treebank dataset [1]. We used the two encoding/masking configurations highlighted in Table 1 (i.e., parcel/channels and channels/channels), and found that the parcel/channels model outperformed the channel/channel model. Our results are presented in the table below (average AUC across the 7 test sessions and 5 finetuning seeds each); as baselines we compared against PopT and randomly initialized BaRISTA models with channel-level and parcel-level encoding.

These results indicate that the flexibility of spatial encoding offered by BaRISTA is beneficial for decoding visual stimuli as well. However, it will be very interesting to use the flexibility of spatial encoding enabled by our model for hypothesis-driven testing of spatial encoding scales on a variety of downstream tasks that exhibit different degrees of complexity in the future – including simpler sensory tasks. We will include discussion of this point as a possible future research direction in our manuscript.

Please clarify earlier on in your intro the working definition of neural tokenization.

We thank the reviewer for their recommendation. We will revise the manuscript accordingly.

References:

[1] Christopher Wang et al. Brain Treebank: Large-scale intracranial recordings from naturalistic language stimuli.

We thank the reviewer for their thorough comments, recommendations, and excellent questions. We reply inline below.

retrain or fine-tune baselines under different scales

This is an interesting point. First, we clarify that there was no notion of spatial encoding based on metadata (e.g., channel coordinates) in the iteration of Brant against which we compare [1]; vanilla index-based positional encoding [2] was used. Introducing spatial encoding to Brant would be adding new functionality and would require validation against the original pretrained model. The fairest way to perform such a comparison and validate the new functionality would be to pretrain on the same dataset the existing model was pretrained on. However we did not have access to that dataset nor did we have access to their pretraining code. Thus retraining/finetuning/validating Brant with different spatial scales would be outside of the scope of this work. The request to pretrain PopT with different spatial scales is possible, as we have the dataset the model was originally pretrained on and their code. We tried pretraining PopT with region-level encoding by replacing the $(x,y,z)$ encoding used in the original work with parcel-level spatial encoding. Using the default pretraining settings available on PopT's GitHub, we were unable to successfully pretrain the model: performance on downstream tasks was chance (~0.5 AUC on average). We suspect that this result may be due to the difference in pretraining tasks (we use a masked reconstruction task while PopT was trained on discriminative tasks). In order to train PopT to support different spatial scales, the pretraining task and/or optimization configurations would most likely need to be modified. Considering this result, we will revise the manuscript to clarify that the impact of spatial scales was seen for masked reconstruction pretraining, and that more work needs to be done to investigate the impact of spatial encoding in other pretraining tasks (e.g., discriminative tasks).

complete Tab 1 & Tab 3 of the main paper

We clarify that Tab 2 is the complete version of Tab 1. We chose to present the comparisons against baseline models separately from the full 3x3 scale combinations in Tab 2 for the sake of readability; we will revise the text to make the relationship between Tab 1 and 2 clearer. Below we present the results of completed Tab 3 which we will add to the manuscript. We will revise section 4.4 to reflect these more comprehensive results, specifically noting that the granularity of spatial encoding may impact the model’s ability to learn individual channel properties, as evidenced by the drop in channel reconstruction performance with the lobe-level spatial encoding model.

Motivate & ablating the design choices of Sect. 3.2; explain additional benefits

The choice to use a dilated CNN was based on prior works modeling uni/multivariate time-series activity [3-5]. We have now also performed an ablation on the temporal encoder choice and present results comparing against a linear projection (linear layer the size of our patch length $L$, similar to [1]) and a single layer univariate CNN with kernel size 5. For all models we used the best-performing encoding/masking pair presented in Tab 1 (parcels/channels) and report average AUC over 3 pretraining & 5 finetuning seeds (as in Tab 7). We will complete a more comprehensive ablation (use all spatial encoding/masking scales) and add it to the supplementary materials. Our results show that dilated CNN encoder achieves higher performance. With respect to computational savings: a dilated CNN is slower to train than a fully connected linear model but requires fewer parameters (302K for the linear projection vs 71.8K for the dilated CNN in our implementation). An RNN or SSM would have also been suitable alternatives but the training time for RNNs/SSMs is longer than that of a CNN. Since we wanted to pretrain several encoding/masking configurations in our analysis, we chose to optimize running time for the sake of experimentation and therefore opted for a CNN over a RNN/SSM. We will include this discussion on to computation savings/design choices in the manuscript.

We decided to use interleaved tokens (i.e., $\mathbf{S}$ vector) to give our model the capacity to learn spatiotemporal relationships between the channels within the same attention module (rather than each axis individually). We now show that this interleaved approach outperforms having separated attention modules through an ablation study in which we first pass our sequences through a temporal attention module (i.e., self-attention on the patches within each channel independently) and then through a spatial attention module (i.e., self-attention on the channels within each patch), similar to [1,6,7]. For fairness of comparison, we split our 12-layer transformer into 2 6-layer transformers each with 4 attention heads and the same hidden dimension 64. We used the best performing encoding/masking pair in Tab 1, but will add a more comprehensive analysis with all pairs in the supp materials. AUC scores are averaged across 3 pretraining & 5 finetuning seeds for each of the 7 test sessions. Our results show that the combined attention module achieves higher downstream performance.

We use a lightweight linear decoder to evaluate the quality of the learned embeddings/representations, as is common practice [8,9]. We were concerned that a more complex downstream decoder would have boosted the downstream performance regardless of the embedding/representation quality. We will clarify our choice of decoder in the text.

counts of training steps taken

We thank the reviewer for the recommendation. For pretraining our model had 19,500 update steps, PopT pretraining involved 500,000 update steps, and Brant had 750,000 update steps. Because we ran training for a fixed number of epochs (as reported in Supp D.2), the total number of finetuning update steps was dependent on the downstream task and session (i.e., the number of training segments available). The average number of updates for BaRISTA across 7 test sessions and 2 downstream tasks was 252 (range 393 steps), for PopT it was 629 (range 982), and for Brant it was 1258 (range 1964). We chose the larger number of update steps for the baseline models to ensure they converged because we wanted to validate our model against their best performance. We trained Brant the longest as its finetuning learning rate was 1e3 times smaller than BaRISTA’s and 1e2 times smaller than that of PopT’s; we note that the LRs used reflect the rates used by the authors in the original works. We will include these numbers in the revised manuscript.

No concrete reasons are provided why such clusterings could not be learned directly…question of where on the spectrum of possible clusterings lies the optimal and how to find it remains open

We agree with the reviewer that future work can explore using clusterings that are learned in a data-driven way, e.g., based on functional relationships or signal correlations. However, the use of handcrafted features in this work does not necessarily reflect a limitation in our model. We designed our model (1) to support both hand-crafted and data-driven (with minor code modifications) channel clusterings, and (2) to enable hypothesis testing by allowing users the flexibility to specify specific clustering labels based on side meta-information, which may not come from the data at hand but instead reflect domain knowledge (e.g., brain anatomy) and/or information from other modalities (e.g., structural images). For example, by utilizing clusters based on brain regions or functional relationships, users of our model can assess the significance and/or involvement of different regions or functional brain networks on various downstream tasks of interest. We will further clarify these points in the Discussion.

References

[1] D. Zhang, et al. Brant: Foundation Model for Intracranial Neural Signal

[2] A. Vaswani et al. Attention Is All You Need

[3] A. Oord, et al.. WaveNet: A Generative Model for Raw Audio

[4] Z. Yue, et al. TS2Vec: Towards Universal Representation of Time Series

[5] S. Bai, et al. An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling

[6] G. Mentzelopoulos et al. Neural decoding from stereotactic EEG: accounting for electrode variability across subjects

[7] G. Chau et al. Population Transformer: Learning Population-level Representations of Neural Activity

[8] F. Pei et al. Neural Latents Benchmark '21: Evaluating latent variable models of neural population activity

[9] T. Chen et al. A Simple Framework for Contrastive Learning of Visual Representations