Geometry Selection Module for EEG Enhancement

• Geometry Selection Module (GSM) is a framework that integrates fixed device geometry constraints with learned soft selection to identify a minimal, informative subset of EEG electrodes.
• It employs a two-stage masking approach that fuses hard selection based on predefined electrode regions with soft channel weights produced by a compact 1D convolutional network.
• Adjustable regularization parameters balance performance and computational efficiency, enabling robust auditory decoding while conforming to device-specific spatial layouts.

A Geometry Selection Module (GSM) refers here specifically to the geometry-constrained convolutional regularization selection (GC-ConvRS) framework introduced for EEG channel selection in brain-assisted speech enhancement (BASE). GSM enforces a device-constraint geometry (e.g., headphone-shaped region) and leverages both hard and soft selection mechanisms to identify a minimal, informative subset of EEG electrodes for optimal performance and reduced cost within deep learning pipelines for auditory attention decoding (Zuo et al., 19 Sep 2024).

1. Selection Objective and Geometry Constraints

The GSM formalizes channel selection with a two-stage mask combining hard device-level constraints and soft channel importance weights. Let $Q$ denote the total number of EEG electrodes; $\mathcal{Q} = \{1,\ldots,Q\}$ the full index set; $\mathcal{S} \subset \mathcal{Q}$ a hard pre-selected geometry-defined subset (e.g., headphone-shaped); $s \in \mathbb{R}^{|\mathcal{S}|}$ the learned soft-selection weights; $m_h \in \{0,1\}^Q$ the hard mask satisfying $(m_h)_i = 1$ for $i \in \mathcal{S}$, otherwise $0$.

The channel mask is constructed as

$m = m_h \odot \widetilde s,$

where

$$\widetilde s_i = \begin{cases} s_j, & i \in \mathcal{S},\text{ and }i \text{ is the }j\text{th in }\mathcal{S} \ 0, & i \notin \mathcal{S} \end{cases}$$

providing $EEG'(t,i) = EEG(t,i) \cdot \widetilde s_i$.

The training loss encompasses:

$$\mathcal{L} = -\alpha \mathcal{L}_{\rm SI\text{-}SDR} + \beta \mathcal{L}_d + \gamma \mathcal{L}_{\rm reg} + \lambda \mathcal{L}_{g},$$

with:

• $\mathcal{L}_{\rm SI\text{-}SDR}$: negative scale-invariant signal-to-distortion ratio,
• $\mathcal{L}_d$: discreteness regularizer favoring binary mask values,
• $\mathcal{L}_{\rm reg}$: $L_2$ regularizer on $s$ for sparsity,
• $\mathcal{L}_g$: optional geometry-regularization (sum of spatial distances weighted by $s_is_j$).

Geometry is enforced primarily by $m_h$ controlling the allowable subset, and optionally by $\mathcal{L}_g$ introducing spatial contiguity penalties indexed by electrode cap distances $D_{ij}$.

2. Soft-Selection Mask Architecture

Within the allowed geometry, soft selection refines which channels most contribute to the decoding task. The vector $s \in [0,1]^{|\mathcal{S}|}$ is learned by a compact 1D convolutional network ("ConvRS") operating on the downsampled EEG time series restricted to $\mathcal{S}$. This network outputs the weight vector via a final sigmoid activation.

Domain constraints force $s_j \in [0,1]$ and $\widetilde{s}_i = 0$ for $i \notin \mathcal{S}$. Channel-wise weighting by $\widetilde{s}$ precedes transmission to the WD-TCN separator, guaranteeing consistency with the geometric device region for all forward passes.

3. Integration with the WD-TCN Backbone

GSM (GC-ConvRS) integrates as follows into the weighted multi-dilation temporal convolutional network (WD-TCN) pipeline for BASE:

