wq2012's picture
v0.5.0: SOTA upgrade, MUSAN+RIR+speed shards, BatchNorm ECAPA, AS-Norm, TFLite exports, and full-data flagships
11ab4ef verified
|
Raw History Blame Contribute Delete
11.7 kB
metadata
license: apache-2.0
library_name: flaxspeaker
pipeline_tag: feature-extraction
tags:
  - audio
  - speaker-verification
  - speaker-recognition
  - speaker-embeddings
  - jax
  - flax
  - safetensors
  - conformer
  - extended_set_softmax
  - cosine
datasets:
  - librispeech_asr
  - cnceleb
  - voxceleb
metrics:
  - eer
  - min_dcf
  - roc_auc

FlaxSpeaker: ls_conformer_small_ext_softmax (conformer - small)

Disclaimer & Reproduction Notice

This is not an officially supported Google product.

This checkpoint and the accompanying FlaxSpeaker library are a from-scratch re-implementation in JAX/Flax, written using only the published papers (no proprietary code or data), of speaker recognition algorithms from:

  • GE2E: Generalized End-to-End Loss for Speaker Verification (arXiv:1710.10467)
  • Dr-Vectors: Decision Residual Networks and an Improved Loss for Speaker Recognition (arXiv:2104.01989)
  • ATP (pooling head repurposed from a language identification paper): Attentive Temporal Pooling for Conformer-based Streaming Language Identification in Long-form Speech (arXiv:2202.12163)
  • PFAS: Parameter-Free Attentive Scoring for Speaker Verification (arXiv:2203.05642)

The experiments in those papers were run on Google-internal infrastructure and data, not with this library. This is a converged small-corpus benchmark checkpoint (2,400 steps / ~13.27 epochs on 5,788 utterances from 232 speakers), not a full-corpus reproduction of the papers' large-scale results -- see Training budget and Deviations from the source papers below.

Training Budget (read this first)

This model was trained to convergence for 2,400 gradient steps (S_safe = X + 2 epochs safety margin, warmup 240 steps) with mini-batches of 32 utterances (8 speakers x 4 utterances), i.e. ~13.27 epochs over the 5,788-utterance LibriSpeech training pool (232 speakers), on a CPU workstation with seed 42. Across the training trajectory, in-domain LibriSpeech EER progressed from step 400: 24.90% (1-utt) / 22.20% (3-utt) -> step 2000 (knee): 14.27% (1-utt) / 9.20% (3-utt) -> step 2400 (final): 14.00% (1-utt) / 9.20% (3-utt) (Universal Pooled final: 22.32% 1-utt / 18.23% 3-utt). While training loss and validation EER have reached their small-corpus convergence plateau on this 232-speaker open split (<= 25 utterances/speaker), absolute EERs remain bounded by training speaker diversity compared to industrial benchmarks trained on thousands of speakers (see below).

Training-crop provenance. This checkpoint was trained with U(90, 180)-frame variable-length crops and an exact validity mask (the current pipeline default); the shipped config.json / YAML (variable_length_training: true) describes the released weights. Step-0 loss fingerprint at the released seed (scripts/fingerprint_released_checkpoints.py, pretrained_models/step0_fingerprints.json): archived 2.4212; current code with fixed 180-frame crops 2.1592, with U(90, 180) crops 2.4212 (match: variable). Evaluating the released weights with the current code reproduces the archived metrics in eval_results.json. At the converged S_safe budget, the seed-to-seed standard deviation of pooled 1-utt EER across seeds 0-2 is 0.08-0.41 % absolute (pretrained_models/r2_seed_and_control_results.json), while the 8,000-trial 95 % bootstrap CI half-width is ~0.8 %.

For orientation only: published ECAPA-TDNN systems trained on VoxCeleb2-dev reach about 0.9% EER on the full VoxCeleb1-O list (Desplanques et al., 2020), and sub-1M-parameter systems about 2% (Wang et al., EfficientTDNN, 2022). The checkpoints released here are trained to convergence (including a +2-epoch safety margin) on the self-contained 381-speaker (<=25 utterances/speaker) benchmark corpus rather than VoxCeleb2-dev (5,994 speakers, ~1.1M utterances), and are intended for controlled relative comparisons within the FlaxSpeaker benchmark, not as state-of-the-art models trained on thousand-speaker corpora.

