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

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

Training Data

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

Hyperparameters

Usage

Quick Start (recommended)

Use the included inference.py — a self-contained script with absolutely no external dependencies beyond torch, torchaudio, and numpy:

# Install dependencies
pip install torch torchaudio numpy

# Single file extraction
python inference.py \
    --checkpoint checkpoint_epoch037.pt \
    --mix mix.wav \
    --enroll target_speaker.wav \
    --output result.wav

# Use GPU
python inference.py \
    --checkpoint checkpoint_epoch037.pt \
    --mix mix.wav \
    --enroll target.wav \
    --output result.wav \
    --device cuda

# Batch mode (process all .wav files in a directory)
python inference.py \
    --checkpoint checkpoint_epoch037.pt \
    --mix-dir ./mixtures/ \
    --enroll-dir ./enrollments/ \
    --output-dir ./results/ \
    --device cuda

# List available CUDA devices
python inference.py --list-devices

Python API

import torch
from inference import BSRNN, load_audio, extract_speaker, load_checkpoint, build_model

# Build model and load weights
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = build_model(device)
load_checkpoint("checkpoint_epoch037.pt", model, device)

# Load audio (16 kHz mono)
mix = load_audio("mixture.wav")
enroll = load_audio("enrollment.wav")

# Run extraction
extracted = extract_speaker(model, mix, enroll, device)

# Save result
torchaudio.save("result.wav", extracted, 16000)

Advanced: Manual Model Loading

import torch
import torchaudio
from inference import BSRNN

# Build model with exact training config
model = BSRNN(
    feat_type="consistent",
    feature_dim=128,
    num_repeat=6,
    spk_emb_dim=192,
    spk_fuse_type="multiply",
    multi_fuse=False,
    spk_model="ECAPA_TDNN_GLOB_c512",
    sr=16000, win=512, stride=128,
    spk_args={"feat_dim": 80, "embed_dim": 192, "pooling_func": "ASTP"},
    spk_model_freeze=True,
    use_spk_transform=False,
    joint_training=True,
    multi_task=False,
    spk_feat=False,
)
model.eval()

# Load checkpoint
ckpt = torch.load("checkpoint_epoch037.pt", map_location="cpu")
state_dict = ckpt["model"] if "model" in ckpt else ckpt
model.load_state_dict(state_dict, strict=False)

# Run inference
mix, sr = torchaudio.load("mixture.wav")
enroll, sr = torchaudio.load("enrollment.wav")
with torch.no_grad():
    extracted, _ = model(mix, enroll)

What's in this repository

The training code, dataset, and full evaluation pipeline are available at:

• Training code: GitHub - TaurenMountain/real-t (see opensource/code/)
• Dataset: TaurenMountain/REAL-PS4

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