litert-community/wav2vec2-keyword-spotting

wav2vec2 Keyword Spotting — LiteRT on-device (all-GPU)

On-device LiteRT conversion of superb/wav2vec2-base-superb-ks (Apache-2.0) — wav2vec2-base keyword spotting, 12 Speech-Commands labels (yes / no / up / down / left / right / on / off / stop / go / unknown / silence). Runs fully on the CompiledModel GPU delegate (LITERT_CL); there is no FFT anywhere (the raw 16 kHz waveform goes straight into a 1D-conv feature extractor — no mel step), and the transformer residual is small enough that the whole model is fp16-exact on GPU (no CPU fallback). Device-verified on a Pixel 8a (Tensor G3).

Files

File	Size	Delegate	In → Out
w2v2_frontend_fp16.tflite	9 MB	GPU	audio [1,16000] → feat [1,49,768]
w2v2_head_fp16.tflite	181 MB	GPU	feat [1,49,768] → logits [1,12]

End-to-end ~19 ms for a 1 s clip (RTF ≈ 0.02); 10/10 keywords correct on real speech, device-vs-CPU logits corr 0.9995.

Two graphs

The model is op-clean for the GPU but the full 1008-node graph exceeds the Mali shader-compile limit (fails to compile fused). Splitting at the conv-frontend / transformer-encoder boundary makes each half compile (frontend 134/134 + head 893/893 LITERT_CL). The frontend output feeds the head:

waveform →[GPU frontend]→ feat[1,49,768] →[GPU head]→ logits[1,12] → argmax

• frontend = feature_extractor (7 strided 1D convs + GroupNorm) + feature_projection.
• head = encoder (12 transformer layers + conv positional embedding) + weighted-layer-sum over all 13 hidden states (use_weighted_layer_sum) + projector + mean-pool + classifier.

Minimal usage

Android (Kotlin, CompiledModel GPU)

// staged into filesDir by the sample's install_to_device.sh (the head is 181 MB)
val opts = CompiledModel.Options(Accelerator.GPU)
val fe = CompiledModel.create(File(ctx.filesDir, "w2v2_frontend_fp16.tflite").absolutePath, opts, null)
val head = CompiledModel.create(File(ctx.filesDir, "w2v2_head_fp16.tflite").absolutePath, opts, null)
val feIn = fe.createInputBuffers(); val feOut = fe.createOutputBuffers()
val hdIn = head.createInputBuffers(); val hdOut = head.createOutputBuffers()
feIn[0].writeFloat(audio)                 // [1,16000] 1 s @ 16 kHz, raw [-1,1]
fe.run(feIn, feOut)
hdIn[0].writeFloat(feOut[0].readFloat())  // feat [1,49,768]
head.run(hdIn, hdOut)
val logits = hdOut[0].readFloat()         // [12] -> argmax = keyword

Python (desktop verification)

import numpy as np, soundfile as sf
from ai_edge_litert.interpreter import Interpreter

LABELS = ["yes", "no", "up", "down", "left", "right",
          "on", "off", "stop", "go", "_unknown_", "_silence_"]

wav, _ = sf.read("clip_16k.wav", dtype="float32")            # mono 16 kHz
x = np.zeros((1, 16000), np.float32); n = min(len(wav), 16000); x[0, :n] = wav[:n]
# raw [-1,1] waveform straight in — this checkpoint uses do_normalize=False

fe = Interpreter(model_path="w2v2_frontend_fp16.tflite"); fe.allocate_tensors()
fe.set_tensor(fe.get_input_details()[0]["index"], x); fe.invoke()
feat = fe.get_tensor(fe.get_output_details()[0]["index"])    # [1,49,768]

hd = Interpreter(model_path="w2v2_head_fp16.tflite"); hd.allocate_tensors()
hd.set_tensor(hd.get_input_details()[0]["index"], feat); hd.invoke()
logits = hd.get_tensor(hd.get_output_details()[0]["index"])[0]  # [12]
print(LABELS[int(logits.argmax())])

Re-authoring (litert-torch, parity corr 1.0)

GELU→tanh-GELU · feature-extractor GroupNorm→GN4D · pos-conv weight_norm fold · create_bidirectional_mask→None · weighted-layer-sum accumulated incrementally with baked softmax(layer_weights) constants (the runtime softmax + per-layer scalar gathers otherwise split the Mali partition).

Sample app

A complete Android sample app + the conversion scripts are in the official LiteRT samples repository under compiled_model_api/audio_classification (google-ai-edge/litert-samples). Push these files to the app's filesDir with that sample's install_to_device.sh.

License follows upstream (Apache-2.0).