Model Summary

Property Value
Model ID ls_conformer_small_ext_softmax
Hugging Face Repository wq2012/ls_conformer_small_ext_softmax
Hugging Face Collection FlaxSpeaker Collection
GitHub Repository wq2012/FlaxSpeaker
PyPI Package flaxspeaker
Backbone Architecture conformer (small)
Temporal Pooling atp
Training Objective Blended Extended-Set Softmax (FlaxSpeaker variant: 0.50 row-softmax + 0.50 extended-set with a -log N logit shift; not the paper's Eq. 2) (extended_set_softmax)
Verification Scorer Cosine similarity (cosine)
Output Representation Dim 192
Inference Parameters (encoder + scoring head; excludes training-only classifier weights) 2,029,894
Training Data LibriSpeech only (232 speakers, 5,788 utterances)
Training Steps / Batch 2,400 steps x 32 utterances (~13.27 epochs), seed 42
Audio Frontend 80-dim natural-log Mel @ 16000 Hz + frame-level energy VAD + CMVN (flaxspeaker.AudioFrontend; not a Hugging Face AutoFeatureExtractor)

Model Files Included

  • model.safetensors: model weights (encoder + GE2E w,b / Dr-Vector / PFAS heads) in SafeTensors format
  • flax_model.msgpack: native Flax Msgpack serialized parameters
  • config.json: FlaxSpeakerConfig (full ExperimentConfig, incl. the frontend/VAD settings)
  • eval_results.json: multi-domain evaluation metrics (1-utt & 3-utt EER, minDCF, ROC-AUC, 95% bootstrap CIs, VoxCeleb1-O 10-speaker subset)

Quickstart Usage

Python API (from Hugging Face Hub or Local Directory)

Requires flaxspeaker >= 0.5.0 (pip install -U flaxspeaker).

from flaxspeaker import load_pretrained

model = load_pretrained("wq2012/ls_conformer_small_ext_softmax")
emb = model.extract_embedding("path/to/utterance.flac")
score = model.verify("path/to/enroll.flac", "path/to/test.flac")
print(f"Embedding shape: {emb.shape}, Verification score: {score:.4f}")

Command-Line Inference

python scripts/inference.py --model_dir wq2012/ls_conformer_small_ext_softmax \
    --enroll_audio path/to/enroll.flac --test_audio path/to/test.flac

Multi-Domain Evaluation Results

1-Utterance Verification Across Evaluation Domains

Trial lists are 50/50 target/impostor over held-out speakers (LibriSpeech 40, CN-Celeb2 30, VoxCeleb1 10 speakers); minDCF@0.01 is close to its trivial ceiling of 1.0 for every model under this trial mix.

Evaluation Split Trials Raw EER (%) AS-Norm EER (%) 95% Bootstrap CI (%) minDCF@0.01 minDCF@0.10 ROC-AUC
LibriSpeech (English Read) 3,000 14.00% 14.43% [12.78, 15.20] 0.9427 0.6953 0.9323
CN-Celeb2 (Chinese Multi-Genre) 3,000 29.40% 29.43% [27.68, 30.98] 0.9973 0.9660 0.7732
VoxCeleb1 (In-the-Wild Celebrity) 2,000 22.80% 22.50% [20.97, 24.73] 0.9240 0.9220 0.8619
Universal Pooled (All 3 Domains) 8,000 22.32% 22.80% [21.42, 23.20] 0.9922 0.8963 0.8585
Universal Macro-Average (3 Domains) — 22.07% 22.12% — 0.9547 0.8611 0.8558
VoxCeleb1-O 10-speaker subset (see caveat) 2,500 24.52% 24.51% [22.21, 26.09] 0.9799 0.9205 0.8389

Caveat: VoxCeleb1-O 10-speaker subset: the official veri_test2.txt list filtered to trials whose enrollment and test speakers are both among the 10 held-out evaluation speakers (2,500 trials, ~72% target). Not comparable to published VoxCeleb1-O numbers; the other 30 of the 40 official test speakers are used as training data by vox_* and univ_* models.

3-Utterance Multi-Enrollment Verification Across Evaluation Domains

Evaluation Split Trials 3-Utt EER (%) 95% Bootstrap CI (%) minDCF@0.01 ROC-AUC
LibriSpeech (3-Utt) 1,000 9.20% [7.20, 11.05] 0.8980 0.9709
CN-Celeb2 (3-Utt) 1,000 25.80% [23.25, 28.70] 0.9720 0.8182
VoxCeleb1 (3-Utt) 1,000 20.40% [17.60, 22.25] 0.9800 0.8902
Universal Pooled (3-Utt) 3,000 18.23% [16.93, 19.64] 0.9640 0.8991
Universal Macro-Average (3-Utt) — 18.47% — 0.9500 0.8931

Single seed; the 95% bootstrap CIs resample trials (not speakers). Differences below about 1% absolute EER between FlaxSpeaker checkpoints are within the seed-to-seed variation measured in pretrained_models/r2_seed_and_control_results.json (see the GitHub README) and should not be interpreted as rankings.

Deviations from the Source Papers

  • Training regime: 2,400 gradient steps x 32 utterances (N = 8 speakers x M = 4, ~13.27 epochs) on <= 25 utterances/speaker from 232 speakers (5,788 utterances); the source papers train for ~500k steps with N = 16-64 / M = 8-10 on thousands of speakers.
  • Augmentation: corpus-backed MUSAN environmental noise, music, and multi-speaker babble + RIRS_NOISES room impulse responses + 3-way polyphase speed perturbation (0.9x, 1.0x, 1.1x) + SpecAugment.
  • Objective: 50/50 blend of row-softmax and -log N-shifted extended-set softmax; the paper's pure Eq. 2 is available via pure_extended_set: true.
  • Conformer: parallel linear + Conv1D stem, no positional encoding, 3 blocks x 192-D, 4 heads, kernel 15, mask-aware GroupNorm(2), no dropout, multi-layer feature aggregation; the Google papers use 12 x 512 blocks, 4-8 heads, kernel 32, relative positional encoding and dropout.
  • ATP pooling: 4 heads with a tanh bottleneck (128-D) and subspace_heads = True; the papers use a single head with a linear-sigmoid weight (attention_dim = 0). Note that arXiv:2202.12163 is a language-identification paper whose pooling head is repurposed here.
  • Frontend: 80-D natural-log Mel + frame-level energy VAD (26 dB dynamic range) + CMVN; evaluation processes up to 3000 frames (30.0 s).

Citation

@inproceedings{wan2018generalized,
  title={Generalized End-to-End Loss for Speaker Verification},
  author={Wan, Li and Wang, Quan and Papir, Alan and Moreno, Ignacio Lopez},
  booktitle={IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  pages={4879--4883},
  year={2018}
}

@inproceedings{pelecanos2021dr,
  title={Dr-Vectors: Decision Residual Networks and an Improved Loss for Speaker Recognition},
  author={Pelecanos, Jason and Wang, Quan and Moreno, Ignacio Lopez},
  booktitle={Interspeech},
  pages={4603--4607},
  year={2021}
}

@inproceedings{wang2022attentive,
  title={Attentive Temporal Pooling for Conformer-based Streaming Language Identification in Long-form Speech},
  author={Wang, Quan and Yu, Yang and Pelecanos, Jason and Huang, Yiling and Moreno, Ignacio Lopez},
  booktitle={The Speaker and Language Recognition Workshop (Odyssey)},
  year={2022}
}

@inproceedings{pelecanos2022parameter,
  title={Parameter-Free Attentive Scoring for Speaker Verification},
  author={Pelecanos, Jason and Wang, Quan and Huang, Yiling and Moreno, Ignacio Lopez},
  booktitle={The Speaker and Language Recognition Workshop (Odyssey)},
  year={2022}
}