MSVTNet: Multi-Scale Vision Transformer Neural Network for EEG-Based Motor Imagery Decoding

Architecture

The proposed MSVTNet network comprises three main blocks: multi-scale spatio-temporal convolutional (MSST) block, cross-scale global temporal encoder (CSGT) block, and auxiliary branch loss (ABL) block. Each branch in the MSST block extracts local spatiotemporal feature representations from the MI-EEG signals. By using multiple independent branches across different scales, more informative encoded representations can be extracted. Before entering the CSGT block, the encoded features from different scales of the MSST block and a class token are concatenated along the feature dimension to create an integrated global spatiotemporal representation enriched with multi-scale features. Internally within the CSGT block, cross-scale global temporal correlations are modeled through a multi-head self-attention (MHSA) mechanism, and the embedded class token is further input into the classifier (CLS) for final decoding. The ABL block, serving as intermediate supervision, addresses the parameter imbalance problem between the MSST and CSGT blocks to prevent overfitting. At the same time, it enhances the feature extraction capabilities of each branch.

Development environment

All models were trained and tested by a single GPU, Nvidia GeForce RTX 3090 (Driver 530.41.03, CUDA 12.1) on Ubuntu 22.04.2 LTS. The main following packages are required:

• Python 3.10.13
• NumPy 1.23.5
• PyTorch 2.0.1 + PyTorch CUDA 11.8
• dpeeg 0.3.6

more detailed dependencies are in environment.yml. Build the development environment using Miniconda:

conda env create -f environment.yml

Implementation

All the core codes are placed in the dpeeg 0.3.6 in the form of package functions, which provides some convenient interface functions to support the experiments in this paper. This repositories only provides the top-level training code. For details about the training code and related experimental methods, please check dpeeg.

After installing dpeeg, you should be able to run the code in the repositories correctly. If you want to run a certain analysis experiment, comment out the corresponding experiment code and run it from the beginning. All training details of the algorithms are also provided in the out folder for reference. The file tree of out is as follows:

out
|- decoding algorithm
   |- KFold
      |- datasets_SD
         |- ...
      |- datasets_SI
   |- LOSO_HO
      |- datasets
|- ...

where decoding algorithm folder contains all experimental analysis results for the corresponding algorithm. Inside, KFold stores the subject-dependent analysis results (datasets_SD and datasets_SI are the session-dependent and session-independent analysis results for the corresponding dataset, respectively), and LOSO_HO stores the subject-independent analysis results for the corresponding dataset. Note: Since the EEG Conformer results exceeded 7G, they were not uploaded.

Results

The overall classification results for MSVTNet and other competing architectures are as follows:

For easy reference, detailed results of all analyses are reorganized in the supporting document.

Cite

If you find this architecture or toolbox useful then please cite this paper:

Liu, Ke, Tao Yang, Zhuliang Yu, Weibo Yi, Hong Yu, Guoyin Wang, and Wei Wu. "MSVTNet: Multi-Scale Vision Transformer Neural Network for EEG-Based Motor Imagery Decoding." IEEE Journal of Biomedical and Health Informatics (2024).

and here is the Bibtex version:

@article{liu2024msvtnet,
  title={MSVTNet: Multi-Scale Vision Transformer Neural Network for EEG-Based Motor Imagery Decoding},
  author={Liu, Ke and Yang, Tao and Yu, Zhuliang and Yi, Weibo and Yu, Hong and Wang, Guoyin and Wu, Wei},
  journal={IEEE Journal of Biomedical and Health Informatics},
  year={2024},
  volume={28},
  number={12},
  pages={7126-7137},
  publisher={IEEE},
  doi={10.1109/JBHI.2024.3450753}
}