End-to-end learnable EEG channel selection for deep neural networks with Gumbel-softmax

Abstract

Objective. To develop an efficient, embedded electroencephalogram (EEG) channel selection approach for deep neural networks, allowing us to match the channel selection to the target model, while avoiding the large computational burdens of wrapper approaches in conjunction with neural networks. Approach. We employ a concrete selector layer to jointly optimize the EEG channel selection and network parameters. This layer uses a Gumbel-softmax trick to build continuous relaxations of the discrete parameters involved in the selection process, allowing them be learned in an end-to-end manner with traditional backpropagation. As the selection layer was often observed to include the same channel twice in a certain selection, we propose a regularization function to mitigate this behavior. We validate this method on two different EEG tasks: motor execution and auditory attention decoding. For each task, we compare the performance of the Gumbel-softmax method with a baseline EEG channel selection approach tailored towards this specific task: mutual information and greedy forward selection with the utility metric respectively. Main results. Our experiments show that the proposed framework is generally applicable, while performing at least as well as (and often better than) these state-of-the-art, task-specific approaches. Significance. The proposed method offers an efficient, task- and model-independent approach to jointly learn the optimal EEG channels along with the neural network weights.