CLD Convex Low-resource Accent-Robust Language Detection in Speech Recognition A lightweight language-detection module for multilingual ASR, optimized via ADMM in JAX.

---

This repository provides the official implementation of CLD, a lightweight language-detection module for multilingual ASR. This codebase contains our pip-installable Python package (jaxcld/) including our training/benchmark scripts implemented in JAX and optimized via ADMM for high performance in low-resource settings. Simply, the package attaches a small language detection head (Convex NN / small NN / linear SVM) to ASR encoder representations, and use it to select the language token (Whisper) or adapter (MMS) before decoding.

The paper PDF is available in paper/ and the pip-installable package is published at pypi.org/project/jaxcld.

Highlights

• High Accuracy: Excels in binary and multiclass language detection (Table 3).
• Low-Resource Robustness: Effective with limited data (Figures 1 & 2).
• Efficient: 13x training speedup from traditional NNs due to ADMM optimization and JAX.

Requirements

The package is published on PyPI as jaxcld. JAX (CPU) is included by default. For GPU support, install the matching extra:

# CPU
pip install jaxcld

# GPU — CUDA 12 (most modern systems)
pip install "jaxcld[cuda12]"

# GPU — CUDA 11
pip install "jaxcld[cuda11]"

If you've cloned this repo, you can instead install from source:

• Package-only install (inference usage):

pip install -e .                  # CPU
pip install -e ".[cuda12]"        # GPU — CUDA 12

• Full training/benchmark environment (recommended if you run the scripts in this repo):

pip install -e ".[train]"         # CPU
pip install -e ".[train,cuda12]"  # GPU — CUDA 12

If you prefer installing from the pinned dependency list instead:

pip install -r requirements.txt

Using the package

Minimal inference example (Whisper)

import numpy as np

from jaxcld import ASRModel, CVXNNLangDetectHead, NNLangDetectHead, SVMLangDetectHead

# 1) Load the base ASR model
languages = ["en", "hi", "id", "ms", "zh"]
asr = ASRModel.from_pretrained("openai/whisper-small", config={"languages": languages})

# 2) Load a language detection head artifact (choose ONE)
# head = CVXNNLangDetectHead.load("path/to/whisper-small_trained_cvx_mlp.pkl", asr)
# head = NNLangDetectHead.load("path/to/openai_whisper-small_nn_head.pkl", asr)
# head = SVMLangDetectHead.load("path/to/openai_whisper-small_linear_svm.pkl", asr)

# 3) Attach head and run inference
asr.set_lang_detect_head(head)

audio_16k_mono: np.ndarray = ...  # shape (T,), sampling rate 16kHz
pred_langs, pred_texts = asr.predict(audio_16k_mono)
print(pred_langs[0], pred_texts[0])

Pre-trained models

Trained convex heads are published on the Hugging Face Hub:

Model	Backbone	Languages	Det. Acc	WER ↓	CER ↓	HF Hub
cld-whisper-small-5lang	whisper-small	en/hi/id/ms/zh	0.98	48.23	27.47	🤗
cld-whisper-large-v3-5lang	whisper-large-v3	en/hi/id/ms/zh	0.98	31.11	19.81	🤗
cld-mms-1b-5lang	mms-1b-all	en/hi/id/ms/zh	0.96	48.10	23.47	🤗
cld-whisper-small-enzh (100–10000 samples/class)	whisper-small	en/zh	0.99–1.00	—	—	🤗

Loading from the Hub

import numpy as np
from huggingface_hub import hf_hub_download
from jaxcld import ASRModel, CVXNNLangDetectHead

languages = ["en", "hi", "id", "ms", "zh"]

# 1) Load the frozen base ASR model
asr = ASRModel.from_pretrained("openai/whisper-small", config={"languages": languages})

# 2) Download the convex head from the Hub and load it
head_path = hf_hub_download("williamhtan/cld-whisper-small-5lang", "model.pkl")
head = CVXNNLangDetectHead.load(head_path, asr)

# 3) Attach the head
asr.set_lang_detect_head(head)

# 4) Grab one clip from the dataset's test split
from datasets import load_dataset

sample = next(iter(load_dataset("williamhtan/cld-multi-dataset", split="test", streaming=True)))
audio_16k_mono = sample["audio"]["array"]   # (T,) float waveform, 16 kHz mono

# 5) Predict
pred_langs, pred_texts = asr.predict(audio_16k_mono)
print("true:", sample["lang"], "| pred:", pred_langs[0], "| text:", pred_texts[0])

For the low-resource binary model, specify the samples-per-class subfolder:

head_path = hf_hub_download("williamhtan/cld-whisper-small-enzh", "1000/model.pkl")

Data format

All training/evaluation scripts expect a Hugging Face DatasetDict saved to disk (loaded via datasets.load_from_disk(...)) with splits like train, valid, test. Use our data_ingestion.py script to prepare your data.

python data_ingestion.py \
  --config configs/en_hi_config.json \
  --out data/en_hi \
  --common-voice-dir /absolute/path/to/CommonVoice \
  --augment

