STPC (Spectral-Temporal Physiological Consistency) EEG

An open-source, physics-informed AI that cleans noisy medical signals, turning untrustworthy data from wearables and clinics into confident, life-saving diagnoses for heart and brain conditions.

Comment

Inspiration

My journey started with a straightforward yet significant insight from my earlier experience with ECG signals: AI models possess immense power, but they are also remarkably simplistic. I noticed that a typical deep learning model, designed to denoise a noisy heartbeat, frequently generated a "clean" signal that appeared unnaturally smooth. The noise had been eliminated, but the 'information' of the signal was also lost: the distinct, rapid spikes essential for a cardiologist's diagnosis. The model lacked an understanding of the fundamental physics behind a heartbeat.

This prompted a crucial inquiry: how can we develop AI systems for science and medicine that are not only precise in theory but also inherently trustworthy? What methods can we use to instruct a model in the principles or reasoning of the real world? This inquiry ignited the foundation of my framework, Spectral-Temporal Physiological Consistency (STPC). I aimed to surpass mere pattern recognition and create an AI rooted in the fundamental principles of biophysics. The primary aim was not merely to purify a signal, but to recreate a signal that was physically feasible and diagnostically dependable, beginning with ECG and subsequently tackling the much more intricate task of EEG

What it does

STPC is a comprehensive research framework aimed at generating reliable, cutting-edge outcomes in the field of biomedical signal processing. Essentially, it is a deep learning framework that is capable of denoising and examining intricate EEG (brain wave) signals, yet its real strength is found in its distinct approach.

The initiative performs three main tasks:

  1. Physiologically Valid Denoising: It receives a noisy, unusable EEG signal (for instance, from a patient experiencing an epileptic seizure) and employs a U-Net model trained with the STPC regularizer to generate a clear signal. In contrast to typical denoisers, the result maintains the vital spatio-temporal patterns and frequency characteristics that are crucial for medical assessment.

  2. Surgical, Frequency-Sensitive Filtering: The STPC framework can be set up to function as a smart filter. I showed that it can be trained to surgically eliminate strong out-of-band noise while accurately maintaining a subtle, neurologically important brainwave (the Alpha rhythm), a challenge that is exceptionally hard for conventional filters.

  3. Unsupervised Representation Learning ("Brain2Vec"): In its most sophisticated use, the framework employs a self-supervised learning objective to uncover the underlying patterns in unlabeled EEG data. By requiring a model to rebuild masked signals while following the physical principles of STPC, it naturally develops the ability to distinguish between healthy and diseased brain states. It efficiently acquires a "vocabulary" of the brain without encountering any labels.

The whole project comes with an easy-to-use Streamlit application that enables anyone to experience the effectiveness of STPC using their own data

How I built it

This initiative was developed from scratch utilizing a contemporary, open-source framework, emphasizing reproducibility and extensibility.

  • Core Library (stpc): The project transitioned from a collection of scripts to a formal Python library. This collection includes the essential elements:

  • Models: A versatile UNet1D structure and an ECGClassifier developed using PyTorch.

  • Setbacks: The core of the initiative. Custom PyTorch loss functions were developed for every element of STPC: a TemporalGradientLoss, a SpatialLaplacianLoss, and a BandMaskedFFTLoss.

  • Data Utilities: A strong data loading and preprocessing framework created with MNE for EEG and WFDB for ECG, intended to manage the irregularities of actual medical datasets.

  • Experimentation Framework: I developed a consolidated experiment executor (run_training.py and run_validation.py) employing Python's argparse library. This enables me to methodically execute and verify all experiments (spatio-temporal, frequency-specific, self-supervised) using straightforward command-line instructions.

  • Training & Validation: Model training took place in the STPC_EEG_Research_Hub.ipynb notebook on Google Colab, using free T4 GPUs. This guaranteed rapid iteration and reproducibility. Validation encompassed producing not just quantitative metrics (RMSE, SSIM, Spectral Coherence, Silhouette Score) but also engaging qualitative visuals such as animated topomaps and PSD plots.

  • Archiving & Collaboration: The initiative is entirely managed with version control through Git. The complete code, document, and outcomes for v1.0.0 have been permanently stored on Zenodo (creating a citable DOI) and are highlighted on a main OSF project page.

