Lightweight Transformer for EEG Classification via Balanced Signed Graph Algorithm Unrolling

This paper proposes a lightweight, interpretable transformer-like neural network for EEG-based epilepsy classification that unrolls a spectral denoising algorithm on balanced signed graphs, achieving performance comparable to deep learning models with significantly fewer parameters.

The Gist

Imagine you are trying to listen to a conversation in a very noisy, crowded room. Some people are whispering in harmony (positive correlations), while others are shouting contradictory things at the same time (negative correlations/anti-correlations). Your goal is to figure out if the room is filled with healthy people chatting or people having a seizure.

This paper presents a new, super-efficient way to do exactly that using brain signals (EEG). Instead of using a massive, heavy, and confusing "black box" computer brain (like the giant AI models everyone uses today), the authors built a lightweight, transparent, and smart detective based on a clever mathematical trick.

Here is the breakdown of how it works, using simple analogies:

1. The Problem: The "Black Box" vs. The "Lightweight Detective"

Current AI models for reading brain waves are like giant, expensive supercomputers. They are incredibly powerful but:

• They are too heavy to run on small, portable medical devices.
• They are "black boxes," meaning even the scientists don't fully understand why they made a decision.

The authors wanted to build a lightweight detective that is:

• Small enough to fit on a tiny chip.
• Transparent: You can see exactly how it thinks.

2. The Core Idea: The "Balanced Signed Graph"

Imagine the brain sensors as people in a room.

• Positive edges: Two people who agree with each other (holding hands).
• Negative edges: Two people who disagree or are doing the opposite (pushing away).

Most AI models only look at people holding hands (positive connections). But in the brain, "pushing away" (anti-correlation) is just as important. The authors realized that if you have a room full of people who are either holding hands or pushing away, you can organize them into a "Balanced Room."

The "Balanced Room" Analogy:
Imagine you can split everyone in the room into two teams: Team Red and Team Blue.

• If two people are on the same team, they hold hands (positive connection).
• If they are on different teams, they push away (negative connection).

If you can do this perfectly without any confusion (no one is holding hands with someone on the other team while also pushing them), the room is "Balanced." This mathematical balance allows the AI to understand the "music" (frequencies) of the brain waves perfectly, even with the negative pushes.

3. The Magic Trick: "Unrolling" an Algorithm

Usually, to clean up a noisy signal, you run a mathematical recipe step-by-step.

• Old way: Run the recipe 10 times, then stop.
• The Authors' way (Unrolling): They took that 10-step recipe and turned every single step into a layer of a neural network.

Think of it like a factory assembly line. Instead of a robot doing a task and then stopping to think, they built a conveyor belt where the task happens, then the product gets tweaked, then it moves to the next station. This turns a slow, iterative math problem into a fast, one-pass neural network.

4. The "Spectral Denoising" (Cleaning the Signal)

Once the "Balanced Room" is set up, the AI acts like a noise-canceling headphone.

• It knows that "healthy brain waves" and "seizure brain waves" have different rhythms.
• It uses a Low-Pass Filter (like a sieve) to let the smooth, slow rhythms pass through and block the jagged, noisy spikes.
• Crucially, the AI learns exactly how big the holes in the sieve should be, rather than being told by a human.

5. How It Classifies: The "Two-Recorder" Test

Here is the clever part on how it decides if a patient has epilepsy or not:

1. The AI trains two different "cleaners":
   • Cleaner A is trained only on healthy brain waves. It learns what "normal" sounds like.
   • Cleaner B is trained only on seizure brain waves. It learns what "seizure" sounds like.
2. When a new, unknown brain signal comes in:
   • Cleaner A tries to fix it. If the signal was actually healthy, Cleaner A fixes it perfectly. If it was a seizure, Cleaner A gets confused and leaves a lot of "noise" (error).
   • Cleaner B tries to fix it. If the signal was a seizure, Cleaner B fixes it perfectly. If it was healthy, Cleaner B leaves a lot of noise.
