Beatrix: Out-of-Distribution Generalization of Large EEG Model via Invariant Contrastive Fine-Tuning

We thank the reviewer for the helpful feedback.

“Lack of Novelty”

Our work introduces significant novelties that set it apart from prior approaches. We combine spectral tokenization with generative pre-training and invariant contrastive fine-tuning (CIFT) to improve domain-invariant embeddings and OoD performance. We conducted extensive experiments across multiple datasets to demonstrate robust OoD performance. We propose novel approaches that are likely to influence future research in EEG data analysis and biomedical signal processing.We curated an extensive list of open EEG data corpora (arguably the largest one to our knowledge) for pretraining, which enhances the reproducibility of this work and will facilitate future research on EEG foundation models.

We would like to highlight that all previous works were cited in the initial version across different sections of the main text. Bottleneck and low-rank adapters are compatible: the bottleneck adapter operates following the main network module, while the low-rank adapter is integrated within it. We refined these in the revised version.

“Factual Errors”

1. The term "low-rankness" refers to parameter-efficient adapters, as clarified in the main text. This phrasing is refined in revised version.

2. We clarify that we use SwiGLU throughout the work as our Feed-Forward component, which is explicitly denoted in Figure 3, see main text discussion in Sec. 3.2

3. We acknowledge the issue with the unclear arrows and rectified it by ensuring a clear and closed connection between the two lines in the revised version.

Experiments

We structured the paper to focus the main text on key contributions and results while placing detailed analyses in the appendix due to space constraints. We would like to highlight that we already ablated it in Table 1, with ERM method (which means no additional OoD auxiliary losses are used except for the cross-entropy loss), i.e., w/ and w/o CIFT loss. We also compared CIFT with other contrastive OoD generalization algorithms.

Representation Problem

1. This will be revised in the final version.

2. We clearly pointed out in Sec. 2.1 that e represents an environment or domain.

3. Semi-causal refers to a hybrid pretraining framework that combines causal (unidirectional, using past tokens to predict future ones) and non-causal (bidirectional, using both past and future tokens for reconstruction) spans. Unlike purely causal models, semi-causal modeling enables both extrapolation and interpolation across temporal and spectral domains. Previous works only use non-causal approach. Our approach balances temporal consistency with broader contextual understanding, enhancing the model’s ability to generalize to complex EEG dynamics.

Other details have been refined in the revised version.

We thank the reviewer for the helpful feedback.

Quality of the Experiments

1. The experiments are conducted following standard practices in EEG foundation model research, where we report the mean performance across multiple test sets. Some baseline pretrained checkpoints were unavailable due to intellectual property and privacy concerns, which restricted our ability to conduct statistical tests.

2. We clarify that full-parameter and linear probing results were presented in Table 1 as ERM-Full and ERM-LP, respectively. Moreover, choice of rank has significant influence on the performance. Given limited computational budget, larger rank leads to better performance. Ablation of different adapters in CIFT was presented in the appendix.

3. Regarding the visual discrepancy between the performance reported in Table 1 and Figure 5, it is not due to differences in the underlying data, but rather an artifact in the visual representation in Figure 5, which arose from the way the graph was rendered in Python using Matplotlib.

4. We have refined the details in the revised version.

Clarity of CIFT

We have addressed and clarified the ambiguities in the CIFT section in the revised version, providing specific clarifications in response to the reviewer’s points as follows.

1. We clarify that the [CLS] token is not used in CIFT or any other OoD generalization algorithms presented in this work. This has been corrected in the revised version.

2. $\rho$ is defined as an auxiliary function for estimating virtual environment labels for each sample. It is not used during the inference stage of CIFT. The additional parameters $\eta$, introduced by $\rho$, are implemented via a two-layer MLP, as already noted in the main text.

3. We highlight that the dimensions for the contrastive fine-tuning loss were provided in the corresponding subsection.

4. We clarify that $\rho$ is solely used to estimate virtual environments $\hat{e}$, not class labels $\hat{y}$. Furthermore, the number of virtual environments $K$ is unrelated to the number of classes, and this distinction is now clearly articulated.

Questions

Introduction

We have polished the introduction section to avoid ambiguities.

