Population Transformer: Learning Population-level Representations of Neural Activity

## Abstract

Abstract We present a self-supervised framework that learns population-level codes for arbitrary ensembles of neural recordings at scal. We address two key challenges in scaling models with neural time-series data: sparse and variable electrode distribution across subjects and datasets. The Population Transformer (PopT) stacks on top of pretrained representations and enhances downstream decoding by enabling learned aggregation of multiple spatially-sparse data channels. The pretrained PopT lowers the amount of data required for downstream decoding experiments, while increasing accuracy, even on held-out subjects and tasks. Compared to end-to-end methods, this approach is computationally lightweight and more interpretable, while still retaining competitive performance. We further show how our framework is generalizable to multiple time-series embeddings and neural data modalities. Beyond decoding, we interpret the pretrained PopT and fine-tuned models to show how they can be used to extract neuroscience insights from massive amounts of data. We release our code as well as a pretrained PopT to enable off-the-shelf improvements in multi-channel intracranial data decoding and interpretability.

## 1 Introduction

Building effective representations of neural data is an important tool in enabling neuroscience research.
We are particularly interested in modeling neural activity recorded from spatially varying electrode layouts, such as intracranial recordings, which use invasive probes placed within the brain to provide high temporal resolution recordings of local neural activity (Parvizi & Kastner, ; Herff et al., ).
But because of dispersed electrode placement within the brain volume, intracranial recordings (iEEG) suffer from data sparsity, and there is significant variability in probe placement across subjects (Parvizi & Kastner, ; Herff et al., ), leading to high variability in the semantics of the input channels.
While these difficulties are particularly relevant for iEEG, they also exist for other neural data modalities such as non-invasive EEG due to variation in electrode layouts.
Furthermore, annotated training data is often scarce, since it is limited by patient availability and labor-intensive labeling.
This is a problem for many existing approaches which rely on supervised learning (Faezi et al., ; Herff et al., ; Martin et al., ; Metzger et al., ; Willett et al., ).

To improve decoding data-efficiency, self-supervised pretraining on unannotated data can be employed to first learn generic representations of the recordings, independent of the downstream decoding task that requires task labels.
This means that the model does not have to use valuable annotated samples to learn how to do feature extraction before it can do classification, improving the reach of neuroscientific research.
In this paper, we are interested in developing generic representations of multi-channel intracranial recordings that enable efficient adaptation to a wide range of downstream decoding tasks.
Prior work has shown how to pretrain subject-specific (Le & Shlizerman, ) or channel-specific (Wang et al., ) models of intracranial data, but such techniques ignore inter-channel relationships or commonalities that might exist across subjects.
The most general approach would be to pretrain using data from multiple datasets, but would require tackling the aforementioned challenges of sparse electrode coverage and variable electrode placement between subjects.

We propose Population Transformer (PopT), a self-supervised pretraining approach that learns subject-generic representations of arbitrary electrode ensembles.
Transformers offer the flexibility to learn aggregate information across channel configurations, but massive data is needed to train the attention weights (Devlin et al., ).
During pretraining, we train on large amounts of unannotated data and simultaneously optimize both a channel-level and ensemble-level objective. This requires the model to (1) build subject-generic representations of channel ensembles and (2) meaningfully distinguish temporal relationships between different ensembles of channels.

Our PopT approach is modular, and builds on top of powerful single-channel temporal embeddings, which provides two key advantages.
First, by separating the single-channel embedding and multi-channel-aggregation into different modules, we make our approach agnostic to the specific type of temporal embedding used, leaving room for future independent improvements along either the temporal or spatial dimension (an approach that has been validated in video modeling (Arnab et al., )).
Second, by taking advantage of learned channel embeddings, PopT training is computationally lightweight compared to their end-to-end counterparts (Appendix B) and baseline aggregation approaches (Figure 4), allowing for adoption in lower compute resource environments.

Empirically, we find that our pretrained PopT outperforms commonly used aggregation approaches (Ghosal & Abbasi-Asl, ), and is competitive with end-to-end trained methods (Zhang et al., ; Yang et al., ; You et al., ).
Moreover, we find that these benefits hold even for subjects not seen during pretraining, indicating its usefulness for new subject decoding.
We also show that the pretrained PopT weights themselves reveal interpretable patterns for neuroscientific study.
Finally, we demonstrate that our proposed framework is agnostic to the underlying temporal encoder further allowing it to adapt to other neural recording modalities.

Our main contributions are:

- 1.
a generic self-supervised learning framework, Population Transformer (PopT), that learns joint representations of arbitrary channel ensembles across neural datasets,
- 2.
a demonstration that pretraining systematically improves ensemble representations for downstream decoding even for held-out subjects,
- 3.
a new method for brain region connectivity analysis and functional brain region identification based on the pretrained and fine-tuned PopT weights,
- 4.
a trained and usable off-the-shelf model that computes population-level representations of high temporal resolution intracranial neural recordings.

## 2 Related Work

Self-supervised learning on neural data Channel independent pretrained models are a popular approach for neural spiking data (Liu et al., ), intracranial brain data (Wang et al., ; Talukder & Gkioxari, ), and general time-series (Talukder et al., ). Additionally, in fixed-channel neural datasets, approaches exist for EEG (Chien et al., ; Kostas et al., ; Yi et al., ), fMRI (Thomas et al., ; Kan et al., ; Ortega Caro et al., ), and calcium imaging (Antoniades et al., ) datasets.
However, these approaches do not learn population-level interactions across datasets with different recording layouts, either due to a single-channel focus or the assumption that the channel layout is fixed.
Several works pretrain spatial and temporal dimensions across datasets with variable inputs (Zhang et al., ; Yang et al., ; Jiang et al., ; Ye et al., ; Cai et al., ), but most simultaneously learn the temporal embeddings with the spatial modeling, which make them challenging to interpret and computationally expensive to train, especially for high temporal resolution signals. To our knowledge, we are the first to study the problem of building pretrained channel aggregation models on top of pre-existing temporal embeddings trained across neural datasets with variable channel layouts, allowing for modeling of high quality neural data.

Modeling across variable input channels Modeling spatial representations on top of temporal embeddings has been found to be beneficial for decoding (Faezi et al., ; Le & Shlizerman, ; Azabou et al., ), but prior works use supervised labels, so do not leverage large amounts of unannotated data. The brain-computer-interface field has studied how to align latent spaces (Pandarinath et al., ; Karpowicz et al., ; Degenhart et al., ; ; Ma et al., ) which either still requires creating an alignment matrix to learn across datasets or only provides post-training alignment mechanisms rather than learning across datasets. Other approaches impute missing channels or learn latent spaces robust to missing channels (Talukder et al., ; Zhang et al., ; Chau et al., ), but these are more suited for the occasional missing channel rather than largely varying sensor layouts. We directly learn spatial-level representations using self-supervised learning across datasets to leverage massive amounts of unannotated intracranial data.

Figure: Figure 1: Schematic of our approach. The inputs to our model (a) are the combined neural activities from a collection of intracranial electrodes in a given time interval. These are passed to a frozen temporal embedding model (dotted red outline: BrainBERT shown), which produces a set of time-contextual embedding vectors (yellow). The 3D position of each electrode (red) is added to these vectors to produce the model inputs (orange). The PopT produces space-contextual embeddings for each electrode and a [CLS] token (blue), which can be fine-tuned for downstream tasks. During pretraining, the PopT is trained on two objectives simultaneously. In the first, (b) the PopT determines whether two different sets of electrodes (orange vs brown) represent consecutive or non-consecutive times. In the second objective, (c) the PopT must determine whether an input channel has been replaced with activity at a random other time that is inconsistent with the majority of inputs.
Refer to caption: extracted/5912780/figures/approach_v2.png

## 3 Population Transformer Approach

Figure 1 overviews our Population Transformer (PopT) approach. The key ideas are: (1) to learn a generic representation of neural recordings that can handle arbitrary electrode configurations; and (2) to employ a modular system design that uses a transformer architecture to aggregate information from existing per-channel temporal embeddings. To do so, we employ a self-supervised pretraining approach to learn ensemble and channel level representations. Afterwards, one can fine-tune PopT on downstream decoding tasks. In addition to offering strong decoding results, including generalization to new subjects with different electrode configurations than training subjects (see Section 5), the modular system design is computationally lightweight (see Appendix B), can benefit from improved temporal representations, and is more readily interpretable (see Section 6).