3. The Verdict: The system asks, "Which cleaner left the least mess?"
   • If Cleaner A left the least mess, it's a Healthy patient.
   • If Cleaner B left the least mess, it's an Epilepsy patient.

6. The Results: Small but Mighty

The paper tested this on real data from 121 people.

• Performance: It was just as good as the giant, heavy AI models (getting about 97.6% accuracy).
• Efficiency: It used less than 1% of the computer memory and parameters required by the big models.
• Speed: It trains and runs much faster.

Summary

The authors took a complex math problem (cleaning brain waves with negative connections), turned it into a "Balanced" structure, and built a tiny, transparent neural network to solve it. It's like replacing a giant, fuel-guzzling truck with a sleek, electric bicycle that gets you to the same destination, faster, cheaper, and you can see exactly how the gears work.

This is a huge step forward for making AI-powered medical diagnosis possible on small, portable devices that doctors can actually use in the real world.

Technical Summary

1. Problem Statement

The paper addresses the binary classification of Electroencephalogram (EEG) signals to distinguish between epilepsy patients and healthy subjects.

• Challenges: Existing deep learning (DL) models, particularly Transformer-based architectures, achieve state-of-the-art (SOTA) accuracy but suffer from two major drawbacks:
  1. High Parameter Count: They are computationally expensive and memory-intensive, making them unsuitable for resource-constrained EEG devices.
  2. Lack of Interpretability: They function as "black boxes," offering little insight into the underlying signal processing mechanisms.
• Specific Signal Characteristic: EEG signals often exhibit pairwise anti-correlations (negative correlations) between sensors. Standard Graph Signal Processing (GSP) methods typically model only positive correlations (positive graphs), failing to capture these critical negative relationships effectively. Furthermore, general signed graphs (with both positive and negative edges) lack well-defined frequency domains, complicating spectral filtering.

2. Methodology

The authors propose a Lightweight, Interpretable Transformer constructed via Algorithm Unrolling. The core idea is to unroll an iterative optimization algorithm for signal denoising on a Balanced Signed Graph into a feed-forward neural network.

A. Balanced Signed Graph Construction

• Signed Graphs: The model represents EEG sensors as nodes and their correlations as edges. Positive edges denote positive correlations; negative edges denote anti-correlations.
• Graph Balance: A signed graph is "balanced" if it contains no cycles with an odd number of negative edges. This property is crucial because it allows the graph to be mapped to a corresponding Positive Graph ($G^+$) via a similarity transform.
  • Transformation: Using the Cartwright-Harary Theorem, node polarities ($\beta_i \in \{1, -1\}$) are assigned. The Laplacian of the balanced signed graph ($L_B$) is related to the Laplacian of the positive graph ($L_+$) by $L_+ = T L_B T^{-1}$, where $T = \text{diag}(\beta)$.
  • Benefit: This mapping ensures that $L_B$ and $L_+$ share the same eigenvalues, allowing the use of well-established spectral filtering techniques (defined for positive graphs) on the signed graph.
• Learning Process:
  1. Polarity Selection: Node polarities are initialized based on empirical covariance signs and iteratively optimized to minimize the Graph Laplacian Regularizer (GLR), ensuring the graph remains balanced.
  2. Edge Weights: Signed weights are computed based on the Mahalanobis distance between node features. Positive distances map to positive weights; negative distances (anti-correlations) map to negative weights, scaled by the node polarities.
  3. PSD Enforcement: Self-loops are added using the Gershgorin Circle Theorem to ensure the Laplacian is Positive Semi-Definite (PSD).

B. Algorithm Unrolling & Network Architecture

The network is built by unrolling an iterative denoising algorithm:

1. Graph Learning Module (BGL): Updates the graph structure (polarities and edge weights) based on current signal features. This module acts as a learnable self-attention mechanism. Unlike standard Transformers that learn dense Key/Query matrices, this module learns a shallow CNN for feature extraction and a metric matrix, significantly reducing parameters.
2. Low-Pass Filter (LPF): Performs spectral denoising on the mapped positive graph ($G^+$).
   • Ideal Filter: Theoretically projects the signal onto the low-frequency subspace of $L_+$.
   • Lanczos Approximation: To avoid the $O(N^3)$ cost of full eigen-decomposition, the authors use the Lanczos method to approximate the low-pass filter in linear time ($O(N)$).
   • Learnable Cutoff: The cutoff frequency ($\omega$) is a learnable parameter per block, replacing hard truncation with a smooth sigmoid function.
