BELT-2: Bootstrapping EEG-to-Language representation alignment for multi-task brain decoding

Q 1:

The training details of the Q-Conformer is only provided in the appendix.

R 1:

We put training details in the appendix mainly due to the page limitation. We will reorganize and put the details into the main paper in the final version.

Q 2:

Question about vocabulary size.

R 2:

The vocabulary in the training split of the ZuCo dataset has 2603 unique words. However, in the final translation result in the the decoding phase, we use the vocabulary of the LLM (T5) which is 32128.

Q 3:

Not clear what loss function is used to train the multi-task query and which part of the network is updated.

R 3:

In the multi-task training experiment, we train 3 tasks simultaneously by randomly sampling tasks in each training step. We also found the best performance came from updating the whole Q-Conformer encoder, including the discrete conformer, the query prompt, and the context transformer.

Q 4:

Question about the BPE method.

R 4:

We provide an additional ablation experiment involving using word-level contrastive learning with a discrete Conformer model involving three different settings.

Results are detailed in the table. The use of word-level improves the BLEU scores to an extent while the BPE-level contrastive learning leads to further improvements for all BLEU scores. This rationalizes the use of BPE-level contrastive learning over word-level contrastive learning.

Q 5:

When combining Q-Conformer with LLM using continuous virtual prompts, how is it trained exactly?

R 5:

In Section 2.4 and the introduction, we briefly introduced the training approach involving the frozen Q-Conformer and the pre-trained Language Model (LLM) decoder using the prefix-tuning method. While this method is commonly employed in fine-tuning LLMs.

The essence of the prefix-tuning is to only tune the prefix tokens to the LLM so that the trained prefix contains enough instruction for the LLM to understand the task [1].

Since the prefix-tuning method is not a primary contribution introduced by our work, we did not delve into extensive details in the main paper.

[1] Prefix-tuning: Optimizing continuous prompts for generation.

Q 6:

Question about the training of the Q-Conformer.

R 6:

We conducted additional ablation experiments were conducted on the Q-Conformer encoder to show performances when trained with or without the VQ and contrastive components. The results indicate that VQ and BPE-Contrastive could improve the model's performance together. We will add the ablation in the final version of the paper.

Q 7:

Question about the experimental setting of word-level modeling.

R 7:

There are two major reasons for using word-level tokens as input in conducting the sentiment classification task.

The first reason is that we strictly follow the experimental setup in the baseline method for a fair comparison.

Another reason is that we want to demonstrate the multi-task capacity of our model. This means that given the same input, our model is able to decode different information from the input EEG tokens (i.e., the translation, summary, and sentiment).

Q 8:

Question about word-level modeling.

R 8:

We add ablation experiments on different contrastive learning levels to answer this question as shown below.

[Translation Task]

We observe that the use of sequence-to-sequence modeling did not yield a significant improvement when compared to word-level modeling.

Q 9:

Question about Speculative Augmentation.

R 9:

The Speculative Augmentation does not incur extra memory requirements. we employ a cache system by saving copies of the intermediate layer output (last output tokens) from the training set at each epoch. The cached outputs from a total of 15 training epochs, are utilized to diversify the training data spectrum during the prefix-tuning process. Thus, we use these cached outputs instead of loading checkpoints.

Extended response to Weakness a:

To further address the problem of the clarity of the training process, we have improved the readability and organized the illustration of the learning process. The training and loss function of Q-Conformer is clarified in Section 2.2 where we summarize the EEG-to-langauge alignment learning. In this stage, we jointly optimize 3 objectives, BPE-level contrastive learning (BPE-CL), Nagative contrastive learning (NCL), and the EEG-to-language matching (ELM) with the loss function clearly presented in the section. The first two objectives are applied to the discrete EEG tokens to improve the semantic (EEG-language alignment) while keeping the distinction among different tokens. The third objective is for training a task-specific query prompt as well as training the context-transformer. For multi-tasking, we provided a clearer description in Sectopm 2.4 and Figure 5. Please refer to the current paper for more details.

Dear Esteemed Reviewer,

