SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Daoze Zhang*
zhangdz@zju.edu.cn
Zhizhang Yuan*
zhizhangyuan@zju.edu.cn
Yang Yang†
yangya@zju.edu.cn
Junru Chen
jrchen_cali@zju.edu.cn
Jingjing Wang
wjxjj@zju.edu.cn
Yafeng Li
Nuozhu Technology Co., Ltd.
yafeng.li@neurox.cn
We propose a foundation model named Brant for modeling intracranial recordings, which learns powerful representations of intracranial neural signals by pre-training, providing a large-scale, off-the-shelf model for medicine. Brant is the largest model in the field of brain signals and is pre-trained on a large corpus of intracranial data collected by us. The design of Brant is to capture long-term temporal dependency and spatial correlation from neural signals, combining the information in both time and frequency domains. As a foundation model, Brant achieves SOTA performance on various downstream tasks (i.e. neural signal forecasting, frequency-phase forecasting, imputation and seizure detection), showing the generalization ability to a broad range of tasks. The low-resource label analysis and representation visualization further illustrate the effectiveness of our pre-training strategy. In addition, we explore the effect of model size to show that a larger model with a higher capacity can lead to performance improvements on our dataset. The source code and pre-trained weights are available at: https://zju-brainnet.github.io/Brant.github.io/.
Brain signals are electrical impulses that are generated by brain neurons and transmit through neural networks. These signals provide important information about brain activity and can usually be monitored in two ways, namely scalp electroencephalography (EEG) and intracranial electroencephalography (iEEG). The former records the electrical brain activity through electrodes placed on the scalp, while the latter implants intracranial electrodes into brain tissue directly to derive neural recordings. Compared with EEG, iEEG manifests significant advantages by providing more abundant, stereotactic and detailed information about brain wave patterns from deeper brain structures, which has been the mainstream method to obtain deep brain information, and is essential in the therapies for some brain diseases (e.g., Parkinson's disease [1], epileptic seizure [2]).
Modeling intracranial recordings has drawn much research attention, but several issues still remain unresolved. Currently, studies for modeling intracranial recordings are mainly divided into two research lines, namely handcrafted feature based methods [3-8] and deep learning based methods [9-12]. Handcrafted feature engineering requires lots of domain knowledge and may only work on specific tasks. And most deep learning based methods are fully supervised, which relies heavily on
* Equal contribution.
labeled data. However, labeling data at scale in medical experiments is often infeasible or expensive, which underscores the importance of maximizing the label efficiency. To overcome these limitations, the paradigm of self-supervised pre-training followed by fine-tuning with few samples can greatly reduce the reliance on labels and enable the model to generalize to various downstream tasks.
Moreover, modeling intracranial recordings requires careful consideration of several key factors. (1) Long-term dependency. Since intracranial neural signals are time series and gradual changes in brain activity may only be captured by the long-period analysis [13], long-term temporal dependency is crucial in modeling intracranial data. (2) Spatial correlation. The electrodes implanted in the brain contain many contacts (also called channels), which are distributed across various brain regions. Due to the fact that brain waves propagate through different brain regions [14], signals recorded from different channels can be spatially correlated, reflecting the underlying neural activity. (3) Time and frequency domains. For neural recordings, the time domain provides information about the amplitude and duration while the frequency domain can reveal underlying oscillatory patterns and rhythms [15]. Therefore, modeling neural signals in both domains can provide information more consistent with the neurophysiological mechanisms [16]. To the best of our knowledge, no existing work on intracranial recordings considers all the three key factors simultaneously.
In view of the unresolved issues above, we propose a foundation model for intracranial neural signal named Brain Neural Transformer (Brant). The design of our model takes all the three key factors (i.e., long-term dependency, spatial correlation, time and frequency domains) for intracranial signal modeling into account. Moreover, Brant contains more than 500M parameters and is pre-trained on a large intracranial dataset with 1.01 TB data, which can be adapted to accomplish various downstream tasks. Compared to other existing methods for modeling brain signals, Brant can achieve better performance with far fewer labeled samples, showing the great benefit of our work in medical scenarios. As an off-the-shelf model along with the code and weights, Brant can participate in other medical research and experiments, which alleviates the issue of sample and label efficiency.
Figure 1: The model scale of existing brain signal models, including RP [17], TS [17], CPC [17], BENDR [18], MVTS [19], BrainBERT [20], LGGNet [21], EEG-GCNN [22], BrainNet [23], SICR [24], SEEG-Net [9], and two works from Banville et al. [25] and Tang et al. [26].
To sum up, the main contributions of our work comprise:
- We propose a foundation model for intracranial neural signals named Brant, which is the largest model on brain signals (shown in Fig. 1) and pre-trained on a large intracranial dataset collected by us, providing a large-scale and off-the-shelf model for medicine.
- To our knowledge, Brant is the first to date that attends long-term dependency and captures spatial correlation across channels, while combining the information from both time and frequency domains.
- Extensive experiments show that Brant generalizes well to various downstream tasks, showing the great potential in neural recordings modeling. Further analysis illustrates the effectiveness of large-scale pre-trained model, demonstrating the medical value of our work.
Model overview. As previously mentioned, we propose Brant to capture long-term dependency and spatial correlation from intracranial recordings, while combining the information from both time and frequency domains. Our model mainly consists of two Transformer encoders, namely, temporal encoder and spatial encoder (shown in Fig. 2). The temporal encoder encodes a sequence of  $L$  consecutive patches which focus on the temporal dependency, and the spatial encoder encodes  $C$  patches with the same time indices from all channels to capture their underlying spatial correlation.
Figure 2: The pre-training framework of Brant. The collected intracranial neural recordings are first processed to a set of patches  $\pmb{p} \in \mathbb{R}^{N_p \times C \times M}$  as the input. Then each time we mask a subset of patches  $\pmb{p}_{j:j+L-1}$  and map the input patches to the hidden space while adding positional and frequency information to derive the input encoding  $\tilde{\pmb{p}}_{j:j+L-1}$ . The temporal encoder encodes the input encoding in each channel to derive the temporal representations  $h_{j:j+L-1}$ . The spatial encoder encodes the temporal representations with the same time indices from all channels to obtain the final representations  $z_{j:j+L-1}$ . Then a linear head is used to obtain the reconstructed patches  $\hat{\pmb{p}}_{j:j+L-1}$  from the final representations.
During pre-training, we randomly mask the input signals recorded from the implanted electrodes, then linearly map the original signals to the latent space and add positional and frequency information through input encoding. Then the input encoding is encoded to latent representations and a linear projection head is added on the top of the spatial encoder which maps the representations to reconstruct the original signals. The details of our model and pre-training task are described as below.
**Patching.** Since neural recordings are electrical signals with high sampling rates, we aggregate timestamps into patches to (1) enhance the locality and extract semantic information; (2) reduce computation and memory usage; and (3) attend a longer temporal dependency [27]. Specifically, given a neural signal  $\pmb{x} \in \mathbb{R}^{N \times C}$ , where  $N$  is the number of timestamps and  $C$  is the number of electrode channels, we divide  $\pmb{x}$  with length  $M$  and stride  $S$  to generate a set of patches  $\pmb{p} \in \mathbb{R}^{N_p \times C \times M}$ , where  $N_p = \left\lfloor \frac{N - M}{S} \right\rfloor$  is the number of patches in each channel.
Frequency encoding. We propose frequency encoding to explicitly inject the information on frequency domain to the observed data. The frequency encoding is mainly based on power spectral density (PSD) which describes the distribution of a signal's total average power over frequency (details about PSD are in App. A). To be specific, we split the frequency domain into several bands according to the standard description for rhythmic activity [28]: (1)  $\theta$  (4-8Hz), (2)  $\alpha$  (8-13Hz), (3)  $\beta$  (13-30Hz), (4)  $\gamma 1$  (30-50Hz), (5)  $\gamma 2$  (50-70Hz), (6)  $\gamma 3$  (70-90Hz), (7)  $\gamma 4$  (90-110Hz), (8)  $\gamma 5$  (110-128Hz). For the  $i$ -th frequency band, a learnable encoding  $\pmb{f}_i$  is set as its representation which is shared across all the patches. Then we compute the absolute spectral power of each patch  $p_{j,c} \in \pmb{p}, j = 1,\dots,N_p; c = 1,\dots,C$  in the  $i$ -th frequency band:
$$
P _ {j, c} (i) = \log \sum_ {\omega \in \operatorname {b a n d} (i)} P S D _ {p _ {j, c}} (\omega), \quad i \in \{1, 2, \dots , 8 \}, \tag {1}
$$
which acts as the weight of  $\pmb{f}_i$ . The frequency encoding  $\mathbf{F}_{j,c} \in \mathbb{R}^D$  of patch  $p_{j,c}$  is obtained as the weighted sum of the learnable encodings  $\pmb{f}_i$ :
$$
\mathbf {F} _ {j, c} = \sum_ {i = 1} ^ {8} \frac {\exp \left(P _ {j , c} (i)\right)}{\sum_ {i ^ {\prime} = 1} ^ {8} \exp \left(P _ {j , c} \left(i ^ {\prime}\right)\right)} \boldsymbol {f} _ {i}. \tag {2}
$$
Encoding process. It contains several steps to encode the input signal to latent representations. Specifically, the input  $\pmb{p}_{j:j + L - 1} \in \mathbb{R}^{L \times C \times M}$  contains  $L \times C$  patches and  $\pmb{p}_{j:j + L - 1,c} \in \mathbb{R}^{L \times M}$ .
denotes the patches in  $c$ -th channel. We first map each sequence of patches  $\pmb{p}_{j:j + L - 1,c}, c = 1,\dots,C$  to the latent space of dimension  $D$  by a linear projection  $\mathbf{W}_{\mathrm{proj}} \in \mathbb{R}^{D \times M}$ , and the projected input will be added with a learnable positional encoding  $\mathbf{W}_{\mathrm{pos}} \in \mathbb{R}^{L \times D}$  which monitors the temporal order of patches and the frequency encoding:
$$
\tilde {\boldsymbol {p}} _ {j: j + L - 1, c} = \left(\mathbf {W} _ {\text {p r o j}} \boldsymbol {p} _ {j: j + L - 1, c} ^ {\mathrm {T}}\right) ^ {\mathrm {T}} + \mathbf {W} _ {\text {p o s}} + \mathbf {F} _ {j: j + L - 1, c}, \tag {3}
$$
where  $\tilde{\pmb{p}}_{j:j + L - 1,c}\in \mathbb{R}^{L\times D}$  denotes the input encoding of the original signals  $\pmb{p}_{j:j + L - 1,c}$ . The input encoding will be fed into the temporal encoder to obtain temporal hidden representations  $\pmb{h}_{j:j + L - 1,c}\in \mathbb{R}^{L\times D}$  and we denote the temporal representations of the whole input as  $\pmb{h}_{j:j + L - 1}\in \mathbb{R}^{L\times C\times D}$ . The spatial encoder further captures the spatial correlation across channels, which takes each of the temporal representations  $\pmb{h}_k\in \mathbb{R}^{C\times D}, k = j,\dots,j + L - 1$  as the input and outputs the final representations  $\pmb{z}_k\in \mathbb{R}^{C\times D}, k = j,\dots,j + L - 1$ . Overall, given the input signal  $\pmb{p}_{j:j + L - 1}$ , the encoding process returns the corresponding latent representations  $\pmb{z}_{j:j + L - 1}\in \mathbb{R}^{L\times C\times D}$ .
Self-supervised pre-training. Self-supervised representation learning is a powerful approach to extract high level abstract representation from unlabelled data. Among those methods to learn representation via self-supervised pre-training, masked autoencoder (MAE) has been proved to be a simple but effective way in many fields [27, 29, 30]. We apply this technique to our self-supervised pre-training, in which the model is trained to reconstruct the whole input given its partial observation.
Given the input patches  $\pmb{p}_{j:j + L - 1}$ , we mask a subset of patches uniformly at random and encode the masked patches to the latent representations  $\mathbf{z}_{j:j + L - 1}$ . During the pre-training stage, the representations will be fed into a flatten layer with linear head  $\mathbf{W}_{\mathrm{rec}} \in \mathbb{R}^{M \times D}$  to reconstruct the original patches. Finally we utilize an MSE loss to measure the discrepancy between the reconstructed patches  $\hat{\pmb{p}}_{j:j + L - 1}$  and the original patches  $\pmb{p}_{j:j + L - 1}$ .
Pre-training dataset. Brant is pre-trained on 1.01 TB neural data, a large clinical intracranial neural signal dataset recorded by stereo-electroencephalography (SEEG) technique from a first-class hospital. The subjects undergo a surgical procedure to implant 4 to 11 invasive electrodes, each with 52 to 153 channels, in their brain. The dataset contains 2528 hours of  $1000\mathrm{Hz}$  recordings with more than 1 trillion timestamps. More details are in App. E. We down sample the original signals to  $250\mathrm{Hz}$  and generate a set of patches of 6s (1500 timestamps).
Downstream dataset. Another neural dataset collected by us with seizure labels is used to fine-tune and evaluate our model. It contains 29.39 GB data with 43 hours of  $1000\mathrm{Hz}$  intracranial recordings and we do the same preprocessing (i.e. down sampling and patching) to it as the pre-training dataset. Professional neurosurgeons participate in the labeling of epileptic seizures and the labels consist of two categories, namely, seizure and normal samples. More details are in App. E. For each downstream task, we sample a small subset from the downstream dataset for fine-tuning and evaluation.
3.2 Pre-training
For the model configurations, the temporal encoder contains a 12-layer Transformer encoder with model dimension 2048, inner dimension (FFN) 3072 and 16 attention heads, and the spatial encoder contains a 5-layer Transformer encoder with model dimension 2048, inner dimension 3072 and 16 attention heads. During the pre-training,  $40\%$  patches in each input sample are masked with zero values uniformly at random. We take 16 input samples as a minibatch and each minibatch contains an average of  $24k$  patches. The model is pre-trained on a Linux system with 2 CPUs (AMD EPYC 9654 96-Core Processor) and 4 GPUs (NVIDIA Tesla A100 80G) for about 2.8 days.
We optimize with Adam [31], updating the model parameters every 4 steps, and the model trains for  $750\mathrm{k}$  updates in total. A cyclic scheduler that adopts a basic triangular cycle without amplitude scaling is utilized to adjust learning rate during pre-training. Specifically, we set the basic learning rate as  $3\times 10^{-6}$  and the maximum learning rate as  $1\times 10^{-5}$ , then the learning rate steps up (down)
for every 8k updates. We apply mixed precision training with FP32 and BF16 to reduce the memory usage in GPUs for acceleration.
3.3 Downstream tasks
As Brant is a foundation model for intracranial recordings, we conduct extensive experiments on several downstream tasks, including short- and long-term signal forecasting, frequency-phase forecasting, imputation and seizure detection, to verify the high capacity of our model in modeling intracranial recordings. We conduct each downstream task for Brant on two settings: (1) fine-tune the model with a learning rate of  $1 \times 10^{-7}$ ; (2) freeze the pre-trained weights. Five random runs were performed to obtain the mean and standard deviation. The detailed setups of these downstream tasks are as follows:
Short- and long-term signal forecasting. Predictive observation of the neural signal values is beneficial for the development of warning systems for patients in need of precautionary measures [32, 33]. Therefore, we adopt short- and long-term signal forecasting, in which the learned representations are fine-tuned to predict future signals with different lengths given a past sequence. The past sequence length is set as 15 patches (90s) and the prediction lengths are set as 2 patches (12s) and 20 patches (120s) in short- and long-term forecasting, respectively. We sample 400 minutes of recordings from the downstream dataset, then randomly split into 320 minutes for fine-tuning and 80 minutes for evaluation. A linear prediction head is used to predict the future signals. We adopt MAE and MSE as the performance metrics.
Frequency-phase forecasting. In medicine, it is essential to predict the physical features like frequency and phase of brain signals to provide guidance for some treatments. For example, in a therapy that modulates brain activity called transcranial alternating current stimulation (tACS), the stimulation control heavily depends on these knowledge about the target brain activity [34]. Therefore, we set up the frequency-phase forecasting task in order for these therapies like tACS to be most effectively used in the treatment of brain disorders [35].
Given a past sequence, this task is to predict the dominant frequency and phase information (see details in App. B) of intracranial signals in the future. The past and prediction lengths are set as 15 patches (90s) and 5 patches (30s), respectively. We use the same sampling and split strategy as that in the short- and long-term signal forecasting to generate the data for the frequency-phase forecasting task. A linear layer is adopted to predict the dominant frequency and phase of the future signals. As for the metrics, following the work by Mansouri et al. [34], we use MAE for dominant frequency forecasting, and phase locking value (PLV) for phase forecasting. The PLV is a value between 0 and 1 calculated by the equation:
$$
P L V = \left\| \frac {1}{T} \sum_ {t = 1} ^ {T} \exp (i (\psi (t) - \hat {\psi} (t))) \right\|, \tag {4}
$$
where  $\psi$  and  $\hat{\psi}$  are the original signals and the forecasted signals, respectively.
Imputation. During brain signal recordings, measurement problems such as artifact contamination or electrode impairment are not easily corrected, thus the neural recordings will be incomplete. Imputation can fill in these contaminated signals so that other medical devices can keep operating under missing values, making use of available data [36]. For the imputation task, we randomly mask the timestamps in each patch with the ratio of  $40\%$  and fine-tune the model to predict the missing values. We adopt the same sampling and split strategy to obtain the imputation data as in the short- and long-term signal forecasting. We add a linear head to make predictions, then apply MAE and MSE as the evaluation metrics to measure the discrepancy between the masked and predicted values.
Seizure detection. As one of the most important applications of intracranial recordings, seizure detection task is to evaluate the model ability to distinguish between epileptic seizures and normal waveforms. We sample 250 minutes of recordings from the downstream dataset with about  $10\%$  positive (seizure) samples, where 200 minutes are randomly selected for fine-tuning and the remaining 50 minutes for evaluation. An MLP is adopted to classify the pre-trained representations. The evaluation metrics we use are accuracy, precision, recall,  $F_{1}$  and  $F_{2}$  scores. The F-measure is a metric defined as the weighted harmonic mean of precision and recall, with the following equation
$$
F _ {\beta} = \frac {(1 + \beta^ {2}) \times p r e c i s i o n \times r e c a l l}{\beta^ {2} \times p r e c i s i o n + r e c a l l}.
$$
3.4Baselines
As Brant is a pre-training work on neural signals, we extensively compare our model with the advanced self-supervised or unsupervised pre-training works on all the downstream tasks, including works designed for brain signals: RP [17], TS [17], CPC [17], BENDR [18], MVTS [19] and BrainBERT [20]; and for general time series: CoST [37], TF-C [38], PatchTST [27] and TS-TCC [39]. Considering that seizure detection is an important application scenario of intracranial recordings, we further compare our model with several supervised methods which conduct seizure detection on EEG or iEEG to evaluate the performance of Brant in detecting epilepsy. These baselines contain handcrafted feature based methods, including spectral power [28], rhythmicity spectrogram [3] and amplitude-integrated EEG [40]; and the SOTA deep learning based method on seizure detection, including SEEG-Net [9]. More details of the baselines are shown in App. C.
4.1 Main Results
Fig. 3 summarizes the results of all the downstream tasks. As a foundation model for intracranial recordings, Brant achieves consistent SOTA performance on a variety of tasks compared with other baseline models. We discuss more detailed comparisons of each task in the following paragraphs, where in all the tables we mark values ranking the first (v), second (v) and third (*v) in each column.
Figure 3: Performance comparison of our model and other baseline models on all downstream tasks. Model performances are plotted on three radar subfigures for clarity with the same coordinate range.
Table 1: Performance on forecasting tasks.
Model\Task	Short-term signal forecasting	Long-term signal forecasting	Frequency-phase forecasting
MAE	MSE	MAE	MSE	ph. PLV	freq. MAE
RP [17]	0.7569±0.0016	1.1769±0.0103	0.6774±0.0017	1.0579±0.0010	0.4929±0.0002	0.4197±0.0042
TS [17]	0.6142±0.0012	0.8771±0.0076	0.6286±0.0014	0.9771±0.0011	0.5609±0.0002	0.0245±0.0008
CPC [17]	0.6318±0.0012	0.8886±0.0088	0.6324±0.0013	0.9604±0.0010	0.5608±0.0003	0.1036±0.0014
BENDR [18]	0.6002±0.0014	0.8720±0.0081	0.5948±0.0020	0.8535±0.0013	0.5685±0.0002	0.0604±0.0017
MVTS [19]	0.6868±0.0019	1.1859±0.0102	0.5867±0.0026	0.8381±0.0014	*0.5843±0.0004	0.0320±0.0014
BrainBERT [20]	0.6161±0.0010	0.8857±0.0094	0.7214±0.0014	1.4727±0.0012	0.5712±0.0002	0.1065±0.0015
PatchTST [27]	0.5050±0.0010	0.6482±0.0062	*0.5792±0.0014	*0.8235±0.0006	0.5905±0.0002	0.0267±0.0006
TS-TCC [39]	0.6042±0.0012	1.0244±0.0096	0.5938±0.0018	0.8506±0.0008	0.5743±0.0002	0.0061±0.0003
TF-C [38]	1.2530±0.0025	2.9441±0.0146	1.0683±0.0030	2.2491±0.0016	0.5493±0.0004	0.1593±0.0029
CoST [37]	0.5531±0.0011	1.0922±0.0092	0.5955±0.0012	0.8546±0.0005	0.5728±0.0003	0.1152±0.0018
Brant-Freeze	*0.5352±0.0008	*0.7161±0.0111	0.5780±0.0004	0.7869±0.0007	0.5777±0.0002	*0.0113±0.0003
Brant	0.4007±0.0014	0.4626±0.0048	0.5693±0.0012	0.7676±0.0005	0.6004±0.0001	0.0041±0.0002
Short- and long-term signal forecasting. The results of short- and long-term signal forecasting are shown on the left and middle of Tab. 1, respectively. Not only does Brant achieve SOTA performance,
but the results of Brant-Freeze are also better than most baselines. Especially on long-term forecasting, Brant-Freeze defeats all the baseline methods, showing the ability to capture long-term dependency of our model. Among the baselines, PatchTST [27] performs well compared to other methods, mainly because it also adopts a patching strategy to attend longer temporal dependency.
Frequency-phase forecasting. Forecasting physical features like frequency and phase of neural signal is essential in some medical techniques (see Sec. 3.3). The results of frequency-phase forecasting are shown on the right of Tab. 1, which shows that our model outperforms all the methods, demonstrating that Brant facilitates effective implementation of therapies for some brain disorders.