| # DPDA-GA-Mamba — implementation notes (2026-10-30) |
|
|
| Same network code is shipped in three self-contained dirs: `XLSR-Mamba-LA/`, `XLSR-Mamba-DF/`, |
| `DPDA-GA-Mamba/` (`_net.py` and `_mamba_v1.py` byte-identical; only the wrapper module / class / |
| checkpoint differ). Not to be confused with `XLSR-MamBo/` (arXiv 2601.02944, a different model). |
|
|
| ## Source |
| - Code: github.com/swagshaw/XLSR-Mamba (MIT) `model.py` + `mamba_blocks.py`. |
| DPDA-GA `nets/mamba.py` + `nets/mamba_blocks.py` are byte-identical (md5 checked). |
| - Training recipe: DPDA-GA (Truong et al., ICASSP 2026, arXiv 2509.20682; github.com/ductuantruong/dpda_ga, |
| MIT): dual-path data-augmented training (each batch = clean + RawBoost-augmented copies) with |
| PCGrad gradient alignment between the two paths' losses, on the XLSR-Mamba architecture. |
| - Checkpoints: HF `tuantd/dpda_ga` @ 531554e0a1c28d6a6ccf22a355281e3d3acefeb5 |
| `pcgrad/xlsr_mamba/best/best_{0..4}.pth` (565 tensors each, same keys as XLSR-Mamba). |
| - **Weight average (`make_avg5.py` -> `dpda_ga_mamba_avg5.pth`)** — maintainer decision 2026-10-30: |
| reproduce main.py (`--average_model True --n_average_model 5`) EXACTLY, bug included. |
| main.py does `sd = model.state_dict()` after loading best_0; state_dict() tensors share |
| storage with the parameters, so the next `load_state_dict(best_1)` overwrites sd in place |
| before the first `sd[key] + sd2[key]` (which then rebinds sd to fresh tensors). Result: |
| **(2·best_1 + best_2 + best_3 + best_4) / 5**, best_0 dropped (verified on torch 2.7 with a toy |
| Linear: weights 0..4 -> 2.2, not 2.0; and on the real file, max |diff| to that formula 1e-7). |
| This is what the authors evaluated. Float tensors summed in float64, stored float32; |
| `first_bn.num_batches_tracked` rounded (unused at eval). `make_avg5.py --truemean` gives the |
| plain 5-way mean instead (not shipped). |
| - Hyper-parameters (main.py defaults, confirmed by the checkpoint shapes): emb_size 144, |
| num_encoders 12 -> 6 forward + 6 backward Mamba blocks; ssm_cfg {} -> d_state 16, d_conv 4, |
| expand 2 (d_inner 288, dt_rank 9); rms_norm, fused_add_norm, residual_in_fp32 all True. |
| |
| ## Parameter count |
| 319,327,542 parameters + 3 BatchNorm buffer elements = 319,327,545 state_dict elements |
| (the 319,327,545 in the brief / HF header counts the buffers). `params_millions: 319.3275`. |
|
|
| ## Edits to the source (all inference-only; listed in `_net.py` header) |
| 1. XLS-R built from `xlsr2_300m.pt` cfg via `merge_with_parent` (W2V2-AASIST pattern; symlink |
| to `../W2V2-AASIST/xlsr2_300m.pt`); the checkpoint overwrites every weight (strict load). |
| 2. `extract_feat`: removed `self.model.train()` (would re-enable dropout on a device/dtype |
| mismatch) and the duplicate XLS-R forward whose `layer_results` was discarded. |
| 3. `Model(emb_size, num_encoders)` instead of argparse args. |
| 4. Kernels: mamba_ssm / causal_conv1d are not installed, and Triton is unreliable on the |
| RTX 5060 Ti (sm_120; XLSR-MamBo hit "computeCapability not supported"). `_mamba_v1.py` |
| vendors the needed mamba_ssm v1 subset (from benchmarks/RawBMamba/mamba_ssm 1.1.3.post1) |
| as ONE pure-PyTorch file (not a `mamba_ssm` package -> no sys.modules clash with other |
| dirs): Mamba (slow path: depthwise causal conv1d + SiLU, x_proj/dt_proj, selective scan, |
| z-gating; == the fused `mamba_inner_fn` math), Block, RMSNorm, fused add+RMSNorm |
| (`rms_norm_fn`, `rms_norm_ref` semantics incl. prenorm/residual_in_fp32). |
| **Pure torch is used for every score on every card.** |
| 5. Selective scan = `selective_scan_chunked`: exact closed form of the recurrence in chunks of |
| 8 (intra-chunk masked-exp sums + inter-chunk sums over chunk totals; every exp argument |
| is a masked difference <= 0 -> no overflow). No Python loop -> fast eager and a compact |
| ONNX graph. Checked against the sequential `selective_scan_ref` (kept in the file): |
| random inputs max |diff| 2e-4 on |y|~128 (rel 2e-6); full-model scores within 1e-4 |
| (`test_chunked_scan_matches_reference`). |
| 6. Only the released path of MixerModel is ported (bidirectional, fused_add_norm). Its quirk |
| is preserved verbatim: the final fused add+norm of each direction adds the *stack input* |
| `hidden_states` to that direction's residual (the last mixer output is not added), and the |
| backward branch's residual stays in flipped time order. BiBlock/BiMixerModel (unused) dropped. |
|
|
| ## Inference conventions |
| - First 66,800 samples; shorter clips repeat-tiled (source `utils.pad`). No resampling, |
| no normalisation, no RawBoost. FP32. Score = `logits[:, 1]` (bona fide = class 1, |
| source `batch_score = batch_out[:, 1]`). `.to(device)` then `.eval()`. |
|
|
| ## Tests (`test_dpda_ga_mamba.py`) |
| pad truncate / repeat, strict load + eval mode + param/buffer count, batched == single |
| (rtol=atol=1e-4), chunked scan == reference scan on the full net, determinism, finite. |
| |
| ## Results (2026-10-30; GPUs shared with other agents' jobs, so throughputs are lower bounds) |
| - PyTorch sweep (`sweep.py`): RTX 4070 Ti SUPER ~27 utt/s (plateau bs 8-32) -> `batch_size = 16`; |
| RTX 5060 Ti ~15 utt/s at bs 2-4 -> 4 (set in load()). |
| - TensorRT (`trt_<module>.py`, sidecar `trt_<slug>.json`): FP16 failed the ranking check on both cards -> |
| **FP32** engines (TF32 on). best_batch / utt/s: RTX 4070 Ti SUPER 32 / 33.9; RTX 5060 Ti 32 / 31.6. |
| Build parity (64 InTheWild clips) Spearman >= 0.9999 on both cards. |
| - Full J-SPAW_LA (2,397 trials): PyTorch EER 20.000% (`results_pt/`); TRT RTX 4070 Ti SUPER 20.038 (Spearman 1.0000) |
| (`results_trt/`); TRT RTX 5060 Ti 20.100 (Spearman 1.0000) (`results_trt3/`). All |dEER| <= 0.10 -> both engines accepted. |
| - Unit tests: 6/6 pass on both cards. Batched-vs-single uses rtol=atol=1e-4: cuDNN chooses |
| batch-dependent conv algorithms (TF32 on by default), giving ~4e-5 relative jitter on logits ~4. |
| - `setup_lane.sh` (in XLSR-Mamba-LA/) is the driver that produced all of the above. |
|
|
