huybik/hibiki-zero-vi-full-sft

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

# evaluate
python finetune/eval_lora.py --device cuda --dtype bfloat16 \
  --adapter model_step055284.safetensors \
  --pairs finetune/pairs/val128.jsonl --limit 128 --batch-size 8

# resume / continue training (phase 2)
python finetune/train_lora.py --full-finetune --device cuda --dtype float32 \
  --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:

1. Phase-2 data scale-up first — mobile ships whatever quality the 3B reaches; distill/quant preserve, never add, quality.
2. 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.
3. 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).
4. 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).
5. 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.