Paper: BELT-2: Bootstrapping EEG-to-Language representation alignment for multi-task brain decoding (ICLR-2024)

Reviewer comment (Reproducibility):
Many of the training details are unclear, which makes it very hard to understand the working mechanism of the proposed BELT-2. I will elaborate in the question section.

Author response:
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.

Reviewer comment (Experiments):
The training of the Q-Conformer is based on seven different loss terms λ1(L_cm + L_cb) + λ2 * L_recon + λ3 * L_div + λ4 * L_bpe + λ5 L_neg + λ6 * L_tr with a list of balancing coefficients of [1, 1, 0.001, 10, 0.001, 10]. Training a discrete encoder (Conformer in this case) like VQ-VAE along with the downstream (translation in this case) is not common. From my experience, balancing the reconstruction quality and the code quality itself while training the VQVAE is already a hard task. Thus, I would very much like to know how this model would perform without the quantization and reconstruction. That is to remove L_cm, L_cb, L_recon, and L_div and only use the Conformer as an EEG encoder. Ideally, an ablation study on these loss terms and corresponding curves during training would help the readers understand the different components of BELT-2.

Author response:
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.

| VQ  | BPE-Contrastive | Bleu-1 | Bleu-2 | Bleu-3 | Bleu-4 |
| --- | ---------- | ------ | ------ | ------ | ------ |
| No  | No         | 40.12  | 23.18  | 12.61  | 6.8    |
| No  | Yes        | 41.9   | 24.57  | 14.2   | 8.278  |
| Yes | No         | 41.57  | 24.02  | 13.80  | 8.06   |
| Yes | Yes        | 43.06  | 25.57  | 15.16  | 9.17   |

Reviewer comment (Experiments):
The experimental setting of word-level modeling is questionable. For a regular sentiment classification task, a continuous EEG signal is used as input, yet this work only extracts EEG segments based on eye tracking. Naively speaking, sentiment-related information does not necessarily exist only when the subject is looking at words. Or else, there might be a delay of the reaction in the brain upon reading. Is this considered while aligning the word and EEG?

Author response:
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).

Reviewer comment (Experiments):
The word-level modeling seems too artificial to fit the word embedding format to LLM. Sequence-to-sequence modeling would be more interesting and practical.

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

[Translation Task]
| Contrastive learning | BLEU-1 | BLEU-2 | BLEU-3 | BLEU-4 |
| -------------------- | ------ | ------ | ------ | ------ |
| word-level           | 42.31  | 25.26  | 14.81  | 8.73   |
| sequence-level       | 42.23  | 24.95  | 14.29  | 8.14   |
| BPE-level            | 43.06  | 25.57  | 15.16  | 9.17   |

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

Reviewer comment (Experiments):
Speculative Augmentation takes k=15 other copies of the Conformer, which is both computationally and memory expensive. An ablation on K should also be provided.

Author response:
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.

Reviewer comment (Reproducibility):
The training details of the Q-Conformer is only provided in the appendix but nowhere to be mentioned in the main text. The authors should not use the appendix as additional pages to the paper.

Author response:
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.

Reviewer comment (Reproducibility):
What is the vocabulary used during training the EEG-language alignment? Is it the vocabulary of the ZuCo dataset or the vocabulary of the LLM (T5)?

Author response:
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.

Reviewer comment (Reproducibility):
It is also not clear what loss function is used to train the multi-task query and which part of the network is updated. Is it the context transformer or the query prompts?

Author response:
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.

Reviewer comment (Experiments):
It is not clear why BPE is adopted. What would happen if the contrastive loss is used without BPE, meaning use whole word tokens?

Author response:
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.
| Contrastive Learning | Bleu-1 | Bleu-2 | Bleu-3 | Bleu-4 |
| --------------------- | ------ | ------ | ------ | ------ |
| None                  |41.57  | 24.03  | 13.8  | 8.06    |
| Word-level            | 42.31  | 25.26  | 14.81  | 8.73   |
| BPE-level             | 43.06  | 25.57  | 15.16  | 9.17   |

Reviewer comment (Reproducibility):
When combining Q-Conformer with LLM using continuous virtual prompts, how is it trained exactly? These are not even mentioned in the appendix.

Author response:
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.

Reviewer comment (Presentation):
Figure 3 presents EEG topography plots for both the input and output during the EEG token quantization process, leading to some ambiguity in interpretation. I would recommend the authors to elucidate this procedure in greater detail. Specifically, it would be insightful to understand whether the spatial arrangement of the EEG sensors played any role in this process.

Author 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.

Reviewer comment (Novelty):
The manuscript introduces BELT-2 as a progression from the prior BELT-1 model. However, the discussion and distinction between the two models are somewhat scanty, especially given their apparent similarities. It would be of immense value if the authors could elaborate on the design improvements made in BELT-2 over BELT-1. A focused discussion highlighting the specific enhancements and their contribution to the performance improvements, as showcased in Table 1 and Table 4, would add depth to the paper.

Author 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**.

Reviewer comment (Presentation):
In Figure 4 and Section 2.4, there is a mention of utilizing the mediate layer coding as 'EEG prompts'. The concept, as presented, leaves some gaps in understanding, primarily because its introduction and visualization seem absent or not explicitly labeled in the preceding figures and method sections. It would enhance coherence and clarity if the authors could revisit Figures 2 and/or 3 and annotate the specific parts illustrating this mediate layer coding.

Author 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. 

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.

