Fake-Mamba / implementation-notes.md
korallll's picture
Wrapper code (speech-spoof-bench)
632131a verified
|
Raw History Blame Contribute Delete
4.11 kB

Fake-Mamba — implementation notes

Arena wrapper for Fake-Mamba (PN-BiMamba, Xuan et al., IEEE ASRU 2025, arXiv 2508.09294), built 2026-10-30 under the 6-model Mamba batch plan (speech-spoof-bench/docs/plans/2026-10-29-mamba-batch-arena-submission.md).

Source

  • Code: https://github.com/xuanxixi/Fake-Mamba @ e966d2f317b187f474d8b2d54b7e0f9d0a124ba0 — no LICENSE.
  • Checkpoint: HF xxuan-speech/fake-mamba-models @ cd99e0efb4d4f3f6dfe610ae8e27de9a73d6bf41, models/Exp65/best.pth (1.29 GB). Local: fake_mamba_exp65_best.pth → symlink to ../_ckpt_dl/xxuan-speech__fake-mamba-models/models/Exp65/best.pth. Not re-hosted (no license); meta.yaml checkpoint is the commit-pinned authors' URL.
  • best.pth is byte-identical (md5 4601685f…) to best/best_0.pth (lowest LA-dev loss).
  • 12 PN-BiMamba blocks (encoder.layers.0-11), paper/source default is 7 → built with depth=12.
  • XLS-R arch: xlsr2_300m_cfg.json = the stored model cfg of fairseq xlsr2_300m.pt (symlink xlsr2_300m.pt → ../W2V2-AASIST/xlsr2_300m.pt, used only to extract it), merged onto the installed Wav2Vec2Config (merge_with_parent, cross-version fix from W2V2-AASIST). The 3.6 GB base checkpoint is not loaded at runtime; the fine-tuned checkpoint overwrites every weight.

Port (_net.py)

  • Only Fake_Mamba + PN_BiMambas(_Encoder) ported; the missing modules imported by the source (conformer.ConformerBlock/BiMambaEncoder, ConExBiMamba.ExBiMamba) belong to other back-ends.
  • Mamba = mamba_ssm 1.x Mamba (d_model 144, d_state 16, d_conv 4, expand 2; checkpoint keys match), non-fast-path branch (nn.Conv1d + selective scan). No RMSNorm anywhere (LayerNorm only).
  • Selective scan: selective_scan_parallel — chunked (8) Hillis-Steele prefix scan + 26-step inter-chunk carry; equals the library's sequential selective_scan_ref (kept, FAKE_MAMBA_SCAN=ref) to ~4e-6 on encoder outputs, 2e-7 on scores. ~3.6× faster than the 208-step loop.
  • Forward + time-flipped sequence go through the shared Mamba as one batch (same math).
  • extract_feat's self.model.train() on device mismatch removed; wrapper does .to() then .eval().

Wrapper (fake_mamba.py:FakeMamba)

  • First 66,800 samples, repeat-tiled if shorter (utils.pad), no resampling/normalisation, FP32.
  • Score = logits[:, 1] (source labels bonafide = 1; eval scripts write out[:, 1]).
  • params: 321,290,117 (321.2901 M).

Parity with the authors' In-the-Wild scores (ref_scores/…/best/SITW/score.txt, label 1 = spoof)

  • Authors' file → EER 5.849% (paper 5.85%).
  • Our best.pth on the full arena In-the-Wild parquet (31,779): EER 5.97%, per-clip Pearson r = 0.989 vs the authors' scores.
  • The authors' SITW-DF_test.py defaults to --average_model True --n_average_model 5 (weight-average of best/best_0..4.pth). Averaging those 5 files moves us to r = 0.999 vs their file, full-ITW EER 5.82% (median |Δ| 0.13 logit; not exact — their best/ set or environment differed slightly). Checked on 256/128 clips; the 16-bit FLAC arena audio vs the original float wavs makes no difference (r identical). CPU vs GPU, TF32 on/off, ref vs parallel scan: identical.
  • Decision (plan): ship the single best.pth (Exp65) as approved; avg-5 noted as an option.

Results (2026-10-30; both GPUs were shared with other unlocked jobs, so throughput is a lower bound)

RTX 4070 Ti SUPER (idx 2) RTX 5060 Ti (idx 3)
PyTorch sweep (bs 4) 34.6 utt/s 17.2 utt/s
TRT FP32 (FP16 Spearman 0.998 < 0.9999 → FP32) 60.2 utt/s @24 38.8 utt/s @16
TRT build parity (64 ITW clips) Spearman 0.99995 0.99995
J-SPAW_LA EER, TRT vs PyTorch 15.000 15.000 (Δ0), Spearman 0.999996 14.966 (Δ0.034), Spearman 0.999996

PyTorch J-SPAW_LA run (results_pt) was on the RTX 5060 Ti. Both engines are accepted.

  • 2026-10-30 (maintainer): weights switched to the authors' exact eval average (2*b1+b2+b3+b4)/5, averaged at load time from the pinned HF files; ITW r=0.99996 vs authors' scores. Old best.pth outputs in old_bestpth/.