BiMamba2 Masked Discrete-Unit Prediction for Multilingual Speech Representation for Unsupervised Speech in the Wild Challenge

1. Introduction

• Many languages lack the labeled data needed for supervised pretraining, while speech models show performance disparities across languages and acoustic conditions.The UPS Challenge evaluates language identification, speech recognition, and speaker clustering on held-out audio without task-specific fine-tuning.Source: §1 · p. 1
• The submission adapts HuBERT-style masked discrete-unit prediction to a bidirectional Mamba2 encoder and trains on about 250 hours across 67 languages.Its objectives combine discrete-cluster prediction, VICReg regularization, and language identification loss.Source: §1 · p. 1
• The paper contributes a multilingual masked-unit model, a shard-based training pipeline with VAD filtering and language-aware sampling, and an evaluation-mismatch analysis.The stated implications of the mismatch analysis concern reproducibility for future challenge participants.Source: §1 · p. 1

2.1. Dataset

• The corpus uses approximately 87,400 utterance chunks from 438 shards, covering about 250 hours and 67 languages, with metadata supplied by MLCommons.The data includes roughly 100 hours of English and 150 hours of non-English speech; no external data or local VAD/LID inference is used.Source: §2.1 · p. 1

2.2. Filtering and Feature Extraction

• Clips pass VAD-based speech-ratio, segment-length, and extractable-speech thresholds before offline 80-bin log-mel features are stored in precomputed shards.At dataloader time, clips also need a speech ratio of at least 0.25 and a VAD segment long enough for a full 10-second chunk.Source: §2.2 · pp. 1–2

2.3. Language-Aware Batch Sampling

• Batches draw 70% of examples from language-balanced sampling and 30% from uniform shard sampling, with English capped at 10%.The strategy addresses English's approximately 40% share of utterances after filtering.Source: §2.3 · p. 2

3.1. Overview

• The model predicts k-means cluster targets at masked positions using a BiMamba2 encoder, with VICReg and language-identification losses on mean-pooled states.A linear head projects encoder outputs to pseudo-label logits; the model has no separate decoder.Source: §3.1 · p. 2

3.2. BiMamba2 Encoder Layer

• Each BiMamba2 layer sums forward and reversed Mamba2 paths with a linear skip path, yielding bidirectional sequence representations.The Mamba2 paths use structured state-space duality, local depthwise convolution, and a gated recurrence; the submitted model has 12 layers and 47.88M parameters.Source: Figure 1 · p. 2, §3.2 · p. 2

3.3. Masking and Pseudo-Labels

• Masking covers at least one valid frame and about 75% of valid frames in 3–5 contiguous spans, replacing them with a learned embedding.Pseudo-labels come from MiniBatchK-Means with k=200 on per-dimension-scaled 80-bin log-mel frames, without a pretrained model.Source: §3.3 · p. 2

3.4. Training Objectives

• The primary loss is cross-entropy on masked-frame pseudo-labels, combined with VICReg variance and covariance penalties and weighted language-ID cross-entropy.The VICReg penalties act on mean-pooled hidden states; the language-ID loss uses precomputed UPS labels and has weight 0.05.Source: §3.4 · p. 2

4. Training

• AdamW [23] optimization used a 10^-4 learning rate and 0.5 gradient-clip max norm, with 2,000-step linear warmup followed by cosine decay to 0.1 times the base learning rate.Training resumed from a step-10,500 checkpoint of the same configuration and ran to step 48,000.Source: §4 · p. 2
• Training used bfloat16 mixed precision, batch size 128, seed 42, and a single NVIDIA RTX A6000 GPU.Input frames and pooled VICReg embeddings were clamped for numerical stability, and hidden states passed through nan-to-num.Source: §4 · pp. 2–3

5. Challenge Submission

• Frame-level encoder outputs are averaged across three non-overlapping start, middle, and end segments, then L2-normalized as embeddings.The outputs come from the post-final-LN encoder.Source: §5 · p. 3
• The frontend resamples audio to 16 kHz, extracts 80-bin log-mel features, and pads or crops them to a fixed 10-second window.It uses FFT 400 and hop 160.Source: §5 · p. 3
• Checkpoint loading uses conservative fallbacks across runtimes, while BiMamba2 inference requires its CUDA backend.Source: §5 · p. 3

6.1. Official evaluation.

• Official Dynabench evaluation measures language identification macro-F1, speech-recognition CER, and speaker-clustering ARI without giving participants probe training details or test labels.Challenge baselines are Whisper [24], HuBERT-large, XLSR, and wav2vec 2.0.Source: §6.1 · p. 3

6.2. Smaller model configuration.