Thank you immensely for your invaluable insights and thoughtful queries that have significantly contributed to enhancing our work. Your feedback has been meticulously incorporated into our paper, elevating its quality and depth.

We earnestly appeal to your esteemed judgment once again, considering the substantial refinements made, to reconsider the rating for our paper. Our relentless dedication has been aimed at pushing the boundaries within the realm of EEG research, aspiring to pave the way for broader applications and advancements.

We hope our work can help improve brain research and make a positive difference in the wider research community. Your support and revised evaluation would not only recognize our efforts but also amplify the potential impact of this research, steering us collectively toward greater progress.

With sincere regards, The Authors

Question 1:

Quantization process of the EEG token is unclear. Whether the spatial arrangement of the EEG sensors played any roles in the process?

Response:

We use the public ZuCo dataset in our experiment. The ZuCo dataset segments the raw EEG wave using the eye fixation for each word and converts each segmented EEG wave into a total of 8 frequency bands.

Originally, they used a 128-channel EEG cap to collect the EEG signal, but they removed 23 channels which resulted in a total of 105 EEG channels. As a result, each word-level EEG segment is processed into a 105*8=840 size frequency embedding. The maximum number of words in a sentence is set to 56. We therefore use the 56-word-level EEG embeddings as input tokens for the Conformer network.

The Conformer network is a transformer-like architecture that outputs the same number of tokens as the input. Afterward, the output tokens from the Conformer, denoted by $h$ in our paper, will be quantized by a vector quantizer, and each token from the Conformer will be converted to a discrete token $b$. These discrete tokens will be used as input to the context transformer where the query prompt will interact with them via a cross-attention layer.

Finally, the output of the context transformer (i.e., the mediate layer coding) can be sent to the LLM to get the final translation outputs.

For the spatial arrangement of the EEG sensors. Since we are using a public dataset, the arrangement of the 105 EEG channel is fixed. Therefore, the effect of different sensor arrangements is not the research focus of this paper. However, we think it would be an interesting topic in the future works.

Question 2:

Question about the difference between BELT-1 and BELT-2.

Response:

Thank you for your feedback. We have briefly mentioned the different between BELT-1 and our BELT-2 model. We will provide more details in the final version. The limited discussion and comparison between the BELT-2 and BELT-1 models in our paper were primarily constrained by page limitations. The key design improvements in BELT-2 over BELT-1 can be summarized as : 1) Enhanced Alignment, 2) Muti-Task Flexibility, and 3) Integration with Pretrained LLMs.

Question 3:

Questions about the 'EEG prompots' and the Mediate layer coding.

Response:

We used a BART model to initialize the context transformer. This makes the context transformer a bidirectional model that captures global information from the EEG tokens. However, participants read through a sentence according to the order they prefer and will also skip some words during the reading. As a result, we utilize the mediate layer coding as a 'prompt' to LLMs and use the finetuned prefix to instruct the LLMs to rearrange the disorder language information and generate more natural translation sentences.

Last but not least, we will add illustrative information about the mediate layer coding and the EEG prompts to the figures to provide a clearer idea of the mediate layer coding.

Question 4:

Question about the dataset split and subject variations.

Response:

Thanks for the question. Currently, the experimental setup splits the dataset on a cross-training sample basis for a fair comparison with other methods. In response to the cross-subject splitting in the question and to explore whether inter-individual variations will impact the EEG-to-Language translation performance of our model, we additionally conducted the experiment in a cross-subject setting (shown in the table below). In this experiment, we leave one subject out as a test set and train the model on all other subjects. The performance is illustrated in the following table:

The table shows that the subject variant on the ZuCo dataset doesn’t have a significant impact on the EEG-to-Language translation performance of our model. We will include this ablation experiment in the supplementary in the final version.

Extended response to Question 3:

In the latest version of the paper, we have improved the readability and provided a clearer explanation of the mediate layer coding. For clarity, we rename the output of the Q-Conformer as the Mid-Layer Coding as these output embeddings represent a crucial midpoint between the EEG Encoder and LLM. Please refer to Section 2.3 and Figure in the current version of the paper for more details.