3. Architecture: The network consists of stacked blocks, each containing a BGL module followed by an LPF module.

C. Classification Strategy

Instead of a standard classification head, the method uses two class-specific denoisers ($\Psi_0$ for healthy, $\Psi_1$ for epilepsy):

• Pretext Task: Each denoiser is trained to reconstruct signals from its specific class, effectively learning the posterior probability distribution $P(x|y, c)$.
• Decision Rule: For a new input $y$, the model computes the reconstruction error for both denoisers. The class with the minimum reconstruction error is selected:
  $$c^* = \arg \min_{c \in \{0,1\}} \| y - \Psi_c(y) \|^2_2$$
• Loss Function: A contrastive loss is used during training to ensure the denoiser for class 0 minimizes error on class 0 data while maximizing error on class 1 data (and vice versa), enhancing class separability.

3. Key Contributions

1. Balanced Signed Graph Unrolling: Extends algorithm unrolling to balanced signed graphs, enabling the modeling of EEG anti-correlations while retaining rigorous frequency definitions via similarity transforms.
2. Efficient Spectral Filtering: Implements an ideal low-pass filter on signed graphs using Lanczos approximation, achieving linear time complexity without full eigen-decomposition.
3. Interpretable Self-Attention: Demonstrates that the graph learning module with normalized edge weights is mathematically equivalent to self-attention but with drastically fewer parameters (replacing large Key/Query/Value matrices with shallow CNNs and a single cutoff frequency).
4. Generative-Discriminative Bridge: Uses a reconstruction-error-based classification strategy where two denoisers learn posterior probabilities, bridging generative modeling and discriminative classification in an interpretable manner.

4. Experimental Results

The method was evaluated on the Turkish Epilepsy EEG Dataset (10,356 recordings, 121 subjects) and the TUH Abnormal EEG Corpus.

• Performance vs. SOTA:
  • Default Task: Achieved 97.57% accuracy and 98.01% F1-score.
  • LOSO (Leave-One-Subject-Out): Achieved 90.06% accuracy and 92.59% F1-score, outperforming other graph-based methods (DGCNN, GIN, EEGNet) and demonstrating strong generalization to unseen subjects.
  • Comparison: Outperformed a standard Transformer model (85.12% accuracy) and matched or exceeded complex CNN-based models (e.g., STFT+CNN at 99.20% accuracy) while using significantly fewer parameters.
• Efficiency:
  • Parameters: Only 14,787 parameters (less than 1% of the Transformer model's ~1.8M parameters and ~0.1% of the STFT+CNN model's ~11.5M).
  • Speed: Training time was 2 hours 14 minutes (vs. 22 hours for FBCSPNet), and inference time was 55 seconds (vs. 613 seconds for EEGNet).
• Ablation Studies:
  • Graph Type: Balanced Signed Graphs significantly outperformed both Positive Graphs (ignoring anti-correlations) and Unbalanced Signed Graphs (ill-defined frequencies).
  • Loss Function: Contrastive loss significantly improved classification metrics compared to standard MSE.

5. Significance

This paper presents a paradigm shift in EEG classification by combining Graph Signal Processing with Algorithm Unrolling.

• Practicality: The extreme reduction in parameters and computational cost makes the model viable for deployment on wearable and resource-constrained medical devices.
• Interpretability: By unrolling a mathematically defined denoising algorithm, every layer corresponds to a specific optimization step, providing transparency often missing in deep learning.
• Signal Modeling: It successfully addresses the unique challenge of anti-correlations in EEG data through balanced signed graphs, a feature often overlooked in standard GSP approaches.
• Future Impact: The framework offers a blueprint for building lightweight, interpretable, and highly efficient neural networks for other time-series signal processing tasks where negative correlations are prevalent.