• Input: Noisy speech $x \in \mathbb{R}^{B \times T_{\rm audio}}$, EEG signals $E \in \mathbb{R}^{B \times T_{\rm eeg} \times Q}$, batch size $B$.
• Hard-masking: $$E_{\mathcal{S}} = E[..., \mathcal{S}] \in \mathbb{R}^{B \times T_{\rm eeg} \times |\mathcal{S}|}$$.
• Soft-selection: $$s \leftarrow \text{ConvRS}(E_{\mathcal{S}}) \in [0,1]^{|\mathcal{S}|}$$.
• Weighted EEG: $\hat E_{ij} = E_{ij} \cdot s_j$, $j = 1,\ldots,|\mathcal{S}|$.
• EEG Encoder: $$e_x = \text{EEGencoder}(\hat E) \in \mathbb{R}^{B \times T' \times C_e}$$.
• Audio Encoder: $$w_x = \text{AudioEnc}(x) \in \mathbb{R}^{B \times T' \times C_a}$$.
• Separator: $m_t = \text{Separator}(w_x, e_x) \in \mathbb{R}^{B \times T' \times N_{\rm mask}$.
• Decoder: $$\hat s_t = \text{Decoder}(w_x \odot m_t) \in \mathbb{R}^{B \times T_{\rm audio}}$$.

The GC-ConvRS operates immediately after raw EEG input, gating signal flow into feature encoding and fusion.

4. Training Procedure and Optimization

Training employs standard gradient-based optimization for both WD-TCN and ConvRS parameters ($\theta_{\rm sep}, \theta_{\rm rs}$), enforcing geometry via the hard mask and regularization via the loss function. Hard constraints are maintained by zeroing gradients for inactive indices ($m_h[j]=0$) or excluding those channels from ConvRS.

A summary of the algorithmic procedure is:

for epoch in 1..N_epochs: for minibatch {x, E}_{b=1}^B: # 1. Hard-select geometry E_S = E[..., i where m_h[i]==1] # 2. Soft-select via ConvRS s = sigmoid(ConvRS_Net(E_S; theta_rs)) # 3. Apply soft × hard mask E_weighted[...,j] = E_S[...,j] * s[j] # 4. Forward through WD-TCN BASE w_x = AudioEncoder(x) e_x = EEGEncoder(E_weighted) m_t = Separator(w_x, e_x; theta_sep) s_hat = Decoder(w_x ⊙ m_t) # 5. Compute loss components L_sisdr = compute_SI_SDR_loss(s_hat,x_target) L_d = k1 * (-(s-0.5)·(s-0.5)/(Q*B) + 0.25) L_reg = k2 * sum(s^2) L_geo = sum_{i,j} s_i*s_j*D_{ij} # optional L_total = -alpha * L_sisdr + beta * L_d + gamma * L_reg + lambda * L_geo # 6. Gradient update (theta_sep, theta_rs) -= lr * grad(L_total)

5. Quantitative Effects of Geometry Constraints

Performance and cost trade-offs can be managed by tuning the sparsity weight $\gamma$ (and geometry regularization $\lambda$ if activated), directly controlling the number of retained channels $|\text{supp}(s)|$. On a 33-subject public EEG speech dataset (Zuo et al., 19 Sep 2024):

Selection strategy	Channels $|\text{supp}(s)|$	SI-SDR (dB)	PESQ	STOI
Hard-only	30	10.8	2.66	0.88
Soft GC-ConvRS, γ=0.1	≈18	10.9	2.65	0.88
Soft GC-ConvRS, γ=0.2	≈16	10.7	2.63	0.87
Soft GC-ConvRS, γ=0.3	≈12	10.4	2.58	0.86
Soft GC-ConvRS, γ=0.4	≈11	10.1	...	...
Soft GC-ConvRS, γ=0.6	≈6	9.8	...	...

As $\gamma$ increases, selected channels decrease, reducing computational cost. Retaining 18 channels instead of 128 yields roughly $7 \times$ speed-up in the EEG encoder and feature fusion. Qualitatively, selected electrodes cluster near left/right temporal regions and the ears, which are known loci for auditory attention signatures.

6. Practical Impact and Qualitative Electrode Patterns

The GSM ensures channel selection strictly adheres to user-defined hardware geometry, supporting application constraints (e.g., headphones, wearable EEG). The method enables trade-offs between hardware cost and speech separation performance by modulating regularization parameters, making it practical for integration into wearable BASE systems.

Selection patterns observed favor electrodes over auditory cortex and periauricular regions, aligning with neurophysiological expectations. A plausible implication is that GSM can generalize across device designs where a fixed geometry must be respected, and offers extensibility to incorporate advanced geometry-dependent penalties should the spatial arrangement of electrodes influence decoding quality.

7. Extensions and Theoretical Significance

While the original implementation disables the geometry-regularization term ($\lambda=0$), the loss definition provides a pathway for further research into spatially-adaptive selection regularly. This suggests GSM can be adapted for situations where within-geometry spatial clustering or dispersion is desirable, for example in minimizing wiring complexity or maximizing cortical spatial coverage.

GC-ConvRS, as a GSM, demonstrates the utility of embedding geometry constraints into differentiable selection modules within deep neural architectures for neurotechnological applications. Its separation of hard (device-driven) and soft (performance-driven) channel selection offers a formalized approach uniquely suited to real-world hardware limitations in brain-computer interface design (Zuo et al., 19 Sep 2024).