Extended response to Question 4:

In the latest version of the paper, we have included the cross-subject experiment in the main paper (Page 7 last paragraph, and Figure 6). Please refer to the current paper for more details.

Question 1:

Unclear whether BPE-level contrastive learning is a novel method, or novel in the EEG domain.

Response:

To the best of our knowledge, there is no prior work employing BPE-level contrastive learning in a manner similar to our approach.

We conducted ablation experiments on the BPE in Table 1 in the main paper. We move the comparison for the BPE here for your convenience. Results in the table below show that the introduction of BPE-level contrastive learning brings significant improvement to the BLEU scores, reaching 43.06(+1.49), 25.57(+1.55), 15.16(+1.36), 9.17(+0.57). We will improve readability and refine the paper in the final version.

Previous works on visual-language alignment are employing sentence-level contrastive learning [1,2]. However, due to the scarcity of EEG-language pairs, we need to incorporate stronger guidance from the language modality during training. Consequently, we consider that the use of BPE-level contrastive learning constitutes a novel method for aligning a modality with language when faced with a limited number of training sample pairs.

[1] Image as a Foreign Language: BEiT Pretraining for Vision and Vision-Language Tasks.

[2] Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models.

Question 2:

Question about the "general bridging between brain signals and languages.". The main difference between BELT and BELT-2?

Response:

Here, "general bridging between brain signals and language" means this work is the first to achieve multi-task EEG decoding and is able to leverage the generative capacity of the existing LLMs.

This includes being able to enhance the general understanding of brain signals by learning from different tasks while being able to decode multiple information from EEG signals (e.g., sentiment or summary). Our work signatures an important step to link the brain signal into existing multimodal understanding models which has the potential for boarder applications.

The main difference between BELT-1 and BELT-2 is that BELT-1 is a single-task architecture while BELT-2 is a multi-task architecture. Our model accomplishes multi-task decoding using query prompts, a strategy aligned with current Language Models (LLMs). Also, BELT-2 allows us to directly combine various LLM (e.g., T5 and LLAMA2, PEGASUS) with our EEG Encoder to enhance the translation quality by merely tuning the prefix tokens. We are also the first work to provide a study (Table 5) of combining EEG encoders with different state-of-the-art LLMs.

Question 3:

Are raw EEG wave segments of uniform length? What is the use of eye-tracking information?

Response:

In our experimental setup, we leverage the ZuCo dataset [1]. Notably, this dataset incorporates eye-tracking data, recording eye fixations on individual words. The fixation duration for each word varies significantly, and some words may be omitted during reading, leading to EEG segments of different lengths.

Although the length for each word-corresponding EEG signal is different, the ZuCo dataset performs a frequency transformation operation to extract a fixed number of frequency values from each EEG segment for each channel.

[1] ZuCo, a simultaneous EEG and eye-tracking resource for natural sentence reading.

Question 4:

Question about "we could easily extend the Q-conformer to a new downstream task ... "

Response:

Here, we don’t need to retrain a new model or the query tokens. For the multi-task setting of our Q-conformer model, we trained 3 tasks together and implemented a sampling strategy to select a task for each training batch.

Question 5:

In section 2.4, it says that "During the EEG representation learning stage, the Q-Conformer extracts, task-specific information from the EEG input signals." But this doesn't seem to be the case? Is Figure 2 left meant to depict the representation learning stage? If so, why aren't there any task-specific terms in the objective function?

Response:

Figure 2 is meant to depict the feature extraction process of the Q-Conformer. In Figure 2 (left), our primary intention was to illustrate that the query prompts will query task-specific information from the EEG tokens via a cross-attention layer in the Context Transformer.

As mentioned in the supplementary, the query prompt and the Q-Conformer are trained with the objective function mentioned in the appendix. We will make refinement to the figure in the final version.

Extened Response to Question 5:

