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).

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

Runtime	Backend	Graph on GPU	Latency
TFLite benchmark_model (TfLiteGpuDelegateV2) — w2v2_frontend_fp16.tflite	GPU (OpenCL)	134 / 134	38.4 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) — w2v2_head_fp16.tflite	GPU (OpenCL)	52 / 893	515.7 ms
TFLite benchmark_model — w2v2_frontend_fp16.tflite	CPU (XNNPACK, 4 threads)	—	XNNPACK declined the graph
TFLite benchmark_model — w2v2_head_fp16.tflite	CPU (XNNPACK, 4 threads)	—	123.3 ms

Any on-device figure recorded when this model shipped came from a different runtime. It was taken through LiteRT's own CompiledModel accelerator (logcat reports it as LITERT_CL), which is the path the Kotlin sample app and the LiteRT API use, and it appears elsewhere on this card. The rows above are the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. The two are not comparable, so read the rows above as a reproducible floor rather than as this model's speed on LiteRT.

XNNPACK declines these fp16 graphs — it reports failed to delegate DEPTHWISE_CONV_2D and then fails to allocate tensors — so there is no usable CPU number. Disabling XNNPACK falls back to reference kernels, which measured about 20× slower than the GPU on models of this size and would not represent CPU inference anyone would ship.

On this delegate the CPU is the faster choice for w2v2_head_fp16.tflite (123.3 ms on CPU against 515.7 ms on GPU) — worth knowing before you reach for the GPU on a mid-range phone.

Note that the GPU does not take the whole graph here (52 / 893 in w2v2_head_fp16.tflite); the remainder runs on the CPU and the split costs a per-partition round trip.