Technical Summary: TFM-Tokenizer

Problem Statement

While foundation models have revolutionized Natural Language Processing (NLP) and are increasingly applied to Electroencephalogram (EEG) analysis, a critical bottleneck remains: effective EEG tokenization. Existing EEG foundation models (e.g., BIOT, LaBraM) typically tokenize signals by directly segmenting continuous EEG into short-duration patches without learning a discrete vocabulary. This approach merely discretizes the signal, failing to capture statistically grounded, data-driven representations. Furthermore, previous attempts at neural tokenizers (e.g., in LaBraM) often serve primarily as training objectives that are discarded during downstream inference, limiting their reusability and the ability to improve input representation quality.

The authors identify three specific challenges in designing an effective EEG tokenizer:

1. Tokenization Target: Real-world EEG varies widely in device, channel configuration, and recording length. Tokenizers should operate at the single-channel level to learn channel-agnostic tokens that generalize to multi-channel and non-standard devices (e.g., ear-EEG).
2. Token Resolution: EEG signals contain diverse oscillatory (e.g., alpha, beta) and transient patterns (e.g., spikes). Effective tokens must represent these underlying time-frequency motifs, which are often distorted by noise and temporal warping.
3. Learning Objective: Relying solely on time-based motifs risks overlooking spectral structure. The learning objective must explicitly incorporate time-frequency representations to capture band-specific and cross-frequency patterns.

Methodology: TFM-Tokenizer

The paper proposes TFM-Tokenizer, a model-agnostic framework that learns a vocabulary of time-frequency motifs from single-channel EEG signals and encodes them into discrete tokens. The framework operates in two phases:

1. TFM-Tokenizer Pretraining (Unsupervised, Single-Channel)

The tokenizer is trained to map continuous single-channel EEG signals into a sequence of discrete tokens. It employs a dual-path encoding design:

• Localized Spectral Window Encoder: Instead of projecting entire spectral windows through a single linear layer, this component divides the spectrogram into frequency patches. It uses a Frequency Patch Encoder followed by a Frequency Transformer to model intra-window cross-frequency dependencies. Crucially, it employs a Gated Patchwise Aggregation mechanism to emphasize relevant frequency bands (e.g., up to 32 Hz for sleep) and suppress irrelevant ones.
• Temporal Encoder: Processes raw EEG patches to capture temporal context.
• Fusion and Quantization: The aggregated frequency embeddings are concatenated with temporal embeddings and processed by a Temporal Transformer to model long-range dependencies. The resulting embeddings are mapped to a discrete vocabulary using Vector Quantization (VQ) based on a learnable codebook, producing motif tokens.
• Training Objective: The model is trained using a joint frequency–temporal masking strategy. Both frequency bands and temporal segments are masked, and the model must reconstruct the masked spectrogram and predict the discrete token vocabulary. This forces the tokenizer to learn distinct, localized time-frequency patterns.

2. Downstream Transformer Training

Once the tokenizer is trained, it is frozen. A lightweight Transformer (using linear attention) serves as the downstream model.

• Input: Multi-channel EEG is processed by the tokenizer independently for each channel to generate token sequences.
• Processing: These token sequences are flattened, combined with channel and position embeddings, and processed by the Transformer.
• Pretraining: The downstream model undergoes Masked Token Prediction (MTP) pretraining to learn intra- and cross-channel dependencies.
• Fine-tuning: The model is fine-tuned for specific downstream tasks (e.g., seizure detection, sleep staging).

Key Contributions

1. Formulation of Single-Channel EEG Tokenization: The authors are the first to investigate learning a discrete token vocabulary specifically for single-channel EEG signals that captures time-frequency motifs and is directly reusable as input for downstream modeling.
2. Novel TFM-Token Framework: The introduction of a framework that transforms EEG into discrete token sequences via a dual-path (time-frequency) encoder. This framework is model-agnostic, seamlessly integrating with existing foundation models (e.g., BIOT, LaBraM) to enhance their performance.
3. Broad Evaluation and Scalability: The method is validated across four diverse EEG datasets (TUEV, TUAB, CHB-MIT, IIIC Seizure) and demonstrated to be scalable to non-standard devices, specifically ear-EEG for sleep staging, where it outperformed baselines significantly.

Experimental Results

The paper reports consistent performance gains across multiple benchmarks:

• Accuracy: On the TUEV dataset (event classification), TFM-Tokenizer achieved up to an 11% improvement in Cohen's Kappa (0.5588 to 0.6189) over strong baselines in a multi-dataset pretraining setting. On the IIIC Seizure dataset, it improved Cohen's Kappa by 36% over LaBraM in the multi-dataset setting.
• Efficiency: The TFM-Tokenizer paired with a lightweight downstream model (approx. 1.9M parameters) outperformed larger baselines (e.g., LaBraM at 5.8M, BIOT at 3.2M), demonstrating that tokenization reduces data complexity effectively.
• Generalization: As a plug-and-play component, integrating TFM-Tokenizer into existing foundation models (BIOT and LaBraM) consistently boosted their performance across single- and multi-dataset settings. For instance, it improved LaBraM's performance on the CHB-MIT dataset by 147% in AUC-PR in the single-dataset setting.
• Scalability (Cross-Device): In an out-of-domain experiment on the EESM23 ear-EEG sleep staging dataset (using electrodes outside the standard 10–20 system), TFM-Tokenizer outperformed baselines by 14%, validating its ability to generalize across different recording devices and configurations.
• Token Quality: Analysis revealed that the learned tokens exhibit strong class-discriminative uniqueness, consistency within classes, and the ability to capture meaningful physiological motifs (e.g., PLEDs and GPEDs) despite noise and amplitude variations.

Significance and Claims

The paper claims that TFM-Tokenizer addresses a fundamental gap in EEG foundation models by providing a principled, data-driven method for tokenization that moves beyond simple signal segmentation. Its significance lies in three areas:

1. Improved Representation Quality: By explicitly modeling time-frequency motifs, the tokenizer provides higher-quality input representations that are more informative and interpretable than raw continuous segments.
2. Plug-and-Play Versatility: The framework is not tied to a specific architecture, allowing it to enhance a wide range of existing foundation models without requiring their retraining from scratch.
3. Device Agnosticism: By operating at the single-channel level, the method decouples tokenization from strict channel configurations (like the 10–20 system), enabling the application of foundation models to emerging, non-standard EEG devices such as ear-EEG sensors.

The authors conclude that this work lays the groundwork for more robust, scalable, and generalizable EEG foundation models capable of adapting to diverse modalities, devices, and clinical tasks.