Self-Supervised Learning with Graph Neural Networks for Seizure Detection

Image generated by DALL·E 3

Introduction

This research project adapts several self-supervised learning (SSL) techniques with Graph Neural Network (GNN) encoders to the domain of seizure detection. This work is part of an ongoing collaborative effort with the Neural Signal & Brain Processing Lab at the Krembil Research Institute. The objective is to train a GNN encoder and then evaluate its performance downstream on seizure detection with the following SSL methods: Relative Positioning, Temporal Shuffling, Contrastive Predictive Coding (CPC), and Variance-Invariance-Covariance Regularization (VICReg). The GNN architecture comprises of Edge-Conditioned Convolution (ECC) and Graph Attention Network (GAT) layers, utilizing PyTorch and PyTorch Geometric libraries for standard deep learning and GNN implementation. For a more detailed description of our research, see projects/ssl-seizure-detection.

For more information on SSL and GNNs, please refer to the relevant papers:

• A Path Towards Autonomous Machine Intelligence (Lecun, 2022)
• Edge-Conditioned Convolution (ECC): (Simonovsky & Komodakis, 2017)
• Graph Attention Layer (GAT): (Veličković et al., 2018)
• Relative Positioning and Temporal Shuffling (Banvile et al., 2021)
• Contrastive Predictive Coding (CPC) (Oord et al., 2018)
• Variance-Invariance-Covariance Regularization (VICreg) (Bardes et al., 2022)

Table of Contents

1. Installation
2. Data
3. Usage
4. File Descriptions
5. License
6. Contact

Installation

For detailed setup instructions and environment configuration, please see our installation guide.

Data

Processed data formatted as PyTorch Geometric Data objects is not yet available for release but will be provided once the paper is released. Meanwhile, the initial intracranial electroencephalogram (iEEG) dataset utilized for this project is publicly accessible on OpenNeuro, identified by Accession Number ds003029. From ds003029, we selected 26 patients out of the 88 available in the dataset. For each patient iEEG signal, we divided the signal into time windows of equal length, and for each window an initial graph representation was constructed with the following method. The initial graph representations are fully connected graphs, with the nodes corresponding to individual electrodes. To construct the edge features, for each electrode pair, we computed the Pearson correlation, phase-lock value (PLV), and the coherence, giving us edge features of dimension 3. To construct the node features, we used the average energy of the electrode, and the average energies at various frequency bands; due to variability in the iEEG data format, there may have been more or less frequency bands available for certain patients, thus node features dimensions varied between $7-9$ (dependent on the patient). From this, we converted the data to the standard PyTorch Geometric Data format of [edge_index, x, edge_attr, y] where edge_index is a tensor defining the graph structure (similar to a binary adjacency matrix), x is the node feature tensor, edge_attr is the edge feature tensor, and y is the target label tensor taking on values $0$ or $1$ for binary classification of ictal (seizure) or nonictal (no seizure); and took on values $0,1,$ or $2$ for preictal (before seizure), ictal (seizure), and postictal (after seizure) respectively for multiclass classification. Due of the heterogeneity of the iEEG data from patient to patient (e.g., different number of electrodes, different placement of electrodes) model training and evaluation was conducted intra-patient, thus a separate model was created for each patient.

Usage

To run the entire pipeline, please refer to the train() function in train.py and the relevant docstring. Please see the notebook train.ipynb for guidance on how to train each model. The main.py script is optimized for HPC on the Cledar cluster (Digital Research Alliance of Canada), and is not recommended for general use. For a tutorial on PyTorch Geometric and customized GNN models, please refer to tutorial.ipynb. To see how the initial graph representations are created, please refer to preprocess.ipynb. For a information on the transfer learning process (implemented in train.py), please refer to transfer.ipynb.

File Descriptions

• models.py: Contains self-supervised models: relative_positioning, temporal_shuffling, CPC (to be added), and VICReg (to be added); and supervised models: supervised (base model), downstream1, and downstream2.
• train.py: Implements the training loop for both self-supervised and supervised models. Includes logging with Weights and Biases (highly recomended). For a guide on how to train each model please see train.ipynb.
• main.py: The primary script to run the training pipeline in parallel on multiple patients, optimized for Cedar cluster resources. Please see the scripts folder for more details on how this is implemented.
• preprocess.py: Includes helper functions for all preprocessing tasks, such as converting initial graph representations to PyG-compatible structures.
• patch.py: Patches pre-existing Numpy data from our lab to the PyG-compatible format, not recommended for general use unless your existing data fits the specifications as outlined in preprocess.ipynb.

Acknowledgements

I would like to sincerely thank both Dr. Alan A. Díaz-Montiel and Dr. Milad Lankarany for their continued support throughout this research endeavour and for their expert guidance. I am also extremely grateful for the support from the Temerty Faculty of Medicine, University of Toronto through the esteemed T-CAIREM Summer Studentship, which allowed me to focus on this work.

I'd also like to extend my gratitude to the researchers and institutions for their generosity in sharing their iEEG data through OpenNeuro ds003029. Special thanks to:

• Department of Biomedical Engineering, Johns Hopkins University, Baltimore, United States
• Epilepsy Center, Cleveland Clinic, Cleveland, United States
• Department of Neurosurgery, University of Miami Miller School of Medicine, Miami, United States
• Department of Neurology, University of Miami Miller School of Medicine, Miami, United States
• Neurology, University of Maryland Medical Center, Baltimore, United States
• Neurology, Johns Hopkins Hospital, Baltimore, United States
• Surgical Neurology Branch, NINDS, NIH, Bethesda MD
• Neurosurgery, and Epilepsy Center, University of Pittsburgh Medical Center, Pittsburgh, United States
• Institute for Computational Medicine, Johns Hopkins University, Baltimore, United States

For additional details on the dataset, the foundational paper is available at with the following link:

• Li, A., et al. (2021). Neural fragility as an EEG marker of the seizure onset zone. Nature Neuroscience, 24(10), 1465–1474. Nature Publishing Group US New York..

Contact

For any queries, please contact xmootoo at gmail dot com.