litert-community/Zipformer-medium-CR-CTC-LiteRT

Zipformer medium CR-CTC — LiteRT (GPU)

English speech recognition with the Zipformer encoder — the k2/icefall architecture — running fully on the LiteRT CompiledModel GPU (ML Drift). This is the CR-CTC medium checkpoint from the official icefall LibriSpeech recipe (64 M params, WER 2.12 test-clean / 4.62 test-other, greedy CTC), converted so that every op in the graph is GPU-compatible: one graph, no CPU fallback, no FFT inside the model.

Real on-device output: greedy-CTC word onsets for J.F. Kennedy's 1961 inaugural address (U.S. National Archives recording, public domain). A 16 s window transcribes in 156 ms on a Pixel 8a GPU (~19 ms enqueue), RTF ≈ 0.01.

All three CR-CTC variants from the recipe are included — identical I/O signatures (fbank [1,1600,80] + 4 mask biases → CTC logits [1,398,500]), so they are drop-in interchangeable in the same app:

File	Params	Size	WER (clean/other, greedy)	Pixel 8a 16 s window	API
zipformer_ctc_small_fp16.tflite	23 M	46 MB	2.57 / 5.95	124 ms	CompiledModel GPU
zipformer_ctc_fp16.tflite (medium)	64 M	132 MB	2.12 / 4.62	156 ms	CompiledModel GPU
zipformer_ctc_large_fp16.tflite	148 M	298 MB	2.03 / 4.37	220 ms	CompiledModel GPU

Pipeline

`16 kHz mono PCM → host kaldi-fbank (80 mel) → [GPU] Conv2dSubsampling + Zipformer2 (6 stacks)

• CTC linear → host greedy-CTC + BPE detokenize`

• Fixed 16 s window: fbank [1, 1600, 80] — torchaudio.compliance.kaldi.fbank with dither=0, snip_edges=False, high_freq=-400, waveform in [-1, 1] float (do not scale to int16 range). Shorter audio is padded with log(1e-10) frames.
• Mask inputs: padding is folded into the graph as additive attention biases (0 = real frame, -1000 = padding), one per internal frame rate: [1,796], [1,398], [1,199], [1,100]. Build the 50 Hz bias with valid = (fbank_frames - 7) // 2, then take [::2], [::4], [::8] slices.
• Output: raw CTC logits at 25 Hz (log_softmax was moved host-side; greedy argmax is unaffected). Blank id = 0, BPE vocab 500 (tokens.txt, bpe.model).

Minimal usage — Python

import numpy as np, torch, torchaudio
from ai_edge_litert.interpreter import Interpreter

wave, sr = torchaudio.load("speech.wav")            # 16 kHz mono, [-1,1]
feats = torchaudio.compliance.kaldi.fbank(
    wave, num_mel_bins=80, sample_frequency=16000,
    dither=0.0, snip_edges=False, high_freq=-400.0)
T = feats.shape[0]
x = torch.full((1600, 80), np.log(1e-10)); x[:min(T, 1600)] = feats[:1600]

valid = (min(T, 1600) - 7) // 2
b = np.full((1, 796), -1000.0, np.float32); b[0, :valid] = 0.0
biases = {796: b, 398: b[:, ::2], 199: b[:, ::4], 100: b[:, ::8]}

it = Interpreter(model_path="zipformer_ctc_fp16.tflite"); it.allocate_tensors()
for d in it.get_input_details():
    s = list(d["shape"])
    it.set_tensor(d["index"], x[None].numpy().astype(np.float32)
                  if len(s) == 3 else np.ascontiguousarray(biases[s[1]]))
it.invoke()
logits = it.get_tensor(it.get_output_details()[0]["index"])[0]  # [398, 500]

tokens = {int(l.rsplit(maxsplit=1)[1]): l.rsplit(maxsplit=1)[0]
          for l in open("tokens.txt", encoding="utf-8")}
out, prev = [], -1
for i in logits[: (valid + 1) // 2].argmax(-1):
    if i != prev and i != 0: out.append(tokens[int(i)])
    prev = i
print("".join(out).replace("▁", " ").strip())

Minimal usage — Kotlin (Android)

val model = CompiledModel.create(modelPath, CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()

// resolve slots by capacity: fbank 1600*80, biases 796/398/199/100 floats
val fbankSlot = inputs.indexOfFirst { it.readFloat().size == 1600 * 80 }
inputs[fbankSlot].writeFloat(fbank)                 // host kaldi-fbank, log(1e-10) padded
for (len in intArrayOf(796, 398, 199, 100)) {       // additive masks: 0 real / -1000 pad
    val slot = inputs.indexOfFirst { it.readFloat().size == len }
    inputs[slot].writeFloat(FloatArray(len) { i -> if (i * 796 / len < valid50) 0f else -1000f })
}

model.run(inputs, outputs)
val logits = outputs[0].readFloat()                  // [398 * 500], readback syncs the GPU
// greedy CTC: per-frame argmax over 500, drop blanks (id 0) and repeats, then BPE detok

On-device performance (Pixel 8a, CompiledModel GPU)

• GPU compile: 1.8 s (first load)
• 16 s window: 156 ms run+readback (19 ms enqueue) → RTF ≈ 0.01
• Device logits vs desktop float reference: corr 0.9993 (valid region), per-frame argmax agreement 99.2 %; transcripts identical on the test sweep.

Conversion notes

Converted from the icefall PyTorch checkpoint with litert-torch. All rewrites are numerically exact re-authorings of the eval path (tflite vs PyTorch: corr 1.000000):

• Swoosh-L/R via a guard-free stable softplus relu(z) + log1p(exp(-|z|)) (the default logaddexp lowering emits GPU-incompatible inf-guard selects).
• Relative-position shift (as_strided) re-authored as pad + reshape + slice.
• Padding masks folded into additive attention biases / multiplicative conv gates (icefall's own -1000 masked-fill semantics), supplied per frame rate as inputs.
• SimpleUpsample/SimpleDownsample expand → concat repetition; downsample weight softmax baked to a constant; final LogSoftmax moved host-side.

Sources & license

• Architecture / recipe: k2-fsa/icefall LibriSpeech zipformer recipe with CR-CTC (icefall PR #1766) — Apache-2.0.
• Checkpoint: Zengwei/icefall-asr-librispeech-zipformer-medium-cr-ctc-20241018.
• Paper: Zipformer: A faster and better encoder for automatic speech recognition (ICLR 2024).
• Hero audio: J.F. Kennedy inaugural address (1961), U.S. National Archives — public domain.