Architecture
A schematic of our Population Transformer (PopT) approach is shown in Figure 1.
We adapt a transformer backbone due to its ability to accommodate variable channel configurations.
Consider a given subject with $N$ channels indexed by $C=\{1,...,N_{c}\}$, and an arbitrary subset of channels $S\subseteq C$.
Activity from channel $i$ at time $t$ can be denoted by $x^{t}_{i}$.
The PopT takes as input an interval of brain activity $X^{t}=\{x^{t}_{i}|i\in S\}$ from a given time $t$ and a special [CLS] token.
Per channel, each interval of brain activity is passed through a temporal embedding model $B$, in the figure’s case BrainBERT, to obtain a representation of each channel’s temporal context.

To allow the model to learn a common brain state representation across layouts, each channel’s embedding is augmented with its 3D position, so that the final input to the PopT is $X^{t}_{B}=\{B(x^{t}_{i})+pos(i)+\mathcal{N}(0,\sigma)|x^{t}_{i}\in X^{t}\}$.
Spatial location is given by the electrode’s Left, Posterior, and Inferior coordinates for iEEG electrodes (Wideman, ), and XYZ positions for EEG electrodes.
We add Gaussian fuzzing to each coordinate location to prevent overfitting to a particular set of coordinates.
Each coordinate is then encoded using sinusoidal position encoding (Vaswani et al., ).
Membership in a particular ensemble (see below: ensemble-wise loss) is also encoded.
The four encodings are concatenated together to form the position embedding $pos(i)=[e_{\text{left}};e_{\text{post.}};e_{\text{inf}};e_{\text{ensemble}}]$.

The core of PopT consists of a transformer encoder stack (see Appendix A: Architectures).
The output of the PopT are spatial-contextual embeddings of the channels $Y=\{y_{i}\}$ as well as an embedding of the CLS token $y_{cls}$.
During pretraining, the PopulationTransformer additionally is equipped with a linear layer head for the [CLS] token output and separate linear layer heads for all other individual token outputs.
These produce the scalars $\tilde{y}_{cls}$ and $\tilde{y}_{i}$ respectively, which are used in the pretraining objective (see Figure 1b and c).

Self-supervised loss Our loss function has two discriminative components: (1) ensemble-wise — the model determines if activities from two channel ensembles occurred consecutively, requiring an effective brain state representation at the ensemble-level, (2) channel-wise — the model identifies outlier channels that have been swapped with a different timepoint’s activity, requiring sensitivity to surrounding channel context.

A key aspect of our method is the fact that our objective is discriminative, rather than reconstructive, as is often the case in self-supervision (Liu et al., ; Wang et al., ).
We found this to be necessary, because in practice, the temporal embeddings often have low effective dimension (see Wang et al. ()), and reconstruction rewards the model for overfitting to “filler” dimensions in the feature vector (see Section 5).

Pretraining
In ensemble-wise discrimination (fig. 1b), two different subsets of channels $S_{A},S_{B}\subset C$ are chosen with the condition that they be disjoint $S_{A}\cap S_{B}=\emptyset$.
During pretraining, the model receives the activities from these channels at separate times $X_{A}^{t}=\{x^{t}_{i}\mid i\in S_{A}\}$ and $X_{B}^{t^{\prime}}=\{x^{t^{\prime}}_{i}\mid i\in S_{B}\}$.
The objective of the task is then to determine whether these states $X_{A}^{t}$ and $X_{B}^{t^{\prime}}$ have occurred consecutively in time or are separated by some further, randomly selected interval.
Given the output of the classification head, the loss function $\mathcal{L}_{N}$ is again the binary cross entropy.
We also vary the number of input channels during sampling to ensure the model handles ensembles of different sizes.

Additionally, we select disjoint subsets for ensemble-wise discrimination to prevent the model from solving tasks through trivial copying.

In channel-wise discrimination (fig. 1c), the model must determine whether a channel’s activity has been swapped with activity from a random time.
Precisely, activity from each channel $i$ is drawn from a time $t_{i}$.
All channels are drawn from the same time $t_{i}=T$, and then 10% of the channels are randomly selected to have their activity replaced with activity from the same channel, but taken from a random point in time $t_{i}\neq T$.
Then, given the token outputs of PopT, the channel-wise loss function $\mathcal{L}_{C}$ is the binary cross entropy.
Then, our complete objective function is $\mathcal{L}=\mathcal{L}_{N}+\mathcal{L}_{C}$.

Fine-tuning
During fine-tuning, the [CLS] intermediate representation, $\tilde{y}_{cls}$ of the pretrained PopT is passed through a single layer linear neural network to produce a scalar $\hat{y}_{cls}$; this forms the input to the binary cross entropy loss for our decoding tasks (see Section 4).

## 4 Experiment Setup

Data We use two types of neural time-series data: intracranial and scalp electroencepholography (iEEG and EEG).
Whereas iEEG probes are surgically implanted within the 3D brain volume, EEG electrodes lie on the scalp.
These cover two resolution extremes of neural time-series data modalities.

iEEG: We use the publicly available subject data from Wang et al. ().
Data was collected from 10 subjects (total 1,688 electrodes, with a mean of 167 electrodes per subject) who watched 26 movies while intracranial probes recorded their brain activity.
Of the sessions, 19 are used for pretraining, and 7 of the sessions are held-out for evaluation.

EEG: We use the Temple University Hospital EEG and Abnormal datasets (TUEG and TUAB (Obeid & Picone, )) for pretraining and task data respectively. We remove all task subjects from the pretraining set and follow the data preprocessing practices in Yang et al. (); Jiang et al. ().

Decoding Tasks
We evaluate on 5 different classification tasks: 4 auditory-linguistic tasks used in the evaluation of Wang et al. () and 1 widely evaluated seizure detection task from Obeid & Picone ().
Of the auditory-linguistic tasks, two of the tasks are audio focused: determining whether a word is spoken with a high or low pitch and determining whether a word is spoken loudly or softly.
And two of the tasks have a more linguistic focus: determining whether the beginning of a sentence is occurring or determining whether any speech at all is occurring.
The EEG TUAB seizure detection task is a binary classification of pathological or normal EEG recording.

To test scaling with arbitrary ensemble sizes, we select subsets of electrodes based on their individual linear decodability, with the smallest subsets containing the electrodes with highest decodability.

Baselines
For controlled baselines, we concatenate the single-channel temporal embeddings and train a linear (Linear) or non-linear (Deep NN) aggregator on the decoding task. These enable us to directly assess how much PopT improves upon existing aggregation approaches (Ghosal & Abbasi-Asl, ).
These approaches cannot be pretrained across subjects due to the changing meaning and quantity of inputs.
To test the effectiveness of pretraining, we also compare against a non-pretrained PopT.
Methods compared
For the iEEG experiments, we also compare against Brant (Zhang et al., ), which is an end-to-end iEEG trained encoder. We take the fully pretrained Brant model, and fine-tune on our iEEG tasks combining channels with linear aggregation.
For the EEG experiments, we compare against BIOT (Yang et al., ) and LaBraM (Jiang et al., ). Both BIOT and LaBraM use discrete channel ID as positional encodings, which only allows them to be used for layouts with discete channel locations such as EEG. They also train end-to-end in contrast to our modular approach.