In the latest revision, we have improved the readability of the paper. In particular, we revised the original Figure to give a more precise and clearer illustration of the Q-Conformer architecture. Please refer to the Figure 2 in the current paper. In Figure 2, we show the overall structure of the Q-Conformer. It consists of a discrete conformer, a context transformer (C-Former), and a query prompt. The input EEG embeddings (EEG embed) are first processed by the conformer into continuous EEG tokens. A vector quantizer is then used to discretize the EEG tokens. Then, a query prompt interacts with the discrete EEG token via the cross-attention layer from in the C-Former to extract task-specific context information from the discrete EEG tokens. While Figure 2 mainly depicts the interaction of submodules in the Q-Conformer, Figure 3 and Section 2.2 gives a more specific illustration of how the query prompt is trained. We summarize the first training stage as the EEG-to-language alignment learning stage where we use the EEG-to-language matching (ELM) objective as the 'base' task in this alignment learning stage. When training on ELM, the query prompt is trained to become task-specific. For multi-task training, as depicted in Figure 5, we assign a query prompt to each decoding task, so that each is trained by the objective of a task (e.g., translation, summary, or sentiment classification). Please refer to the current paper for more details and consider raising the rating for our work.

Question 1:

The interaction of all the components in Equation 2 and hpyerparameters is unclear.

Response

We combine the loss terms in Equation 2 using a weighted combination as mentioned in the supplementary (page 13 to 14). The combination of these loss terms is a common operation for training a Vector Quantised (VQ) encoder. Specifically, $L_{cm}$, and $L_{cb}$ are the default losses for training the discrete codebook as illustrated in VQVAE paper [1]. The $L_{div}$ is a term to encourage diversity in the usage of the discrete codebook entries which is also frequently used in training a VQ encoder, an example of using the diversity loss is in the wav2vec2.0 paper [2]. $L_{recon}$ is the necessary term used for training a reconstructive encoder as in the latent diffusion model [3]. We also provided ablation experiments on different hyper-parameter settings in the supplementary (Page 17). We will organize these information from the supplementary to the main paper to improve the readability for boarder audiences.

Question 2:

Not clear how the query prompt is trained and whether a new task requires complete retraining.

Response

The Q-Conformer model does not require a complete retraining. We borrow the idea of query prompts from the recent visual-language model BLIP-2. The query prompt is able to select and extract information from the EEG tokens through a cross-attention layer in the context transformer.

As mentioned in Page 4, the Q-Conformer efficiently adapts to new downstream tasks by initializing a set of query tokens, enabling the extraction of task-specific information. This adaptation is achieved through training a query prompt using the objective function tailored to the specific task. Given that BELT-2 is a multi-task model, individual tasks are trained concurrently without the need for complete model retraining.

Question 3:

How the Frequency domain EEG embedding is transformed into continuous EEG tokens.

Response

As mentioned in the paper (Page 3), EEG data from the ZuCo dataset is first segmented according to the gaze fixation on each word using eye-tracker information. A frequency domain transformation is applied to each EEG segment to extract value from 8 frequency bands. The ZuCo dataset contains 105 EEG channels. These frequency bands from the 105 channels form an equal-length word-level EEG embedding (size=840).

Similar to previous works [4], the word-level EEG embedding is processed by a transformer-like EEG encoder that processes the word-level EEG embedding. We use the Conformer model to process the word-level EEG embeddings. After processing by the Conformer model, a vector quantizer to discretize the output EEG tokens into discrete tokens.

Question 4:

Lack of Specificity for Section 2.3.

Response

The reason why we adopt the negative contrastive loss is that the difference in human brain waves when looking at different words is very small, increasing the difficulties of EEG decoding. Therefore, we use the negative contrastive loss to enlarge the differences among EEG tokens and consequently enhance the final translation performance. The method of using negative contrastive loss among encoded tokens is also employed in the training of speech encoders [2].

Question 5:

Unclear Ablation Experiments: The paper does not provide a clear description of the ablation experiments, particularly the absence of ablation on BPE-level Contrastive learning and key components of Formula 2.

Response