• The smaller submission uses dmodel = 512, 8 layers, batch size 64, and step 49,000 with the primary model’s objectives and pipeline.Source: §6.2 · p. 3

6.3. Local evaluation.

• Local holdout validation and post-hoc linear probes use a 45-language split; shard-level exclusion prevents recording overlap, but neither metric is a valid proxy for official performance.Holdout validation covers 6,646 samples, and the linear probe uses precomputed embeddings.Source: §6.3 · p. 3

7.1. Official Challenge Results

• Step 48,000 yields macro-F1 0.070, ARI 0.710, and CER 0.998; the smaller d=512, 8-layer model performs below the primary model on all three tasks.Source: §7.1 · p. 3
• ARI 0.735 exceeds all four baselines; macro-F1 0.073 trails Whisper (0.950) and HuBERT-large (0.690), while CER 0.870 falls between XLSR (0.900) and HuBERT-large (0.570).Official evaluation ranks step 19,500 above step 48,000 on all three metrics, making it the primary result.Source: Table 2 · p. 3, §7.1 · p. 3

7.2. Local Diagnostic Evaluation

• Locally, step 48,000 leads on macro-F1 (0.395 vs. 0.342) and ARI (0.232 vs. 0.225), reversing the official ranking, especially on CER (0.870 vs. 0.998).The local holdout ranking did not reliably predict which checkpoint generalized better to official evaluation.Source: §7.2 · p. 3
• At step 48,000, inter-class cosine similarity is −0.002 and the separability ratio is −366.84; at step 19,500, they are 0.004 and 208.98.The negative ratio at step 48,000 arises from negative mean inter-class cosine similarity.Source: §7.2 · p. 3, Table 3 · p. 3
• Post-hoc linear-probe macro-F1 reaches 0.915, which the authors attribute to the expressive capacity of a trained probe on in-distribution data.Source: §7.2 · p. 3

8.1. Speaker Clustering vs. LID and ASR

• ARI exceeds all baselines while macro-F1 is 0.073; the proposed speaker-discriminative explanation is post-hoc and untested.The authors suggest masked frame recovery may preserve speaker-related acoustic traits, while language coverage limits LID generalization.Source: §8.1 · p. 4

8.2. Late-Stage Embedding Geometry and CER Degradation

• Official ARI stays relatively stable (0.735 vs. 0.710), but CER worsens from 0.870 to 0.998 between steps 19,500 and 48,000.The authors suggest late training reorganizes embedding geometry in a way that may not preserve phonetic similarities needed by CER; this account is a hypothesis, and controlled experiments are needed.Source: §8.2 · p. 4

8.3. Bidirectional Local-vs-Official Mismatch

• Local scores diverge from official scores in both directions, and the locally favored step-48,000 checkpoint ranks worse on all three official metrics.At step 19,500, local and post-hoc probe macro-F1 are 0.342 and 0.915 versus official 0.073; local ARI is 0.225 versus official 0.735.Source: §8.3 · p. 4
• All 45 local evaluation languages occur in training, making the local LID scores optimistic for generalization.The discussion also attributes mismatch to differing clustering methods and evaluation sets; local clusters align with language groups rather than speakers.Source: §8.3 · p. 4
• Valid local proxies should use languages absent from training, evaluation data disjoint by construction, and probes matched to official methodology.Source: §8.3 · p. 4

8.4. Effect of Model Scale

• Scaling from d=512/8L to d=768/12L improves official ARI from 0.291 to 0.735 and CER from 0.996 to 0.870, with macro-F1 rising from 0.052 to 0.073.Because batch size and total training steps also differ, the gains cannot be attributed to model scale alone; controlled ablations were not conducted.Source: §8.4 · p. 4

• The model combines masked k-means pseudo-label prediction with auxiliary LID supervision and VICReg regularization, trained through a shard-based pipeline with language-aware sampling.Source: §9 · p. 4
• The 47.88M-parameter primary checkpoint achieves official ARI 0.735, exceeding all four published speaker-clustering baselines, while macro-F1 0.073 and CER 0.870 remain below supervised baselines.Source: §9 · p. 4
• Local evaluation overestimates LID, underestimates speaker clustering quality, and does not reliably rank checkpoints as official evaluation does.Source: §9 · p. 4
• The study is limited by absent component ablations, single-run results without variance estimates, a smaller training corpus than the baselines, and one-time pseudo-labels without iterative refinement.Source: §9 · p. 4

• Generative AI assisted with outlining and readability improvements in parts of the manuscript.Source: §12 · p. 5
• The authors state they are fully responsible for the manuscript content and checked its claims, numbers, and citations.Source: §12 · p. 5
• Generative AI was not used to produce experimental results.Source: §12 · p. 5