**Table 1: Pretraining PopT is critical to downstream decoding performance (iEEG data). We test on a variety of audio-linguistic decoding tasks (see Section 4) with 90 channels as input. The temporal encoder used for aggregation in sections 1 and 2 are denoted in the section header. We also evaluate against an end-to-end pretrained iEEG model in section 3. Shown are the ROC-AUC mean and standard error across subjects. Best per section are bolded.**
| Model | Pitch | Volume | Sent. Onset | Speech/Non-speech |
| --- | --- | --- | --- | --- |
| BrainBERT: |  |  |  |  |
| Linear Agg. | $0.59\pm 0.08$ | $0.66\pm 0.08$ | $0.70\pm 0.09$ | $0.71\pm 0.11$ |
| Deep NN Agg. | $0.58\pm 0.08$ | $0.67\pm 0.08$ | $0.71\pm 0.10$ | $0.72\pm 0.10$ |
| Non-pretrained PopT | $0.53\pm 0.06$ | $0.61\pm 0.13$ | $0.74\pm 0.10$ | $0.70\pm 0.08$ |
| Pretrained PopT | $\mathbf{0.69}\pm\mathbf{0.07}$ | $\mathbf{0.84}\pm\mathbf{0.06}$ | $\mathbf{0.86}\pm\mathbf{0.05}$ | $\mathbf{0.89}\pm\mathbf{0.07}$ |
| TOTEM: |  |  |  |  |
| Linear Agg. | $0.55\pm 0.02$ | $0.66\pm 0.03$ | $0.79\pm 0.04$ | $0.77\pm 0.05$ |
| Deep NN Agg. | $0.57\pm 0.02$ | $0.67\pm 0.03$ | $0.78\pm 0.03$ | $0.75\pm 0.05$ |
| Non-pretrained PopT | $0.53\pm 0.02$ | $0.64\pm 0.02$ | $0.79\pm 0.03$ | $0.77\pm 0.05$ |
| Pretrained PopT | $\mathbf{0.60}\pm\mathbf{0.02}$ | $\mathbf{0.73}\pm\mathbf{0.02}$ | $\mathbf{0.86}\pm\mathbf{0.03}$ | $\mathbf{0.84}\pm\mathbf{0.06}$ |
| End-to-end: |  |  |  |  |
| Brant (Zhang et al., ) | $\mathbf{0.61}\pm\mathbf{0.03}$ | $\mathbf{0.74}\pm\mathbf{0.03}$ | $\mathbf{0.80}\pm\mathbf{0.04}$ | $\mathbf{0.80}\pm\mathbf{0.03}$ |

**Table 2: Pretraining PopT is critical to downstream decoding performance (EEG data). We test on a seizure detection task (TUAB in Obeid & Picone ()) with 21 channels as input. The temporal encoder used for aggregation in sections 1 and 2 are denoted in the section header. We also evaluate against end-to-end pretrained EEG models in section 3 (values from the original works). Shown are the ROC-AUC mean and stdev across 5 random seeds. Best per section are bolded.**
| Model | Balanced Accuracy | ROC AUC |
| --- | --- | --- |
| Chronos |  |  |
| Linear Agg. | $0.7754\pm 0.0008$ | $0.8563\pm 0.0003$ |
| Deep NN Agg. | $0.7881\pm 0.0057$ | $0.8678\pm 0.0049$ |
| Non-pretrained PopT | $0.7763\pm 0.0047$ | $0.8631\pm 0.0016$ |
| Pretrained PopT | $\mathbf{0.7976}\pm\mathbf{0.0022}$ | $\mathbf{0.8821}\pm\mathbf{0.0016}$ |
| TS2Vec |  |  |
| Linear Agg. | $0.7649\pm 0.0005$ | $0.8533\pm 0.0003$ |
| Deep NN Agg. | $0.7853\pm 0.0021$ | $0.8721\pm 0.0015$ |
| Non-pretrained PopT | $0.7896\pm 0.0037$ | $0.8782\pm 0.0018$ |
| Pretrained PopT | $\mathbf{0.8063}\pm\mathbf{0.0010}$ | $\mathbf{0.8907}\pm\mathbf{0.0019}$ |
| End-to-end: |  |  |
| BIOT (Yang et al., ) | $0.7959\pm 0.0057$ | $0.8815\pm 0.0043$ |
| LaBraM (Jiang et al., ) | $\mathbf{0.8258}\pm\mathbf{0.0011}$ | $\mathbf{0.9162}\pm\mathbf{0.0016}$ |

Figure: Figure 2: Compared to common aggregation approaches, pretrained PopT consistently yields better downstream decoding across tasks, data modalities, and temporal embedding types. NPopT = Non-pretrained PopT. (a) performance on four audio-linguistic iEEG tasks with 90 electrodes. Grey bars denote standard error across subjects. (b) performance on a seizure detection EEG task with 21 electrodes. Grey bars denote standard deviation across 5 random seeds.
Refer to caption: extracted/5912780/figures/cross_dataset_performance.png

Figure: Figure 3: Pretrained PopT downstream performance scales better with ensemble size. Increasing channel ensemble size from 1 to 50 (x-axis), we see pretrained PopT (green) decoding performance (y-axis) not only beat non-pretrained approaches (orange, purple, grey), but also continually improve more with increasing channel count. Shaded bands show the standard error across subjects.
Refer to caption: extracted/5912780/figures/fig2_e50_square.png

## 5 Results

Decoding performance
We find that using a pretrained PopT significantly benefits downstream decoding compared to baseline channel aggregation techniques across tasks, data modalities, and temporal encoding models (Tables 1, 2 and 2).
The pretrained PopT scales well with increasing ensemble sizes (Figure 3a), a challenging task for the baseline aggregation approaches due to limited downstream task data and increasing input size.
To test our method’s ability to handle multiple types of channel encodings, we applied our framework to 4 different channel encoders: (1) an iEEG-specific temporal encoder (BrainBERT (Wang et al., )), (2) a general tokenization-based time-series encoder (TOTEM (Talukder et al., )), (3) a pretrained general time-series encoder (Chronos (Ansari et al., )), and a general convolution-based time-series encoder (TS2Vec (Yue et al., )).
We see significant improvements in performance with the pretrained PopT in all cases when comparing with baseline aggregation approaches (Figure 2).

We also find that PopT can achieve competitive performance against pretrained end-to-end models, such as Brant (Zhang et al., ) for iEEG, and BIOT (Yang et al., ) and LaBraM (Jiang et al., ) for EEG (Tables 1 and 2).
For instance, PopT outperforms Brant (Zhang et al., ) in decoding iEEG data with our pretrained PopT + BrainBERT combination, likely due to PopT’s ability to leverage spatial relationships.
Whereas Brant leaves the channel aggregation problem open.
PopT is competitive with recent end-to-end trained EEG models (Yang et al., ; Jiang et al., ) at the EEG TUAB seizure detection task.
This is impressive, since models such as LaBraM were specifically developed for this application, whereas PopT was trained on top of generic time-series embeddings.
We find that PopT can offer an efficient and competitive alternative to large end-to-end models for these decoding tasks, due to the effectiveness of our pretraining task for learning spatial and functional relationships between channel input embeddings.

