WakeHuBERT: wide mel-TCN

A 2.01M-parameter speech feature extractor for wake-word detection, distilled from HuBERT-base: twelve dilated depthwise-separable convolution blocks with 384 channels and 256 output features, on a fixed log-mel front end (64 bins) and a strided convolution to 50 frames per second. It turns 16 kHz audio into 256-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, 256]. It is strictly causal: frame t depends only on audio before sample 320·(t+1); streamed in chunks it matches offline features once 4 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.984.

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

How it was made

It was trained for 50,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. Spans of the student's input log-mel were masked during training (probability 0.065 per frame, 10-frame spans).

Training speech: LibriSpeech train-clean-360 and train-other-500 and a language-balanced sample of Multilingual Spoken Words (50 languages), cut into 200,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 on separate calibration audio (LibriSpeech dev-clean and babble made from it) for 0.5 false activations per hour, and false activations then measured on 6.5 h of held-out streams (LibriSpeech test-clean, three-talker test-other babble, held-out non-speech):

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 (2,008,704 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 (tcn). 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-mel-tcn-wide"); 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, 256)
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 256-dimensional features (largest absolute difference 2.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 1e-05. The distillation and training code is in wakeforge.

Full training checkpoints

training/best.pt (step 47,500, the checkpoint wakehubert.onnx and model.safetensors come from) and training/last.pt (step 50,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 and Multilingual Spoken Words (both CC BY 4.0), AudioSet-derived noise, non-speech sound collections and room impulse responses, and distilled from facebook/hubert-base-ls960.