Challenges I ran into

This project was an exploration of the tough truths of practical data science. The most significant difficulties lay not in the AI modeling, but in the data itself.

  1. Data Discrepancies: The CHB-MIT EEG dataset, although essential, consists of files with varying channel names, repeated channels, and case-sensitivity problems. This necessitated creating an extremely resilient data pipeline capable of dynamically identifying a consistent array of channels in all files and managing various edge cases. Resolving this issue was a lengthy and challenging task that required the development of tailored inspection tools to grasp the actual foundation of the data.

  2. The Limitations of Basic Metrics: A particularly intriguing challenge was a scientific one. My preliminary quantitative findings indicated that my visually enhanced STPC model achieved a comparable or possibly slightly lower SSIM score than the baseline model. This made me understand that basic metrics often overlook the structural and physiological traits that are truly significant. It was an important lesson: don't overly rely on your metrics. This obstacle transformed the whole story of my paper, converting it from a straightforward "my model is superior" conclusion, into a more nuanced analysis of assessment methods in scientific AI.

  3. Resource Limitations: Working with extensive datasets for training and validation in a free Colab environment presents ongoing challenges related to memory restrictions. The validation script, designed to produce numerous high-resolution plots for a video, repeatedly failed because of memory overloads. This necessitated meticulously restructuring the plotting loops to enhance memory efficiency, intentionally releasing figure objects from memory once each frame was produced.

Accomplishments that I am proud of

  1. The "Striking Visuals": The spatio-temporal GIF, the PSD chart, and the UMAP scatter plot serve not merely as figures; they provide clear, intuitive evidence supporting our hypotheses. They convey the narrative of STPC's strength more compellingly than any numerical data could.

  2. The Self-Supervised Outcome: Demonstrating that my STPC-regularized model could autonomously distinguish between seizure and non-seizure states with a high Silhouette Score (0.6350) was a genuine "wow" moment. It confirmed my most profound hypothesis: that by instructing a model in the physics of a signal, it can grasp its semantics.

  3. Creating an Exceptionally Strong Pipeline: The process of resolving all the data discrepancies was difficult, yet the end product is a polished, durable data pipeline that withstands real-world data issues. I take pride in having developed a framework that serves not only as a proof of concept but also as a robust base for upcoming research.

  4. Finalizing the Entire Research Cycle: I progressed this project from an initial concept to a thoroughly documented, refined, and publicly archived scientific output with a citable DOI. My greatest achievement is this dedication and adherence to the entire process.

What I learnt

. Physics-Informed > Naive: The key takeaway is that for scientific AI, a model's alignment with the domain's physical principles typically outweighs its efficacy on a rudimentary metric. An inductive bias serves as a strong guide rather than a limitation.

  1. Challenge Your Metrics: I discovered that metrics may not always be truthful, or at the very least, may not provide the complete picture. A comprehensive assessment that integrates various, focused quantitative measures with thorough qualitative examination is crucial. The human eye remains the finest evaluator of scientific credibility.

  2. The Importance of Data Engineering: I dedicated more time to troubleshooting data loading and preprocessing than to creating neural networks. This reinforced the timeless lesson: the cornerstone of any effective AI initiative is clear, coherent, and thoroughly comprehended data.

What's next for STPC (Spectral-Temporal Physiological Consistency) EEG

This initiative marks the initial launch of what I foresee as a lasting, open-source research platform. I am currently bringing on board our initial contributors to address the upcoming challenges. My plan incorporates:

  1. Ablation Study on Self-Supervised Learning: Demonstrate quantitatively that the STPC regularizer was crucial for the successful clustering in Phase 3 through a comparison with a baseline L1 autoencoder.

  2. Investigating Advanced Self-Supervised Techniques: Progress from masked reconstruction to robust contrastive learning approaches such as SimCLR to acquire more comprehensive neural representations.

  3. Enhancing the Spatial Model: Substitute the basic Laplacian loss with an advanced Graph Convolutional Network (GCN) layer to capture intricate, long-range spatial relationships in the brain.

  4. Towards a "Brain GPT": Our ultimate goal is to leverage the robust, reliable embeddings produced by our framework as the essential lexicon for a generative transformer model that can forecast and generate likely patterns of neural activity.

Built With

Share this project:

Updates