To verify that the weights of the pretrained PopT capture neural processing well even without fine-tuning, we also train a linear-encoder on top of the frozen PopT [CLS] token and find the same trends (Figure 11.
This point in particular is important in building confidence in the results of our interpretability studies (see Section 6), in which we use the frozen pretrained weights to analyze connectivity.
Finally, for the remaining analyses described below, we use a PopT with BrainBERT inputs.

Sample and compute efficiency
Our PopT learns spatial relationships between channels, in a way that makes downstream supervised learning more data and compute efficient (see Figure 4 and Figure 5).
Compared to the non-pretrained baseline models, fine-tuning the pretrained PopT can achieve the same decoding performance as other aggregation techniques with an order of magnitude fewer samples.
The pretrained PopT surpasses the performance achieved by all other aggregation techniques by 500 samples out of the full dataset (roughly 5-10k examples depending on subject and task) Figure 4.
The pretrained PopT also converges at a low number of steps, competitive with Linear and Deep NN compute efficiency.
This greatly contrasts with the non-pretrained PopT and occasionally the Linear and Deep NN baselines, which may require 2k or more steps (Figure 5).

Figure: Figure 4: Pretrained PopT is more sample efficient when fine-tuning. Varying the number of samples available to each model at train time (x-axis), we see that the pretrained PopT is highly sample efficient, requiring only a fraction of samples (fewer than 500 samples out of 5-10k of the full dataset) to reach the full performance level of baseline aggregation approaches (dashed lines). Bands show standard error across test subjects. Stars indicate performance with full fine-tuning dataset.
Refer to caption: extracted/5912780/figures/sample_efficiency_v2.png

Figure: Figure 5: Pretrained PopT is consistently compute efficient when fine-tuning. Number of steps required for each model to reach final performance during fine-tuning (dashed lines). The pretrained PopT consistently requires fewer than 750 steps (each step is an update on a batch size of 256) to converge, in contrast to the 2k steps required for the non-pretrained PopT. Linear and Deep NN aggregation can be similarly compute efficient, but occasionally require more training steps depending on task. Bands show standard error across subjects. Stars indicate fully trained performance.
Refer to caption: extracted/5912780/figures/compute_efficiency_v2.png

Figure: Figure 6: Gains in decoding performance are available to new subjects. A minimal decrease in downstream decoding performance is found if the subject is held-out from pretraining (Held-out vs All). This is in stark contrast to the achievable downstream performance with a non-pretrained PopT (Non-pretrained PopT).
Refer to caption: extracted/5912780/figures/generalizability.png

Generalizability
To test if our pretrained weights will be useful for subjects not seen during training, we conduct a hold-one-out analysis.
We pretrain a model using all subjects except for one, and then fine-tune and evaluate on the model downstream.
We find that missing a subject from pretraining does not significantly affect the downstream results (see Figure 6).
This raises our confidence that the pretrained weights will be useful for unseen subjects and for researchers using new data.

Scaling with number of pretraining subjects
To investigate the effect of scaling pretraining data on our model, we pretrain additional versions of PopT using only 1, 2, or 3 subjects.
We find a consistent improvement in downstream decoding when we increase the number of pretraining subjects available across all our downstream decoding tasks Figure 7.
A significant improvement is found with just 1 pretraining subject already, potentially due to adaption to the temporal embeddings used.
The decoding performance using all our pretraining data is significantly higher in most decoding tasks than with just 1 or 2 subjects in the pretraining data, suggesting the potential for our framework to continue scaling with more subjects.

Ablation of loss components and position information
An ablation study confirms that both the ensemble-wise and channel-wise component of the pretraining objective contribute to the downstream performance (Table 3).
Furthermore, including the 3D position information for each channel is critical for decoding.
These findings also hold when the PopT is kept frozen during fine-tuning (see Appendix H: Frozen ablation).
Additionally, we find that the discriminative nature of our loss is necessary for decoding.
Attempting to add an L1 reconstruction term to our pretraining objective results in poorer performance, perhaps because the model learns to overfit on low-entropy features in the embedding.
Our discriminative loss requires the model to understand the embeddings in terms of how they can be distinguished from one another, which leads the model to extract more informative representations.

Figure: Figure 7: Pretraining with more subjects leads to better downstream performance. We pretrain PopT with different number of subjects (colors) and test on our decoding tasks (x-axis). Bars indicate mean and standard error of performance across channel ensembles 5-30 on a held out test subject. Pretraining with one subject gives a considerable benefit compared to no pretraining (red to yellow), but the addition of more subjects to pretraining consistently improves performance (yellow $\rightarrow$ green).
Refer to caption: extracted/5912780/figures/pretrain_scale.png

**Table 3: PopT ablation study. We individually ablate our losses and positional encodings during pretraining then decode on the resulting models. Shown are ROC-AUC mean and standard error across subjects evaluated at 90 electrodes. The best performing model across all decoding tasks uses all of our proposed components, showing that they are all necessary. Removing our positional encoding during pretraining and fine-tuning drops the performance the most, indicating that position encoding is highly important for achieving good decoding. Additionally, we attempt adding or only using a reconstruction component to the loss as a regularizing term, but find that this leads to poorer performance (last two rows).**
|  | Pitch | Volume | Sent. Onset | Speech/Non-speech |
| --- | --- | --- | --- | --- |
| PopT | $\mathbf{0.69}\pm\mathbf{0.07}$ | $\mathbf{0.84}\pm\mathbf{0.06}$ | $\mathbf{0.86}\pm\mathbf{0.05}$ | $\mathbf{0.89}\pm\mathbf{0.07}$ |
| PopT w/o group-wise loss | $0.66\pm 0.07$ | $0.83\pm 0.06$ | $0.84\pm 0.04$ | $0.88\pm 0.08$ |
| PopT w/o channel-wise loss | $0.67\pm 0.06$ | $0.81\pm 0.08$ | $0.84\pm 0.06$ | $0.87\pm 0.09$ |
| PopT w/o position encoding | $0.59\pm 0.07$ | $0.67\pm 0.10$ | $0.75\pm 0.08$ | $0.79\pm 0.08$ |
| PopT with reconstruction loss | $0.60\pm 0.11$ | $0.73\pm 0.11$ | $0.81\pm 0.05$ | $0.83\pm 0.09$ |
| PopT with L1 reconstruction only | $0.56\pm 0.04$ | $0.65\pm 0.08$ | $0.73\pm 0.10$ | $0.74\pm 0.10$ |

Figure: Figure 8: Probing the pretrained model for inter-channel connectivity Traditionally, connectivity analysis between regions is done by computing the coherence between electrode activity (left). We propose an alternative analysis purely based on the contextual sensitivity learned during pretraining. Briefly, we select an electrode, mask out its activity, and then measure the degradation in the channel-wise objective function for the remaining electrodes. Plotting the values of this delta (right) recovers the main points of connectivity. Plots for all test subjects can be seen in Appendix F: Connectivity.
Refer to caption: extracted/5912780/figures/connectivity_v2.png

## 6 Interpreting Learned Weights

Finally, we conduct two interpretability studies of the Population Transformer’s learned weights.

Connectivity
Traditional neuroscience analyses typically use cross-correlation as a measure of region connectivity (Wang et al., ).
Our PopT allows for an alternative method of determining connectivity, based on the degree to which channels are sensitive to each other’s context.
In this method, each channel is masked in turn, and then model performance on the pretraining channel-wise objective for the remaining unmasked channels is measured.
We use the degradation in performance as a measure of connectivity.
We can construct plots as in Figure 8, that recapitulate the strongest connectivity of the cross-correlation maps.
Note that while some approaches for modelling brain activity explicitly build this into their architecture (Cai et al., ), we recover these connections purely as a result of our self-supervised learning.
Additional method details available in Appendix E.

Candidate functional brain regions from attention weights
After fine-tuning our weights on a decoding task, we can examine the attention weights of the [CLS] output for candidate functional brain regions.
We obtain a normalized Scaled Attention Weight metric across all subjects to analyze candidate functional brain regions across sparsely sampled subject datasets Figure 9.
The Scaled Attention Weight is computed from raw attention weights at the [CLS] token passed through the attention rollout algorithm (Abnar & Zuidema, ).
The resulting weights from each channel are then grouped by brain region according to the Destrieux atlas (Destrieux et al., ).
A full description of the method is available in Appendix E.

The resulting weights reveal expected functional brain regions related to the tasks decoded Figure 9, with low-level auditory tasks highlighting primary auditory cortex and higher-level language distinction tasks highlighting language-specific areas.
Given the massive pretraining PopT undergoes, these scaled attention weights provide a valuable new tool for discovering candidate functional regions.

Figure: Figure 9: Attention weights from a fine-tuned PopT identify candidate functional brain regions Candidate functional maps can be read from attention weights of a PopT fine-tuned on our decoding tasks. For the Volume and Pitch tasks, note the weight placed on the primary auditory cortex (black arrows), but not in Wernicke’s area. For the Speech vs Non-speech and Sentence onset tasks, note the weight placed on regions near Wernicke’s area (black arrows). Center brain figure highlight regions related to auditory-linguistic processing (aph, )).
Refer to caption: extracted/5912780/figures/attention_v2.png

## 7 Discussion

We presented a self-supervised scheme for learning effective joint representations of neural activity from temporal embeddings.
Our approach improves decoding and reduces the samples required to learn downstream tasks, which is especially critical for neural data modalities given patient constraints.
A key aspect of our approach is the fact that we focus on spatial aggregation of existing channel embeddings, rather than training a large end-to-end model.
By decoupling temporal and spatial feature extraction, we are able to leverage existing temporal embeddings to learn spatiotemporal representations efficiently and with a smaller number of parameters.
This makes our model available for use in low compute-resource settings.
Furthermore, this separation of considerations opens up the possibility for future independent improvement in temporal modeling, whether that be from a domain specific model or a more general time-series encoder.
The generality of this approach allowed us to train on two very different neural modalities: scalp EEG and invasive iEEG.
Our success in these domains suggest that this approach could even be extended to settings outside of neuroscience that also contend with sparsely and variably distributed time-series data channels, as is often the case with geophysical or climate data.