We provided ablations about the BPE-level contrastive learning (Table 1), about using pretrained EEG encoder for downstream tasks (Table 3 and Table 4), about connecting the EEG encoder with different LLMs (Table 5), and about hyper parameters (Figure 9) in our main paper and in the supplementary. Due to the page limit, we moved the rest ablation experiments to the supplementary. We will reorganize more ablation information in the final version of the paper.

Regarding to the effect of BPE-level contrastive learning, we have already provided this ablation study in Table 1.

This ablation shows the performance of our model with or without the BPE-level Contrastive learning term. Results show that the introduction of BPE-level contrastive learning brings significant improvement to the BLEU scores, reaching 43.06(+1.49), 25.57(+1.55), 15.16(+1.36), 9.17(+0.57).

[1] Neural discrete representation learning.

[2] wav2vec 2.0: A framework for self-supervised learning of speech representations.

[3] High-resolution image synthesis with latent diffusion models.

[4] Open vocabulary electroencephalography-to-text decoding and zero-shot sentiment classification.

Question 1:

Difference between BELT and BELT-2.

Response 1:

Thanks for your question. The original paper mentioned the differences between these two models on page 2. Key design improvements in BELT-2 over BELT-1 can be listed as follows:

1. Enhanced Alignment: BELT-1 applies contrastive learning on word and sequence levels to bootstrap EEG encoder training. It has significant shortcomings: 1. Sequence-level alignment is compared coarser and would increase the difficulties for learning the exact relationship between EEG tokens and the corresponding word. 2. Word-level alignment would lead to a decreased performance in the decoding of rare words or out-of-training-vocabulary words when training data size or training vocabulary is limited.

   BELT-2 introduces a significant improvement through BPE-level contrastive learning. It operates at a more granular level than word-level alignment as it breaks down words into smaller size of root words. It is especially useful for handling rare words or words out of the vocabulary of the training set because of the shared subword units. Experimental results in Table 2 improvement lead to a clear improvement in the open-vocabulary setting.

   This discussion is compressed due to the page limit. We will improve the writing to emphasize this point in the final version.

2. Muti-Task Flexibility: BELT-1 follows the same setting of all previous works that only explore brain-to-text translation, a single task. BELT-2 makes a significant improvement and extends the brain decoding to the multi-task setting for the first time. The multi-task setting makes the decoding supervision denser and also enhances the learning process. Additionally, it explores seamlessly switching between tasks by employing different query prompts, reducing computational resource requirements, and utilizing knowledge from diverse tasks.

3. Integration with Pretrained LLMs: BELT-2 offers the ability to connect the EEG encoder to a pre-trained Language Model (LLM), such as T5 or LLAMA2. This work provides a detailed comparison of different LLMs. It also includes an intuitive comparison between single-direction and bi-direction LLMs for follow-up researchers. This integration allows BELT-2 to leverage the strong generative capacity of pre-trained LLMs, providing enhanced performance.

Question 2:

More precision on the training details and the loss function.

Response 2:

Sorry for any inconvenience caused, due to the page limit, we moved some training details and the loss function to the appendix.

For the training of the EEG encoder as illustrated in the appendix, i.e., the Q-Conformer, we used a weighted combination of the VQ loss, BPE-level contrastive loss, and the negative contrastive loss. For multi-task training, we trained all tasks together using a random sampling method on tasks. To bridge the trained Q-Conformer with the LLM decoder, we froze both the Q-Conformer and the LLM decoder and used the prefix-tuning method to train a set of prefix prompts to 'instruct' the LLM to decode fluent sentences from the output tokens of the Q-Conformer model. For the loss function used to train the translation task and the summarization task, we used a common machine translation loss, as in the BART paper [1]. For the sentiment classification task, we only used cross-entropy loss to train the last output token of the Q-Conformer, similar to the method in the XLNet paper [2].

We will further concise the writing to include more information in the main paper.

[1] Lewis, Mike, et al. "Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension." arXiv preprint arXiv:1910.13461 (2019).

[2] Yang, Zhilin, et al. "Xlnet: Generalized autoregressive pretraining for language understanding." Advances in neural information processing systems 32 (2019).