Reviewer comment (Experiments):
In the section detailing the experimental setup, the authors introduced the dataset split. Was this split done on a cross-subject basis or just cross-training samples? Given the well-documented variations in EEG signals across different individuals, understanding this aspect is crucial. Will inter-individual variations impact the EEG-to-Language translation performance of the proposed method?

Author 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:

| test subject | bleu-1 | bleu-2 | bleu-3 | bleu-4 | r    | p    | f1   |
| ------------ | ------ | ------ | ------ | ------ | ---- | ---- | ---- |
| ALL          | 43.1   | 25.6   | 15.2   | 9.2    | 30.6 | 34.3 | 32.2 |
| ZPH          | 50.4   | 33.3   | 22.6   | 15.6   | 36.2 | 40.8 | 38.2 |
| ZMG          | 50.5   | 33.2   | 22.3   | 15.4   | 35.1 | 39.2 | 36.9 |
| ZKW          | 51.2   | 33.9   | 22.7   | 15.6   | 35.9 | 39.8 | 37.7 |
| ZKH          | 51.4   | 34.5   | 23.2   | 16.0   | 36.8 | 40.8 | 38.6 |
| ZKB          | 50.1   | 32.8   | 21.9   | 15.0   | 35.3 | 39.3 | 37.1 |
| ZJS          | 48.4   | 30.8   | 19.7   | 12.9   | 33.6 | 37.7 | 35.4 |
| ZGW          | 51.2   | 34.8   | 24.0   | 16.7   | 36.3 | 40.3 | 38.1 |
| ZJN          | 48.0   | 30.5   | 19.7   | 13.0   | 33.6 | 37.8 | 35.4 |
| ZDM          | 50.0   | 33.0   | 22.2   | 15.2   | 35.5 | 39.7 | 37.4 |
| ZAB          | 50.1   | 33.0   | 22.3   | 15.3   | 34.6 | 38.6 | 36.4 |

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. 

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.

Reviewer comment (Writing):
The last sentence on the first paragraph of page 1 says that previous methods have not achieved a "general bridging between brain signals and languages." Can you say more precisely what this means? Does "general bridging" refer to the multi-task setting? Is that the main difference between BELT and BELT-2?

Author 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.

Reviewer comment (Reproducibility):
I have a question about the input to the conformer. Page 3 says that the input is created by "segmenting the raw EEG waveform signal into windows using eye-tracking information." Are the segments of uniform length? If so, then why is the eye-tracking information necessary? If not, then does some sort of truncation or down-sampling occur?

Author 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.

Reviewer comment (Evaluation):
I have a question about the multi-task query described on page 4. It says that "we could easily extend the Q-conformer to a new downstream task by initializing a new set of query tokens to extract information related to the new task, obviating the need for training an entirely new model." But is this really a saving in time? Wouldn't the same amount of time be needed to train the new set of query tokens, even if you start with the existing model weights?

Author 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.

Reviewer comment (Theory):
The interaction of all the components in Equation 2 is unclear, and it remains uncertain whether the introduction of certain hyperparameters is necessary.

Author 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. 

In the latest revision, we have improved the readability of the paper and provided a clearer introduction to the VQ loss term in Equation 2 in the main paper. Also, we include more details of the training process, loss function, and combination of different loss functions in Figure 3, Section 2.2, Section 2.3, Section 2.4, Appendix C.2 and C.3. Please refer to the current main paper for a more detailed introduction of the training and loss function used in this paper.

Reviewer comment (Novelty):
If EEG representations are employed using a contrastive learning approach, the method proposed in Section 2.3 appears to be a conventional operation.

Author 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]. 

We have improved readability and understanding for a general audience in the latest revision. In particular, the negative contrastive learning loss is more clearly explained in Section 2.2 in the latest revision which is used to improve the distinction between discrete EEG tokens (page 4 last paragraph).

Reviewer comment (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.

Author 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**.

| BPE Contrastive | BLEU-1 | BLEU-2 | BLEU-3 | BLEU-4 |
| --------------- | ------ | ------ | ------ | ------ |
| No              | 41.57  | 24.02  | 13.8   | 8.06   |
| Yes             | 43.06  | 25.57  | 15.16  | 9.17   |

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).

Reviewer comment (Reproducibility):
Some details related to the training and the loss function are not provided.

Author response:
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).

Reviewer comment (Reproducibility):
It is suggested to highlight the difference between BELT and BELT-2, and to provide more precision on the training details and the loss function.

Author response:
Due to the page limit, we could only summarize the differences between BELT-2 and BELT-1 in the paper. In our latest revision of the paper, we have emphasized the difference between BELT-2 and BELT-1 in three key aspects. (1) BELT-1 is designed for a single task (e.g., the translation or the sentiment classification) while BELT-2 is for multi-task decoding (Page 2, line 8-9). (2) Compared to BELT-1, BELT-2 aligns EEG with language in a more fine-grained subword level (page 4). (3) Compared to BELT-1 and all other models, BELT-2 is the first framework that achieves the bridging between the EEG encoder and LLM (Page 2, second paragraph). 

We have provided more precise training details and loss functions in Figure 3, Section 2.2, Section 2.3, and Section 2.4 in the latest revision of the paper. Please refer to the current main paper for more details. Also, additional details on training are also provided in the supplementary in Appendix C.2 and C.3. We also significantly improved the main paper for better readability. Please reconsider raising the rating for our paper as it signatures a significant step
forward in integrating human brain signals with LLMs. This paper also presents solid experimental results with open-source code for ensuring reproducibility. We believe our research could have the potential to inspire further exploration and innovation in this research field.