• Required: --config JSON (see example below), --out save directory.
• Optional: --augment enables audiomentations; --musan-dir for background noise; --common-voice-dir for local Common Voice.
• Output: a saved DatasetDict at data/en_hi with columns: audio, text, lang, accent.

Minimal config example (see more in configs/):

{
  "name": "English-Hindi example",
  "languages": {
    "en": {
      "accents": [
        { "code": "us", "column_name": "United States English", "dataset": "common_voice" }
      ]
    },
    "hi": {
      "accents": [
        { "code": "hi", "column_name": "", "dataset": "common_voice" }
      ]
    }
  },
  "params": {
    "samples_per_class": 1000,
    "split": { "train": 0.8, "val": 0.1, "test": 0.1 }
  }
}

Notes:

• Common Voice selection uses column_name against accents in validated.tsv. Use override_code to point to alternative folders (see configs/final_config.json).
• Lahaja examples match by native_language (e.g., "Telugu", "Konkani").

Training

Train language detection heads

All heads are trained on pooled encoder embeddings extracted by ASRModel.load_data(...) from a dataset on disk.

CVXNN (convex head, JAX + ADMM/CRONOS)

python train_cvxnn.py \
  --model_name openai/whisper-small \
  --dataset_path data/multiclass \
  --languages en,hi,id,ms,zh \
  --output_dir models/lang_heads \
  --neuron 64 \
  --beta 0.001 \
  --rho 0.1 \
  --admm_iters 6

This produces a pickled artifact like:

• models/lang_heads/openai/whisper-small/openai_whisper-small_trained_cvx_mlp.pkl

NN head (PyTorch)

python train_nn.py \
  --dataset_path data/multiclass \
  --model_name openai/whisper-small \
  --languages en,hi,id,ms,zh \
  --output_dir models/lang_heads \
  --num_train_epochs 10 \
  --learning_rate 1e-3 \
  --per_device_train_batch_size 256

This produces a pickled artifact like:

• models/lang_heads/openai/whisper-small/openai_whisper-small_nn_head.pkl

Linear SVM head (sklearn)

python train_linear_svm.py \
  --model_name openai/whisper-small \
  --data_dir data/multiclass \
  --languages en,hi,id,ms,zh \
  --output_dir models/lang_heads \
  --C 1.0 \
  --max_iter 5000

This produces a pickled artifact like:

• models/lang_heads/openai/whisper-small/openai_whisper-small_linear_svm.pkl

Fine-tune Whisper

Use train_whisper.py to fine-tune a Whisper checkpoint on a preprocessed dataset directory:

python train_whisper.py \
  --data_dir data/multiclass \
  --model_id openai/whisper-small \
  --output_dir models/whisper-small-finetuned \
  --num_train_epochs 3 \
  --learning_rate 1e-5 \
  --per_device_train_batch_size 8 \
  --per_device_eval_batch_size 8 \
  --gradient_accumulation_steps 1 \
  --eval_strategy steps \
  --eval_steps 1000 \
  --save_steps 1000

Optional logging:

python train_whisper.py ... \
  --wandb_project CLD \
  --run_name whisper-small-finetune-final_dry

Evaluation

Use benchmark_cld.py to evaluate language detection and transcription quality (WER/CER) on the test split.

Whisper + CVXNN head

python benchmark_cld.py \
  --dataset_path data/multiclass \
  --model_name openai/whisper-small \
  --cld_type cvx \
  --cld_path models/lang_heads/openai/whisper-small/openai_whisper-small_trained_cvx_mlp.pkl \
  --languages en,hi,id,ms,zh \
  --batch_size 32 \
  --no_wandb

Whisper + NN head

python benchmark_cld.py \
  --dataset_path data/multiclass \
  --model_name openai/whisper-small \
  --cld_type nn \
  --cld_path models/lang_heads/openai/whisper-small/openai_whisper-small_nn_head.pkl \
  --languages en,hi,id,ms,zh \
  --batch_size 32 \
  --no_wandb

Whisper + linear SVM head

python benchmark_cld.py \
  --dataset_path data/multiclass \
  --model_name openai/whisper-small \
  --cld_type linear_svm \
  --cld_path models/lang_heads/openai/whisper-small/openai_whisper-small_linear_svm.pkl \
  --languages en,hi,id,ms,zh \
  --batch_size 32 \
  --no_wandb

Whisper vanilla language ID (no head)

python benchmark_cld.py \
  --dataset_path data/multiclass \
  --model_name openai/whisper-small \
  --cld_type vanilla \
  --languages en,hi,id,ms,zh \
  --batch_size 32 \
  --no_wandb

Citation

If you use this code in your work, please cite the paper:

@inproceedings{feng2026cld,
  title     = {Convex Low-resource Accent-Robust Language Detection in Speech Recognition},
  author    = {Feng, Miria and Tan, William and Pilanci, Mert},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
  year      = {2026},
  series    = {Proceedings of Machine Learning Research},
  publisher = {PMLR},
  url       = {https://icml.cc/virtual/2026/poster/64615}
}