scribeair-diar-pruned-156m

ScribeAir Diar Pruned 156M

A structurally-pruned speaker diarization pipeline based on DiariZen (WavLM-Large + Conformer

• VBx clustering). The pruned variant has 156M parameters — half of the original 315M — while retaining most of the accuracy on meeting-style data.

Built for the ScribeAir desktop speech-to-text app where the smaller footprint matters for CPU/edge deployment.

Architecture

The Conformer head, weighted sum of WavLM layers, VBx clustering and PLDA are unchanged from the reference DiariZen pipeline.

Benchmark

Aggregate DER on an 18-file meeting corpus (AMI Array1 / NOTSOFAR / Russian podcasts), collar=0, with overlap, AHC threshold = 0.6:

(ahc_threshold = 0.75 in config.toml was tuned via a 5-value sweep on the bench corpus; it gives a small but consistent −0.05 pp over the DiariZen-default 0.60.)

The pruned model beats the full-size own-trained baseline by 0.79 pp and ends up just 1.71 pp behind the upstream teacher while shipping half the weights.

Training recipe

1. Structured pruning at sparsity = 0.5 on the WavLM-Large encoder using the recipe from Han et al., Fine-tune Before Structured Pruning. Distill+prune trained with pruning_units = "conv,head,interm", distill_lr = 2e-4, reg_lr = 2e-2, 5-epoch warmup, 8 epochs total.
2. Dual-optimiser continued training of the full EEND-Conformer model on the pruned backbone (run_dual_opt_continued_training.py, optimizer_small_lr = 2e-5, optimizer_big_lr = 1e-3, 7 epochs).
3. Seed-net-style data cleaning: the upstream DiariZen v1 teacher was used to re-label files where it disagreed with the ground truth by

   25 % DER (35 / 242 files), producing train_clean.

4. Clean fine-tune on train_clean from the dual-opt best checkpoint (1 epoch at the same LR; further epochs over-fit).
5. Low-LR continuation at LR / 10 (2e-6 / 1e-4) for 1 more epoch.
6. Weight averaging between the clean-e1 and low-LR-best checkpoints at α = 0.75 (target side weighted), shipped as pytorch_model.bin.

Training data

Only permissively-licensed meeting data:

AliMeeting is not used due to gated license terms, so this checkpoint is safe for commercial use.

License

MIT for the trained weights in this repository. The pipeline code lives in BUT-FIT's DiariZen (MIT) and Microsoft's WavLM (MIT). Embedding extractor used at inference time is pyannote/wespeaker-voxceleb-resnet34-LM (MIT).

Files

• pytorch_model.bin — full pipeline state dict (Conformer + pruned WavLM)
• wavlm.bin — pruned WavLM weights + architecture config (loaded first to build the WavLM sub-module with the right pruned shapes)
• config.toml — pipeline configuration for DiariZenPipeline
• plda/{plda,xvec_transform}.npz — VBx PLDA artefacts (borrowed from the upstream BUT-FIT DiariZen v1 release)

Usage

pip install git+https://github.com/BUTSpeechFIT/DiariZen.git

from huggingface_hub import snapshot_download, hf_hub_download
from diarizen.pipelines.inference import DiariZenPipeline

# Pull this repo
local_dir = snapshot_download(repo_id="coriollon/scribeair-diar-pruned-156m")

# WeSpeaker embedding extractor (separate, MIT-licensed)
emb_path = hf_hub_download(
    repo_id="pyannote/wespeaker-voxceleb-resnet34-LM",
    filename="pytorch_model.bin",
)

pipe = DiariZenPipeline(diarizen_hub=local_dir, embedding_model=emb_path)
diarization = pipe("meeting.wav")
print(diarization)  # pyannote.core.Annotation

Caveats

• Tested on conversational English (AMI) and lightly on Russian podcasts. No fine-tuning on AliMeeting or other Chinese meetings — Mandarin performance falls back to whatever AISHELL-4 + WavLM-Large captures pre-pruning.
• The pipeline expects 16 kHz mono input. Multi-channel audio is collapsed to channel 0.
• wavlm.bin is loaded before pytorch_model.bin to instantiate the pruned encoder architecture, so do not edit / replace it.

Citation

If you use this checkpoint, please cite the upstream papers:

@article{han2025fine,
  title={Fine-tune Before Structured Pruning: Towards Compact and Accurate
         Self-Supervised Models for Speaker Diarization},
  author={Han, Jiangyu and Landini, Federico and Rohdin, Johan and
          Silnova, Anna and Diez, Mireia and Cernocky, Jan and Burget, Lukas},
  journal={Proc. INTERSPEECH},
  year={2025}
}

@article{chen2022wavlm,
  title={WavLM: Large-Scale Self-Supervised Pre-Training for Full Stack
         Speech Processing},
  author={Chen, Sanyuan and Wang, Chengyi and Chen, Zhengyang and Wu, Yu
          and Liu, Shujie and Chen, Zhuo and Li, Jinyu and Kanda, Naoyuki
          and Yoshioka, Takuya and Xiao, Xiong and others},
  journal={IEEE Journal of Selected Topics in Signal Processing},
  year={2022}
}

Acknowledgements

Thanks to BUT-FIT (Han et al.) for the DiariZen pipeline and the pruning recipe, and to Microsoft for releasing WavLM under MIT.