Hibiki-Zero 3B — Vietnamese→English full-model SFT
Full-model finetune of
kyutai/hibiki-zero-3b-pytorch-bf16 that adds
Vietnamese as a source language for simultaneous speech-to-text translation (VI speech → EN text + EN speech). The base model supports FR/ES/PT/DE → EN; Vietnamese is unseen.
Headline result: val128 chrF 9.2 → 19.61 vs the best LoRA attempt on the same data family, with the LoRA failure mode (degenerate looping, empty outputs) essentially eliminated: 0/128 empty, 126/128 clean EOS, 8/128 with any 4-gram repetition. Full-model SFT routes the new source language where LoRA (r32, incl. text/audio heads) structurally could not.
Architecture
Unchanged from upstream Hibiki-Zero (diagram: docs/hibiki-pipeline.svg):
- Mimi codec (streaming, 24 kHz in/out, 12.5 frames/s): a neural audio codec that turns speech into discrete tokens. Each frame is
n_q=16 codebook slots: 8 for the incoming source stream, dep_q=8 generated for the target speech.
- Temporal main transformer (3B, GQA
kv_repeat=2, hidden_scale=6): runs once per 80 ms frame over the joint sequence of [text token, 8 target codebooks, 8 source codebooks]; produces the frame's context vector and the next text token (48k SentencePiece vocab).
- Depth transformer ("depformer"): a small transformer that, conditioned on the main transformer's output, decodes the frame's 8 target audio codebooks sequentially (one pass per codebook, with a learned per-slice output LayerNorm). This serial loop is the mobile latency bottleneck — see the mobile plan below.
- Simultaneity: target text/audio are emitted while source audio still streams in (~6 s translation lag; an 8 s silence tail flushes the remaining buffer at end of input). Training bakes in a randomized coarse target delay (0–0.5× source duration) for alignment robustness.
This finetune trains every LM parameter (3.13B) — no LoRA, no frozen backbone. Mimi and the text tokenizer are untouched, so all existing inference stacks (PyTorch/moshi, MLX, moshi-swift) load it like the base model.
Training
| |
|---|
| Data | 147,421 cached pairs: anquachdev/PhoMT-en-vi-speech (222 VI-hours, Kokoro/VieNeu TTS) + FLEURS vi→en train (1,449 real-speech pairs) |
| Hardware | 1× H100 80 GB (vast.ai), torch 2.12.1+cu126 |
| Precision | fp32 master weights + bf16 autocast forward, fused AdamW, TF32 |
| Schedule | 3 epochs = 55,284 steps @ batch 8, max 280 frames (22.4 s); LR 1e-4 → 3e-5 at 50%, 500-step warmup; text-loss weight 5 → 2 at 60% |
| Selection | greedy decode on a fixed 128-row FLEURS validation gate (val128) every 9000 steps, best-on-chrF checkpoint |
| Wall clock | ~4.5 h ≈ $9 |
chrF trajectory on val128: 9.1 → 7.7 → 7.8 → 16.4 → 7.7 → 15.6 → 18.6 (final step = best). The mid-run dip is the grounding transition (outputs move from fluent-but-generic English to source-grounded content); consolidation begins with the LR decay. Val CE was still falling at termination — the model is undertrained at 224 h.
Final protocol eval (finetune/eval_lora.py --stop-on-eos, text temp 0.4, 128 rows): chrF 19.61, BLEU 0.22, WER 1.67. Known residual defect: over-generation (predictions average 1.9× reference length).
Files
| file | contents |
|---|
model_step055284.safetensors | full LM weights, bf16-storable fp32 train state, best & final (12.5 GB) |
trainer_step055284.pt | AdamW optimizer state + step counter for exact training resume (25 GB) |
mimi_caches.tar.gz | pre-encoded Mimi codes + text tokens for the full 224 h corpus + FLEURS train/val, plus all pairs jsonl (726 MB) — training never needs the raw audio again |
logs/ | train_log.jsonl, val_log.jsonl (teacher-forced CE every 2000 steps), greedy_eval_log.jsonl (all chrF gates), run_config.json, best.json |
eval_val128/ | final eval: per-row hypothesis texts + predictions.csv + metrics.json |
docs/ | training-repo docs: architecture diagram, finetune stack, VI training plan/diagnosis, training explainer |
Usage
Training/eval code:
huybik/hibiki-zero-mlx (
finetune/).
1# evaluate
2python finetune/eval_lora.py --device cuda --dtype bfloat16 \
3 --adapter model_step055284.safetensors \
4 --pairs finetune/pairs/val128.jsonl --limit 128 --batch-size 8
5
6# resume / continue training (phase 2)
7python finetune/train_lora.py --full-finetune --device cuda --dtype float32 \
8 --resume-checkpoint trainer_step055284.pt ...
Next (phase 2)
Scale toward 1,000 VI-hours: pool these caches with the new tranches, warm-start from this checkpoint, fresh warmup + longer low-LR tail (late-run chrF was volatile). Data priority: voice diversity / real Vietnamese speech (the mid-run train/val divergence came from single-voice TTS overfit), and reining in over-generation.
Mobile release plan (parallel codebook head)
Target: on-device simultaneous VI→EN inside the iPhone 80 ms/frame budget. The training repo's MLX stack and Track B scaffold already exist and are measured on M4 Pro; numbers below are from reports/ there.
Why a parallel head: the AR depformer's 8 sequential slice passes are launch-bound, not compute-bound — quantization buys ~0 on the depformer. Replacing it with a parallel codebook head (227 M params: shared 6-layer bidirectional trunk over the 16 codebook positions, per-codebook in/out projections + embeddings, each codebook conditioned on its own prev-frame token, warm-started from the AR depformer) is the only lever that removes the serial passes. Self-distilled from the frozen finetuned model: dump per-frame teacher logits (distill/dump_teacher.py), train head-only with KL(teacher‖student) + CE (distill/train_head.py).
Measured on the 3B (M4 Pro): codebook decode 12.7 → 5.3 ms (−58%), total LM step 25.9 → 18.5 ms (−29%), 3.1× → 4.3× realtime — with text quality untouched (text bypasses the head).
Latency budget (projected iPhone ≈ 0.5× M4-Pro GPU):
| config | M4 LM step | iPhone proj. | fits 80 ms |
|---|
| 3B q4, AR depformer | 22.2 ms | ~44 ms | ✅ |
| 3B q4 + parallel head | 18.5 ms | ~37 ms | ✅ |
| 1B q4, AR | 15.1 ms | ~30 ms | ✅ |
| 1B q4 + parallel head | ~13.4 ms | ~27 ms | ✅ |
Memory: 1B ≈ 1.8 GB total (fits a base 6 GB iPhone); 3B ≈ 3.4 GB (needs an 8 GB Pro + extended-memory entitlement).
Release steps for this VI model:
- Phase-2 data scale-up first — mobile ships whatever quality the 3B reaches; distill/quant preserve, never add, quality.
- Convert to MLX q4 (
scripts/convert_mlx_q4.py, group_size=32 — moshi-swift hardcodes gs32 for .q4.safetensors); verify with the silence-in + coherence gates.
- Distill the parallel head from the finetuned 3B on VI data (scaled run — the existing head is a smoke-scale scaffold), q4 the head (it is bandwidth-bound at 227 M bf16; q4 is the path from 5.3 ms toward the 1–4 ms target).
- Optionally distill 3B → Hibiki-M-style 1B for the base-iPhone tier (the FR→EN 1B artifact + staging/compat tooling
check_swift_compat.py already exist, 7/7 PASS pattern).
- Stage the phone artifact (config + q4 LM + mimi + tokenizer + model card), validate with
check_swift_compat.py, publish.
Honest open risks (device-side): no on-device moshi-swift build has been run yet; ANE-vs-GPU placement (current path is Metal GPU; the AR depformer maps poorly to ANE — one more reason for the parallel head); sustained thermals; launch-bound dispatch may scale worse than bandwidth on a phone GPU.
License follows upstream: CC BY-NC-SA 4.0.
