AdityaRaikar/mpse-smoe-speech-enhancement

MP-SENet + Supervised Mixture of Experts (S-MoE) for Dual-Bandwidth Speech Enhancement

Speech enhancement model handling both Narrowband (NB, 8kHz) and Wideband (WB, 16kHz) audio using Supervised Mixture of Experts.

Based on:

• MP-SENet — STFT-domain magnitude-phase speech enhancement backbone
• S-MoE — Supervised MoE with deterministic hard gating by bandwidth metadata

Pre-trained Baseline

The official pre-trained MP-SENet baseline (from yxlu-0102/MP-SENet) is included at pretrained/g_best_vb.pt:

• Trained on VoiceBank+DEMAND, 16kHz, 100 epochs
• 2.263M params, 247 state dict keys
• No need to train baseline from scratch — S-MoE training initializes directly from this checkpoint

Training (S-MoE only)

pip install -r requirements-gpu.txt

# S-MoE training (downloads pre-trained baseline automatically)
SMOE_EPOCHS=60 BATCH_SIZE=2 GRAD_ACCUM_STEPS=2 python train.py

The script automatically:

1. Downloads VoiceBank+DEMAND (~5GB)
2. Pre-extracts WB/NB data into separate folders (no on-the-fly resampling)
3. Downloads official pre-trained MP-SENet baseline from this repo
4. Initializes S-MoE experts from baseline (shared → both WB+NB experts)
5. Trains S-MoE on joint WB+NB data for 60 epochs
6. Pushes trained model to HF Hub

Data Pipeline

All NB/WB data is pre-extracted to separate folders before training:

Folder	Type	Processing	bandwidth_id
original_wb/	WB	Original 16kHz (symlinks)	0
simple_nb/	NB	16k→8k→16k resampled	1
codec_nb/	NB	G.711 A-law/mu-law (round-robin)	1

Architecture

Component	Baseline	S-MoE
Total params	2.263M	3.587M
Active params (inference)	2.263M	2.263M (same!)
FFN type	Single GRU-FFN	2× GRU-FFN experts, hard gated by bandwidth_id
Attention	Shared	Shared (unchanged)

Environment Variables

Variable	Default	Description
SMOE_EPOCHS	60	Epochs for S-MoE training
BATCH_SIZE	2	GPU micro-batch size
GRAD_ACCUM_STEPS	2	Gradient accumulation (effective BS = BATCH_SIZE × GRAD_ACCUM_STEPS)
PESQ_EVERY_N	10	Compute PESQ for discriminator every N steps (expensive)
LR	5e-4	Learning rate
EVAL_INTERVAL	5	Evaluate every N epochs
BASELINE_REPO	AdityaRaikar/mpse-smoe-speech-enhancement	HF repo with pre-trained baseline
BASELINE_FILE	pretrained/g_best_vb.pt	Checkpoint filename in repo

Files

File	Description
train.py	Combined data prep + S-MoE training (self-contained)
pretrained/g_best_vb.pt	Official pre-trained MP-SENet baseline (VoiceBank+DEMAND)
smoe_models.py	Model definitions with detailed docstrings
data_preparation.py	Standalone NB/WB codec pipeline
config.json	Default hyperparameters

References

@article{lu2023mpse,
  title={MP-SENet: A Speech Enhancement Model with Parallel Denoising of Magnitude and Phase Spectra},
  author={Lu, Ye-Xin and Ai, Yang and Ling, Zhen-Hua},
  journal={arXiv preprint arXiv:2305.13686},
  year={2023}
}