Limitations and Future Work
We propose an effective strategy for aggregating signals with meaningful spatial coordinates, but it remains to be seen how we can generalize this further to aggregations without such coordinates.
Individual brains are also highly variable, so improving this positional encoding scheme could provide an even higher resolution latent space of each individual’s brain.
Future work could experiment with automatic functional identification for each channel, such as that explored in neural spiking data (Azabou et al., ), but it is currently unclear how to do so with neural recordings that have lower SNR.

## 8 Conclusion

We introduced a pretraining method for learning representations of arbitrary ensembles of intracranial electrodes.
We showed that our pretraining produced considerable improvements in downstream decoding and efficiency, that would not have been possible without the knowledge of spatial relationships learned during the self-supervised pretraining stage.
These benefits were found across data modalities, decoding tasks, and temporal encoders used, speaking to the generality of our approach.
We further showed that this scheme produces interpretable weights from which connectivity maps and candidate functional brain regions can be read.
Finally, we release the pretrained weights for our PopT with BrainBERT inputs as well as our code for plug-and-play pretraining with any temporal embedding.

## 9 Acknowledgements

We would like to thank our funding sources: Caltech Chen Institute, Caltech Carver Mead New Adventures Fund. This work was also supported by the Center for Brains, Minds, and Machines, NSF STC award CCF-1231216, the NSF awards 2124052 and 2123818, the MIT CSAIL Machine Learning Applications Initiative, the MIT-IBM Watson AI Lab, the CBMM-Siemens Graduate Fellowship, the DARPA Artificial Social Intelligence for Successful Teams (ASIST) program, the DARPA Knowledge Management at Scale and Speed (KMASS) program, the DARPA Machine Common Sense (MCS) program, the United States Air Force Research Laboratory and the Department of the Air Force Artificial Intelligence Accelerator under Cooperative Agreement Number FA8750-19-2-1000, the Air Force Office of Scientific Research (AFOSR) under award number FA9550-21-1-0014, NIH grant R01EY026025, and the Office of Naval Research under award number N00014-20-1-2589 and award number N00014-20-1-2643. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Department of the Air Force or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein. We would also like to thank Xuefei Wang and Amy Li for helpful comments and discussion.

## Appendix A Architectures and training

Pretrained PopT
The core Population Transformer consists of a transformer encoder stack with 6 layers, 8 heads.
All layers ($N=6$) in the encoder stack are set with the following parameters: $d_{h}=512$, $H=8$, and $p_{\text{dropout}}=0.1$.
We pretrain the PopT model with the LAMB optimizer (You et al., ) ($lr=1e-4$), with a batch size of $n_{\text{batch}}=256$, and train/val/test split of 0.98, 0.01, 0.01 of the data.
We pretrain for 500,000 steps, and record the validation performance every 1,000 steps.
Downstream evaluation takes place on the weights with the best validation performance.
We use the intermediate representation at the [CLS] token $d_{h}=512$ and put a linear layer that outputs to $d_{out}=1$ for fine-tuning on downstream tasks.
These parameters for pretraining were the same for any PopT that needed to be pretrained (hold-one-out subject, subject subsets, ablation studies).

Non-pretrained PopT
The architecture for the non-pretrained PopT is the same as the pretrained PopT (above).
However, no pretraining is done, and the weights are randomly initialized with the default initializations.

Linear The linear baseline consists of a single linear layer that outputs to $d_{out}=1$. The inputs are flattened and concatenated BrainBERT embeddings $d_{emb}=756$, TOTEM embeddings $d_{emb}=64$, Chronos embeddings $d_{emb}=512$, or TS2Vec embeddings $d_{emb}=320$ from a subset of channels $S\subset C$. Thus, the full input dimension is $d_{input}=d_{emb}*|S|$.

Deep NN
The inputs are the same as above, but the decoding network now consists of 5 stacked linear layers, each with $d_{h}=512$ and a GeLU activation.

Downstream Training For both PopT models, we train with these parameters: AdamW optimizer (Loshchilov & Hutter, ), $lr=5e^{-4}$ where transformer weights are scaled down by a factor of 10 ($lr_{t}=5e^{-5}$), $n_{batch}=256$, a Ramp Up scheduler (ildoonet, ) with warmup 0.025 and Step LR gamma 0.99, reducing 100 times within the 2000 total steps that we train for. For Linear and DeepNN models, we train with these parameters: AdamW optimizer (Loshchilov & Hutter, ), $lr=5e^{-4}$, $n_{batch}=256$, a Ramp Up scheduler (ildoonet, ) with warmup 0.025 and Step LR gamma 0.95, reducing 25 times within the 17,000 total steps we train for. For all downstream decoding, we use a fixed train/val/test split of 0.8, 0.1, 0.1 of the data.

Compute Resources To run all our experiments (data processing, pretraining, evaluations, interpretability), one only needs 1 NVIDIA Titan RTXs (24GB GPU Ram). Pretraining PopT takes 2 days on 1 GPU. Our downstream evaluations take a few minutes to run each. For the purposes of data processing and gathering all the results in the paper, we parallelized the experiments on 8 GPUs.

## Appendix B Model and Compute Requirements

**Table 4: Parameter counts. Since PopT takes existing temporal embeddings as input, the number of parameters that must be trained is an order of magnitude less than recent end-to-end approaches.**
|  | e5 | e50 | e90 |
| --- | --- | --- | --- |
| PopT |  | 20M |  |
| Deep NN | 3M | 20M | 36M |
| Linear | 3.8k | 38k | 69k |
| Brant (Zhang et al., ) |  | 500M |  |
| LaBraM (Jiang et al., ) |  | 350M |  |

**Table 5: Pretraining compute requirements Based on published train times (none were given for LaBraM) it is evident that PopT has smaller hardware and shorter training time requirements.**
|  | GPU count | GPU type | Time to train | TFLOPS |
| --- | --- | --- | --- | --- |
| PopT | 1 | NVIDIA TITAN RTX (24GB) | 2 days | 2.1M |
| Brant (Zhang et al., ) | 4 | NVIDIA Tesla A100 (80G) | 2.8 days | 18.8M |
| LaBraM (Jiang et al., ) | 8 | NVIDIA Tesla A800 (40G) | – | – |

## Appendix C Decoding tasks

We follow the same task specification as in Wang et al. (), with the modification that the pitch and volume examples are determined by percentile (see below) rather than standard deviation in order to obtain balanced classes.

Pitch
The PopT receives an interval of activity and must determine if it corresponds with a high or low pitch word being spoken.
For the duration of a given word, pitch was extracted using Librosa’s piptrack function over a Mel-spectrogram (sampling rate 48,000 Hz, FFT window length of 2048, hop length of 512, and 128 mel filters).
For this task, for a given session, positive examples consist of words in the top-quartile of mean pitch and negative examples are the words in the bottom quartiles.

Volume
The volume of a given word was computed as the average intensity of root-mean-square (RMS) (rms function, frame and hop lengths 2048 and 512 respectively).
As before, positive examples are the words in the top-quartile of volume and negative examples are those in the bottom quartiles.

Sentence onset
Negative examples are intervals of activity from 1s periods during which no speech is occurring in the movie.
Positive examples are intervals of brain activity that correspond with hearing the first word of a sentence.

Speech vs. Non-speech
Negative examples are as before.
Positive examples are intervals of brain activity that correspond with dialogue being spoken in the stimuli movie.

## Appendix D Data

