WakeHuBERT: MixConv mel-TCN

A 0.64M-parameter speech feature extractor for wake-word detection, distilled from HuBERT-base: eight dilated blocks whose depthwise convolutions mix kernel sizes 3, 5, 7 and 9, with 256 channels, on a fixed log-mel front end (64 bins) and a strided convolution to 50 frames per second. It turns 16 kHz audio into 128-dimensional features, so a small classifier trained on those features from synthetic speech alone can detect a wake word. It is one of the extractors that wakeforge can load; TigreGotico/wakehubert-tiny is the 0.64M HuBERT-distilled mel-TCN.

Using it

The ONNX graph takes waveform [batch, samples] (float32, 16 kHz, -1..1) and returns features [batch, samples // 320, 128]. It is strictly causal: frame t depends only on audio before sample 320·(t+1); streamed in chunks it matches offline features once 5 s of context are kept. It was trained to reproduce each teacher frame five frames (100 ms) later, so its features describe the speech with a delay of about 100 ms. wakehubert_int8.onnx is a static int8 version whose features agree with float at a mean cosine of 0.991.

import numpy as np, onnxruntime as ort
from huggingface_hub import hf_hub_download
sess = ort.InferenceSession(hf_hub_download("TigreGotico/wakehubert-mixconv", "wakehubert.onnx"))
feats = sess.run(None, {"waveform": np.zeros((1, 24000), np.float32)})[0]  # (1, 75, 128)

How it was made

It was trained for 30,000 steps to predict standardised HuBERT-base layers 4, 8 and 12 (L1 plus log-sigmoid cosine, as in DistilHuBERT), with the teacher hearing clean speech while the student heard it with noise, room reverberation and one to three background talkers never louder than the voice; a quarter of the training items were non-speech sounds heard identically by both.

Training speech: LibriSpeech (960 h), Multilingual LibriSpeech (seven languages) and a language-balanced sample of Multilingual Spoken Words (41 languages), cut into 600,000 two-second crops.

Results

Single run. Classifiers trained only on 900 synthetic "alexa" clips (TTS voices cloned with voice conversion, with noise, babble, reverberation, speed and gain augmentation) were scored on the real speakers of the Picovoice wake-word benchmark (315 recordings), with the threshold chosen for 0.5 false activations per hour on the 6.5 h of held-out streams it is reported on (LibriSpeech test-clean, three-talker test-other babble, held-out non-speech), so false activations hold by construction:

Differences of a few points between extractors are within noise, and the two classifier types on the same extractor can differ by several points.

PyTorch weights

model.safetensors holds the float32 weights of the student (643,840 trainable parameters, plus the batch-norm statistics and the fixed log-mel filters) and student.py defines the network in plain PyTorch with no dependency beyond torch and safetensors; the pytorch.init_args block of config.json selects this variant (mixconv). The weights are the checkpoint that wakehubert.onnx was exported from, so they can be fine-tuned or exported to other formats. The three linear heads that predicted the teacher's HuBERT-base layers 4, 8 and 12 during distillation are not included.

from huggingface_hub import snapshot_download
import sys, torch; d = snapshot_download("TigreGotico/wakehubert-mixconv"); sys.path.insert(0, d)
from student import load_student
model = load_student(f"{d}/model.safetensors")
feats = model(torch.zeros(1, 24000))  # (1, 75, 128)
torch.onnx.export(model, torch.zeros(1, 16000), "wakehubert.onnx", input_names=["waveform"], output_names=["features"], dynamic_axes={"waveform": {0: "batch", 1: "samples"}, "features": {0: "batch", 1: "frames"}}, opset_version=17, dynamo=False)

The PyTorch model and wakehubert.onnx agree to within 1e-4 on the 128-dimensional features (largest absolute difference 1.3e-05 over five random three-second waveforms and one ten-second speech recording), and exporting the PyTorch model again gives an ONNX graph whose outputs match the published one to within 7e-06. The distillation and training code is in wakeforge.

Full training checkpoints

training/best.pt (step 22,000, the checkpoint wakehubert.onnx and model.safetensors come from) and training/last.pt (step 30,000) are the complete PyTorch checkpoints of the distillation run, for continuing the distillation rather than only fine-tuning the student. Each is a dict with student (the student state_dict), norm (the standardisation statistics of the teacher targets), opt and sched (optimizer and learning-rate scheduler state), step, best and meta; the three linear heads that predicted the teacher layers are inside student under proj.*. They are pickle files: load them with torch.load(path, map_location="cpu", weights_only=False) only if you trust this repository. The distillation script that reads them is in wakeforge.

Licence and attribution

Released under the Apache License 2.0, the licence of its teacher (HuBERT-base, Apache 2.0). Trained with LibriSpeech, Multilingual LibriSpeech and Multilingual Spoken Words (all CC BY 4.0), MUSAN, AudioSet-derived noise and room impulse responses, and distilled from facebook/hubert-base-ls960.