Preliminaries

1. It is defined as a sample and its label belonging to environment $e$ within a potential set of environments.

2. Relevant works were cited in the OoD baselines section, as well as additional related work section in the appendix in the initial version. Additional citations have been included in the revised version for completeness.

Methods

1. We clarify that for a 4-second EEG sample at a sampling rate of 256 Hz, raw wavelet features are downsampled to $64 \times 256$, with a patch size of $8 \times 8$ adopted.

2. The issue has been corrected in the revised version.

3. We highlight that $\hat{e}$ represents the predicted virtual environment label in softmax probability, with dimensionality $K$.

Experiments

1. The ambiguity in the figure was due to an error during its creation. This has been corrected in the revised version.

2. The experiments were conducted as an average over multiple subject-independent folds, following dataset splits used in previous works and adhering to multi-fold cross-validation. However, we note that similar practices were not mandatory in all prior studies.

3. We clarify that experiments without a temporal domain tokenizer, such as InfoNCE and HSIC in Table 1, were included. However, the CIFT method inherently relies on this tokenizer, making it inseparable.

4. Numbers and classes were provided in the initial version based on prior works, with clear specifications and citations. Further refinements have been added in the revised version.

5. "SF" specifically refers to seizure forecasting. Table 1 reports OoD performances for both seizure forecasting and detection tasks.

Wording, Grammar and Organization

We have refined the wording, grammar, and overall organization in the revised version. Additionally: SEEG refers to stereo-electroencephalography. PEFT stands for "parameter-efficient fine-tuning." ERM means "empirical risk minimization," which involves training without any OoD generalization algorithms.

We thank the reviewer for the helpful feedback.

Weaknesses

(1) We clarify that we also discussed performance on sleep staging, motor imagery and brain decoding tasks in the main text, with detailed results in the appendix.

Questions

(1) The Temporal Tokenizer is implemented using a 1D CNN, trained during the fine-tuning stage as part of the CIFT method, which is the primary OOD approach proposed in this work. Spectral tokenization, on the other hand, is utilized throughout both pretraining and fine-tuning stages, employing a frozen spectral tokenizer implemented via a VQGAN. We included ablation experiments in the appendix, which demonstrate significant performance improvements when both tokenizers are used together, as is done in CIFT.

(2) In this work, our primary focus is not on addressing discrepancies between different OOD tasks. Instead, we aim to propose a unified solution that demonstrates improved performance across distinct OOD tasks, emphasizing generalizability and robustness.

(3) High-frequency activities, such as gamma-band and high-frequency oscillations (HFOs) above 50 Hz, are indeed crucial for understanding cognitive functions and seizure development. Our approach does not impose specific frequency or time resolution constraints on the input data. For consistency, we adopted a 4-second segment length with a sampling rate of 256 Hz as the default configuration throughout our experiments.

We thank the reviewer for the helpful feedback.

Weaknesses

(1) While we agree that presentation could be further refined, we emphasize that the main contributions are clearly highlighted in the Introduction section of the manuscript.

(2) We provided a comprehensive review of prior works in the Related Work section, with additional details included in Appendix F for further context.

(3) A significant portion of our experiments was conducted on publicly available datasets, including seizure detection, seizure forecasting, sleep staging, motor imagery, and brain decoding. Detailed results are provided in the Appendix. For the self-collected datasets, we have made them accessible via an anonymous link. Ablation studies are also thoroughly presented in the Appendix.

(4) We clarify that we compared our results with several state-of-the-art foundation models in Tables 3, 4, 9, 10, and 11, demonstrating the competitiveness and robustness of our approach.

Questions

(1) We clarify that the spectrotemporal integration of EEG representation is not only distinct but also critical, as it underpins the effectiveness of our CIFT method. By leveraging both spectral and temporal tokenization, our approach enables robust domain-invariant representations essential for EEG's non-stationary nature. Widespread experiments have shown that CIFT achieved substantial performance gains by integrating these features in an environment-aware manner, highlighting its novelty and transformative impact across tasks.

(2) The analytic wavelet transform is a well-established technique widely discussed in college-level signal processing textbooks, and relevant references have been included in the revised version.