**Table 6: Subject statistics Subjects used in PopT training, and held-out downstream evaluation. Table taken from Wang et al. (). The second column shows the number of uncorrupted, electrodes that can be Laplacian re-referenced. The average amount of recording data per subject is 4.3 (hrs).**
| Subj. | Age (yrs.) | # Electrodes | Movie | Recording time (hrs) | Held-out |
| --- | --- | --- | --- | --- | --- |
| 1 | 19 | 91 | Thor: Ragnarok | 1.83 |  |
|  |  | Fantastic Mr. Fox | 1.75 |  |  |
|  |  | The Martian | 0.5 | x |  |
| 2 | 12 | 100 | Venom | 2.42 |  |
|  |  | Spider-Man: Homecoming | 2.42 |  |  |
|  |  | Guardians of the Galaxy | 2.5 |  |  |
|  |  | Guardians of the Galaxy 2 | 3 |  |  |
|  |  | Avengers: Infinity War | 4.33 |  |  |
|  |  | Black Panther | 1.75 |  |  |
|  |  | Aquaman | 3.42 | x |  |
| 3 | 18 | 91 | Cars 2 | 1.92 | x |
|  |  | Lord of the Rings 1 | 2.67 |  |  |
|  |  | Lord of the Rings 2 (extended edition) | 3.92 |  |  |
| 4 | 9 | 135 | Megamind | 2.58 |  |
|  |  |  | Toy Story | 1.33 |  |
|  |  |  | Coraline | 1.83 | x |
| 5 | 11 | 205 | Cars 2 | 1.75 | x |
|  |  | Megamind | 1.77 |  |  |
| 6 | 12 | 152 | Incredibles | 1.15 |  |
|  |  | Shrek 3 | 1.68 | x |  |
|  |  | Megamind | 2.43 |  |  |
| 7 | 6 | 109 | Fantastic Mr. Fox | 1.5 |  |
| 8 | 4.5 | 72 | Sesame Street Episode | 1.28 |  |
| 9 | 16 | 102 | Ant Man | 2.28 |  |
| 10 | 12 | 173 | Cars 2 | 1.58 | x |
|  |  |  | Spider-Man: Far from Home | 2.17 |  |

## Appendix E Interpretation Methods

Connectivity analysis
We start with a pretrained PopT.
To test a particular channel’s contribution to connectivity, we mask it with all zeros.
Then, we consider the remaining unmasked channels and ask, how much increase do we see in the pretraining channel-wise loss?
Recall that this objective is to determine whether or not a channel has had its inputs swapped with random activity.
If the increase in loss is large, then we infer that the masked channel provided important context for this task.
Using this delta as a measure for connectivity, we then average across the Desikan-Killiany atlas regions (Alexander et al., ) and produce a plot using mne-connectivity (Gramfort et al., ).

Scaled Attention Weight
First, we obtain an attention weight matrix across all trials which includes weights between all tokens.
Then, we perform attention rollout (Abnar & Zuidema, ) across layers to obtain the contributions of each input channel by the last layer.
We take the resulting last layer of rollout weights for all channels, where the target is the [CLS] token, normalize within subject, and scale by ROC AUC to obtain the Scaled Attention Weight per channel.
Finally, we plot the 0.75 percentile weight per region, as mapped by the Destrieux atlas (Destrieux et al., ) using Nilearn ().

## Appendix F Connectivity

Figure: Figure 10: Full connectivity for our 7 test subjects. We compare between traditional connectivity analaysis performed via coherence (top row in each section) and the analysis based on our PopT pretrained weights (bottom row in each section). We note that our analysis usually recovers the strongest points of connectivity fromt the traditional analysis. Coherence was computed using scikit-learn’s (Pedregosa et al., ) signal.coherence.
Refer to caption: extracted/5912780/figures/connectivity_all.png

## Appendix G Frozen scaling

Figure: Figure 11: Pretraining is critical to frozen PopT performance that scales with the number of channels. As in Figure 3, we see that pretraining results in better downstream decoding and better scaling with the number of added channels. However, unlike in Figure 3, the PopT weights are frozen during fine-tuning, and only the linear classification head is updated. Bands show standard error across subjects. Results are shown for a frozen PopT with BrainBERT inputs.
Refer to caption: extracted/5912780/figures/baseline_ablations_vs_channels_frozen.png

## Appendix H Frozen ablation

**Table 7: An ablation study of the components of our approach for the frozen PopT. During pretraining, we alternate using either only the CLS or token contrastive component of the loss. We fine-tune these weights on all subjects. We find that both components contribute to the full model’s performance.**
|  | Sentence onset | Speech/Non-speech | Pitch | Volume |
| --- | --- | --- | --- | --- |
| Frozen PopT | $\mathbf{0.73}\pm\mathbf{0.06}$ | $\mathbf{0.72}\pm\mathbf{0.08}$ | $0.59\pm 0.06$ | $\mathbf{0.63}\pm\mathbf{0.07}$ |
| w/o cls | $0.67\pm 0.08$ | $0.68\pm 0.07$ | $0.58\pm 0.04$ | $0.60\pm 0.07$ |
| w/o replace loss | $0.69\pm 0.07$ | $0.69\pm 0.09$ | $\mathbf{0.59}\pm\mathbf{0.06}$ | $0.62\pm 0.06$ |
| w/o position encoding | $0.70\pm 0.07$ | $0.69\pm 0.07$ | $0.56\pm 0.08$ | $0.61\pm 0.06$ |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: - Figure 1 has a short-time Fourier transform, but I did not find it to be discussed anywhere in the text.
- It would be great to see how much Gaussian fuzzing contributes to the final performance. EEG is spatially very smoothed.
soundness: 3
strengths: - Validation across two types of neural recordings: EEG and intracranial EEG (iEEG)
- Thorough comparison against common aggregation methods and state-of-the-art models
- Strong ablations to test the importance of ensemble-wise loss, channel-wise loss, position encoding, and reconstruction loss
summary: The manuscript introduces the Population Transformer (PopT). This self-supervised framework tackles two key challenges in neural time-series analysis: sparse electrode distribution and variable electrode placement across subjects and datasets. PopT uses a transformer architecture that combines temporal embeddings with spatial information and trains through self-supervised learning to aggregate signals across variable electrode layouts. The model outperforms baselines on neural decoding tasks, requires less training data, generalizes to new subjects, and reveals interpretable brain connectivity patterns.
weaknesses: - In Tables 1, 2, and 3, consider adding statistical testing (e.g., Wilcoxon) between the best model and others. Then, correct p-values for multiple comparisons (e.g., Holm). Otherwise "We see significant improvements in performance" is not justified.
- I am confused about construction for " (1) ensemble-wise —the model determines if activities from two channel ensembles occurred consecutively, requiring an effective brain state representation at the ensemble-level". How do we know if activities from two channels occurred consecutively? In other words, could you explain how you construct pairs? How do you get the states?

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: In Fig. 8, how does the connectivity shown compare with cross-correlation or some other metrics that are used in the field?
soundness: 3
strengths: - They tackle an important and difficult problem for the field

- Interesting approach backed with strong empirical results.

- I appreciate the work authors put in doing new analyses and interpreting the weights. 

-Well written and interesting.
summary: The authors introduce a self-supervised learning framework called PopT to learn embeddings for EEG and iEEG recordings for the improving downstream decoding analysis. In general, the problem is hard because of the spatially varying channel placements in different experiments. Their method leverages unannotated data to optimize channel-level and ensemble-level objectives, which helps them build generic representations and also allows to capture some dynamical relationships. Their tests show improved decoding on held-out subjects. Based on the interpretation of the weights, they propose a new method for brain region connectivity analysis and for identifying candidate brain regions.
weaknesses: I have some questions about their connectivity analyses below

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: # **Major Questions**

