PS4: Proxy-Supervised Joint Training for Real Target Speaker Extraction

PS4 is a Target Speaker Extraction (TSE) model that jointly optimizes speech separation quality and ASR transcription accuracy through proxy supervision. It is fine-tuned from a pretrained BSRNN + ECAPA-TDNN backbone on the REAL-PS4 dataset — a multi-domain real-recording corpus covering Chinese and English meeting scenarios.

Model Architecture

Component	Details
Separation backbone	BSRNN (Band-Split RNN), feature_dim=128, num_repeat=6, stride=128, win=512
Speaker encoder	ECAPA-TDNN (ECAPA_TDNN_GLOB_c512), embed_dim=192, feat_dim=80, ASTP pooling
Speaker fusion	Element-wise multiply (spk_fuse_type: multiply)
ASR backbone (proxy)	Whisper large-v3 (frozen, used only during training)
Speaker encoders (eval)	ResNet34-LM — voxceleb_resnet34_LM (EN), cnceleb_resnet34_LM (ZH)
Sample rate	16 kHz

Training

Loss Function

PS4 uses a combined proxy-supervised loss that simultaneously optimizes four objectives:

L = λ_ce · L_CE  +  λ_sim · L_sim  +  λ_vad · L_VAD  +  λ_dnsmos · L_DNSMOS

Loss term	Weight	Description
L_CE	1.0	ASR cross-entropy loss (Whisper large-v3 teacher-forcing)
L_sim	5.0	Speaker similarity ranking loss: hinge(margin − (sim(tse,enroll) − sim(mix,enroll))), margin=0.5
L_VAD	0.5	Target speaker activity detection loss (frame-level energy supervision)
L_DNSMOS	0.2	Differentiable DNSMOS-OVRL loss (ONNX model, no reference audio needed)

Training Data

Trained on REAL-PS4, which aggregates four real-recording meeting datasets:

Dataset	Language	Scenario
AISHELL-4	Chinese	Multi-speaker meeting
AliMeeting	Chinese	Multi-speaker meeting
AMI	English	Multi-speaker meeting
CHiME6	English	Dinner party / far-field

Hyperparameters

Parameter	Value
Optimizer	AdamW
Learning rate	1e-5
Weight decay	1e-5
LR scheduler	ExponentialDecrease (1e-5 → 1e-6)
Batch size	1
Max audio length	30 s (480,000 samples)
Max enrollment length	10 s (160,000 samples)
Gradient clipping	5.0
Precision	FP32
Checkpoint	Epoch 37

Usage

Load the Model

import torch

# Load checkpoint
ckpt = torch.load("checkpoint_epoch037.pt", map_location="cpu")

# The checkpoint contains the full TSE model state dict
# Compatible with the wesep BSRNN + ECAPA-TDNN framework
model_state = ckpt["model"] if "model" in ckpt else ckpt

Inference Example

import torch
import torchaudio
from wesep.models import get_model

# Initialize model (must match training config)
model = get_model("BSRNN")(
    feat_type="consistent",
    feature_dim=128,
    num_repeat=6,
    spk_emb_dim=192,
    spk_fuse_type="multiply",
    spk_model="ECAPA_TDNN_GLOB_c512",
    sr=16000,
    win=512,
    stride=128,
)

# Load PS4 weights
ckpt = torch.load("checkpoint_epoch037.pt", map_location="cpu")
model.load_state_dict(ckpt["model"] if "model" in ckpt else ckpt)
model.eval()

# Load mixture and enrollment audio (16 kHz mono)
mixture, sr = torchaudio.load("mixture.wav")
enrollment, sr = torchaudio.load("enrollment.wav")

# Run extraction
with torch.no_grad():
    extracted = model(mixture, enrollment)

Note: The speaker encoder is frozen during training (spk_model_freeze: true). For inference, use the same ResNet34-LM speaker encoder as the evaluation pipeline: voxceleb_resnet34_LM for English, cnceleb_resnet34_LM for Chinese.

Pretrained Backbone

PS4 is fine-tuned from bsrnn_ecapa_vox1, a BSRNN + ECAPA-TDNN model pretrained on VoxCeleb1. The proxy-supervised fine-tuning on REAL-PS4 adapts the model to real far-field multi-speaker meeting conditions.

Citation

@misc{ning2026ps4,
  title         = {PS4: Proxy-Supervised Joint Training for Real Target Speaker Extraction},
  author        = {Wanyi Ning and Wei Zhou and Yingpeng Li and Yinshang Guo and Haitao Qian and Yiming Cheng},
  year          = {2026},
  eprint        = {2607.08111},
  archivePrefix = {arXiv},
  primaryClass  = {cs.SD},
  url           = {https://arxiv.org/abs/2607.08111}
}

Related Resources

• Dataset: TaurenMountain/REAL-PS4