1. You show that the decoding performance increases with more subjects and decreases with holding out a subject. However, these experiments are limited because they consider a validation on specific channels of one specific subject, or holding out one specific subject. Is this also valid across arbitrary channels and arbitrary subjects?  
2. Figure 4: Do the baseline aggregation approaches have fewer free parameters and nonlinearities than the PopT? Could you report the number of free parameters and nonlinearities of all networks and factor this into the discussion of your results?  
3. Results 366-370: The model's ability to generalize is a central claim. However, the hold-one-out analysis seems somewhat limited. What if the held-out subject was a lucky draw? Can this comparison be done in a crossvalidation? Can the performance decrease be computed for hold-k-out and presented as a function of k? Separately, is there a different model class that can generalize to provide a comparison group in this test?  
4. 371-377 and Figure 6 and 7: I would have expected that "All subjects" and "0 subjects" correspond to "All" and "Non-pre-trained PopT" but the numbers are different. Can I rely on this data? Why is there a mismatch?  
5. Figure 7: What if "across channel ensembles 5-30 on a held out test subject" is a lucky draw? Could this be more robustly supported with an extended analysis?  
6. 414-416: To say this, I would expect an analysis of representations resulting from different loss functions. Can one show that the contrastive learning objective results in the clearest association/separation of inputs?  
7. 468-470: One obtains N-squared performances, where N is the number of channels. Can one present the raw functional connectivity matrices inferred this way in addition to the aggregate? This could help to gauge how effective the approach is.  
8. 468-470: Can one compare the results to the functional connectivity inferred with traditional cross-correlation methods on the measurements? This would help to illustrate the true utility of the PopT-based analysis.  
9. 475-485: Same here. Can one compare to neuroscience ways to characterize the activated brain area, e.g., z-score based?  
10. Results 296-297: What are differences in, e.g., architecture or training that makes LaBraM outperform the PopT? Could one provide possible explanations for why LaBraM outperforms the PopT on EEG decoding (whereas end-to-end approaches specifically designed for iEEG do not)?  
11. Figure 4: there is a much more significant improvement for the sentence onset and speech vs. non-speech tasks over the improvement in the pitch and volume tasks. Could one provide an explanation for this in the results section?  
12. In contrast to one of the central claims in the abstract, 23-24, the codebase is messy, and lacks documentation, clear examples, instructions, and tutorials. The readme title is BrainBert 2.0. Could the codebase please be significantly cleaned and documented?

# **Clarifications and Minor Questions**

13. Abstract 16-18: The sentence "lowers the amount of data for downstream decoding experiments" is compatible with the presented results (Fig. 4 and 5). However, the verb 'lowers' requires a comparison group, and 'downstream decoding experiments' is vague. The meaning of 'downstream' is also ambiguous. Could you please rephrase this?  
14. Abstract 18-20: Calling the PopT computationally lightweight in comparison to end-to-end approaches is misleading because the PopT only works with pretrained weights from BrainBert. Could you please rephrase this?  
15. Abstract 18-20: Calling the PopT "more interpretable" when not comparing the interpretations to existing traditional neuroscientific methods for interpretability is misleading. Could you please rephrase this?  
16. Abstract 18-20: What does 'retaining competitive performance' mean? In all iEEG data, it outperforms the end-to-end method used, in all EEG data it is outperformed by LaBraM. Could you please provide this information without ambiguity?  
17. Abstract 21-23: What are "massive" amounts of data according to the authors?  
18. Discussion 524-525: "\[ … \] could provide an even higher resolution latent space \[ … \]". Higher than what? Could you please rephrase this?  
19. Conclusion 538-539: I wouldn't know how to use it "for plug-and-play pretraining". The codebase is messy and lacks clear examples, instructions, or tutorials. Is this a work in progress or the cleaned-up codebase? 
20. 206-208: The explanation is not reader-friendly and not comprehensible for someone who does not know the types of neurorecordings. Could you please explain to the reader why intracranial and scalp electroencepholography lead to two different types of time-series data, and why and which resolution extremes these represent?  
21. 3 Population Transformer Approach 178-179: these are rather "contrastive" components and not "discriminative" components. Could one please adopt "contrastive" instead of discriminative throughout the text?  
22. Figure 2, 6, 7: Could one use violin (or box) plots instead of bar plots to convey the variability and scatter the individual data points on top to convey the sample size?  
23. 375: "potentially due to adaption to the temporal embeddings used" meaning unclear. Could you please rephrase this?  
24. 363-365: These sentences seem contradictory, which might be easily resolved with rephrasing them. Could you please rephrase this? Part of the Figure 5 caption could go here into the main text.  
25. Results Figure 4: What is 'highly' sample efficient?

# **Minor Formatting**

26. 214-215: Could you please remove the parenthesis inside the parenthesis?  
27. Table 1: Could you please emphasize your result better by bolding the label Pretrained PopT or star it for "best overall"?  
28. Results 264-269: hard to read. Could you please remove the parenthesis inside the parenthesis?  
29. Figure 3 is referred to before Figure 2 in the text. Could you please renumber the figures to maintain clarity?  
30. Table 1 and Figure 1: Could you please make latex keep them inside the section where they are referred to?
soundness: 3
strengths: The claimed contributions of the PopT would present an exciting leap forward in decoding neural time series data. This is because unreliable and slow decoding of neural activity limits applications and scientific insights from models. A method that requires comparably few data, no labeled data, and learns representations that generalize across subjects regardless of electrode layout makes decoding more fast and reliable for new subjects and unlocks new applications. The ability of the PopT to generalize across subjects is a major strength of the approach. Another strength is that the PopT relies on pre-trained encoding models that it is agnostic to. The formulations of the contrastive learning objectives seem original. The suggested ways to interpret the model to understand the data are interesting.
summary: The authors introduce Population Transformer (PopT), a transformer-based and contrastive learning approach for decoding neural time series data. The PopT aggregates and decodes channel-wise temporal encodings of neural time series data from, e.g., a pretrained BrainBert but is agnostic to the encoder. The authors introduce two contrastive loss terms for training the PopT. One contrastive loss term requires the PopT to learn temporal representations for each channel, and the other contrastive loss term requires the PopT to learn spatial representations across channels. The PopT receives electrode placement as input, allowing it to decode neural activity across varying electrode placements (e.g., from several subjects). In an application with EEG and iEEG, the paper presents that the PopT converges with fewer samples than other shown methods, achieves superior decoding performance compared to other shown methods, and provides evidence for the PopT's ability to generalize across subjects. The authors also suggest measuring and interpreting functional connectivity based on the ablation of individual channel input to their PopT and performance degradation for all other channels. They also suggest interpreting the weights of the model to map the task-associated function to specific electrodes.
weaknesses: The paper presents multiple weaknesses that could be addressed to improve its clarity and reliability. 

First, the abstract describes PopT's benefits, such as "computationally lightweight" and "more interpretable," which lacks precision without clear comparison groups or quantitative metrics. For example, calling PopT computationally lightweight in comparison to end-to-end trained models is misleading since it relies on pre-trained, e.g., BrainBert weights, which already bear significant computational demands. As another example, the interpretability analysis does not include any comparison method.

In addition, the generalizability claims lack thorough validation: the hold-one-out subject analysis does not sufficiently establish robustness. This experiment could benefit from a broader validation, e.g., cross-validation approaches. 

Additionally, when comparing the performances to other models the analysis does not account for the differences in free parameters and nonlinearities across models. 

Last, the codebase requires significant improvements, as it is messy, poorly documented, and lacks proper instructions or tutorials. For example, the readme is incorrectly titled BrainBert 2.0. This contrasts the authors claim of it being "plug-and-play."

I am not up-to-date and fully familiar with related work. This is why my expertise might be insufficient to fully weigh the significance of the central contribution against the current limitations of the paper in the clarity and reliability of the results. 

I am generally excited about the direction of the work and happy to give a higher evaluation once the presentation of the current version of the paper is significantly improved and the detailed questions below are addressed.

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 1
questions: 1. Table 3 refers to “PoPT w/o group-wise loss”, is the group-wise loss the same of ensemble-wise loss. To improve clarity, could the authors consider using consistent terminology throughout?
2. It would be insightful to see how one of the baseline models, specifically BrainBERT from Table 1, performs on the hold-out dataset, similar to the performance reported for PoPT in Figure 6. This comparison would provide additional context for the robustness of PoPT relative to other methods. Could the authors report the performance of BrainBERT on the hold-out dataset in a format similar to Figure 6?
soundness: 3
strengths: This paper demonstrates strong rigor by evaluating the proposed method across two types of neural time-series data (iEEG and EEG), enhancing the generalizability and robustness of its findings. The authors plan to share both the data and code upon acceptance, promoting transparency and reproducibility within the community. Additionally, they test their method using hold-out subject tests to validate the model’s performance on unseen data, which is essential for assessing real-world applicability. The method’s effectiveness is benchmarked against a diverse set of models, providing a comprehensive view of its performance relative to other approaches. Furthermore, ablation studies are conducted to examine the contribution of different components in the proposed framework, offering valuable insights into how each part enhances overall model performance.
summary: This paper introduces a self-supervised framework, the Population Transformer (PopT), designed to learn population-level representations for large neural recording datasets.  PopT addresses challenges related to sparse and variable electrode distribution across subjects and datasets by using pretrained embeddings. By using a modular approach the model is more flexible and lightweight. The authors claim that the approach is computationally efficient, interpretable, and performs competitively against end-to-end methods.
weaknesses: The study offers a comprehensive evaluation of PopT, and assess it using various experiments. However, the paper would benefit from clearer organization and writing, particularly in emphasizing the significance of each experiment.

- **Introduction and Field Context**: Despite the extensive experimentation, the paper lacks a clear introduction to the field and a focused statement of its goals. Important background information and foundational definitions are often buried in the appendix rather than integrated into the main text.
    - For instance, a key innovation of the study is the use of channel-level and ensemble objectives. Providing a more detailed literature review on current approaches to these objectives would help the reader understand the limitations of existing methods and the advantages of the proposed improvements. This contextualization would make the study’s contributions clearer and more impactful.
    - The connection between experiments is also unclear, with critical information about the pretrained PopT model only appearing in the appendix. Providing a clear description of the data used for pretraining within the main text would help readers understand the study’s foundation and goals more effectively.
    - Similarly, Table 3 introduces PopT with an L1 reconstruction loss as an additional experiment. However, this experiment is not discussed in the text, and it appears tangential to the core contributions of the study. Omitting such details could allow space to expand on more relevant analyses.
- **Benchmarking Choices**: While Chronos/TS2Vec and BrainBERT/TOTEM are used as benchmarks for PopT, it remains unclear why these specific models were selected. The authors could strengthen the study by discussing the criteria for choosing these benchmarks.
- **Minor Presentation Issue**: In Figure 8, the plot slightly overlaps with the title on the left image.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 2
questions: **Clarity**:  the paper is not really neat. Moreover, there is some sloppiness in the mathematical notation and technical explanation of the method. Finally, the figures are uneven and quite messy. Follows a list of detailed complaints:
- general lack of dimensionality: in many points, it is not clear what the dimensions are of the variables. for example, lines 160-161, what is the dimension of $X_B^t$? Is the “+” used for concatenation or for an actual sum? What is the output dimensionality of the temporal embedding model $B$? Some of these are in Appendix A, but it would be useful to streamline and clarify.
- $[CLS]$ token: it is stated nowhere that this is the token used for downstream decoding/classification tasks, and what space it belongs to. Is it only binary? Please clarify in the text or figure legend.
- Self-supervised Loss (line 173-197): the losses are described in words, but it would be useful and much clearer to write them in formulas. Right now, this paragraph is a bit messy and somewhat colloquial.
- Fine-tuning (line 199-201): colloquial and relies on details only available in the appendix. It would also be useful to formally define what is meant by a decoding task.
- Figures: Fig 1 is unclear and messy. Why does it go from bottom to top, rather than the opposite? Currently, the “Population Transformer inputs” title of Fig 1a is just above the output, making it very confusing. Moreover, the figure has at least four different font sizes. The “+” used in temporal embedding + 3D position is confusing; does it refer to concatenation or addition. The title of the colorbar near the STTF is nearly illegible. Panels b and c could provide clearer information on the loss. Fig 3-4: font sizes widely different across figures; Fig 3 text is hard to read on print. Same for Figs 6-7-8. Fig 8 left title is partially covered by the plot.

**Interpretability**: the claim of interpretability requires some work to be convincing/useful. For example, the connectivity analyses are far from clear (also after reading the Appendix!). It is not immediately clear how one infers (pairwise) connectivity by masking out the activity of one channel at the time and measuring the degradation in channel-wise loss. Can you add a few details on that? And do we learn anything new from this analysis as compared to previous coherence or cross-correlations analyses? Can we have a few more qualitative and/or quantitative comparisons? Same questions go for the functional brain regions from attention weights. Moreover, what are the possible caveats of inferring connectivity or functional regions from PopT?

**Minor comments**:
- line 153: *adapt* -> adopt?
- lines 132, 158, …: BrainBERT was used but never cited.
- line 163: *sinusoidal position encoding*: is there a simple formula or a couple of words to explain this method without necessarily having to read the cited paper?
- line 171: and $\tilde y_i$ *and(?)* respectively, …
- line 188: …is *again(?)* the binary cross-entropy…
- line 374, 478, 483: Figure X -> (Figure X)

**Additional questions**:
*[These questions popped up while first reading the paper; some got clearer on a second read, some not. I don’t think it’s necessary for the authors to answer all these questions, but hopefully, they’ll be useful for improving discussion and intro, or future experiments, or mybe just food for thoughts.]*

The paper claims to learn representations of neural activity; do the authors mean the final output of the PopT or only the CLS output? Can one also learn a neural representation without using a CLS token? Can this representation be thought of and used as a dimensionality reduction method? Who are the intended users of this method? Can this method be used for non-human invasive research, such as calcium imaging, neuropixels, etc? Is it possible that the discriminative nature of the pre-training objective leads to potentially misleading representations in case of noisy/faulty channels?
soundness: 3
strengths: This paper has many merits. The idea and approach, from a computational neuroscience perspective, are novel and very interesting. The proposed framework could benefit the neuroscience community at large: combining data from multiple datasets, and learning brain-wide representations that are not idiosyncratic to a specific individual are key to advancing neuroscience research. Many points made throughout the paper and several experiments are convincing and valuable, and I don’t think much work is required in terms of experiments/compute.
summary: The Population Transformer (PopT) is a self-supervised framework designed to learn brain-wide neural activity representations. The approach aims to address common challenges faced by invasive iEEG and non-invasive EEG neural recordings, namely 1) sparsity of recording channels within the brain and 2) variability in channel positioning and reliability. The authors propose a smart data embedding and combining two loss functions that allow pre-training of the network using data from multiple datasets and individuals. Pre-training is shown to be fundamental to lowering the amount of data necessary for fine-tuning neuroscientific tasks, such as decoding experiments. Finally, the authors show experiments and analyses to support the claim that this framework is interpretable and generalizable.
weaknesses: On the flip side, the paper is, at some points, lacking in precision and clarity, at some points a bit colloquial, and at other points, unnecessary jargon. I think that the paper would be stronger, and more accessible, if the authors put some effort into clarifying and streamlining it. Moreover, the claims of interpretability are currently unclear and potentially overstated.

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 2
questions: Page 3, line 161: According to the paper's description, the input to the temporal encoding is a single value of a channel at a time t. Is this correct?
soundness: 3
strengths: In this paper, the author tried to solve the problem of sparse and variable electrode distribution across subjects and neural time series data datasets by modeling with a single-channel embedding method.

The effectiveness of the proposed method was supported through various experiments and analyses.

It achieved high performance with fewer computational resources compared to existing pre-trained models.
summary: In this paper, the authors propose a self-supervised learning framework for neural time series data. The model is trained using ensemble and channel-wise discrimination strategies to address the problem of sparse and variable electrode distribution across subjects and datasets. The authors presented the validity of the proposed method on two datasets with diverse analyses.
weaknesses: It is unclear how exactly the position of each channel is defined in spatial position encoding. The reference provides coarse positions only, which is important information in spatial relation learning. 
Page 4, line 166: The notations used in position embedding are not defined. 

Supporting materials must be presented for the Gaussian fuzzing in channel embedding to ensure its effectiveness in avoiding overfitting.

The experiments are somehow limited. More rigorous and thorough experiments should be conducted over diverse datasets. Other more challenging tasks, e.g., TUEV in the TUH EEG corpus, should be conducted and compared with the comparative methods.

The results in the ablation study (Table 3) suggest that position encoding played a significant role in performance improvements, while other proposed loss functions had marginal effects. In this regard, the related contexts should be provided in detail. 

Regarding the connectivity construction in Section 6, it is not straightforward to understand how the connectivity is measured. More detailed explanations and rationales should be provided regarding why the proposed method can estimate channel connectivity.

Meta-review: The paper introduces a self-supervised model called Population Transformer to model brain-wide neural activity sparsely and variably measured across subjects and datasets. Representations generated by this pre-trained model can then be used to perform downstream decoding tasks, leading to superior accuracy compared to models only trained on one specific dataset. The paper received overwhelmingly positive reviews, especially after the revision to improve the presentation and add additional evaluation and ablation components.

Decision: Accept (Oral)