| --- |
| language: en |
| license: other |
| base_model: |
| - JetLM/SDAR-1.7B-Chat |
| - Qwen/Qwen3-ASR-0.6B |
| datasets: |
| - librispeech_asr |
| tags: |
| - automatic-speech-recognition |
| - speech |
| - diffusion |
| - lora |
| - block-diffusion |
| pipeline_tag: automatic-speech-recognition |
| --- |
| |
| # Diffusion ASR |
|
|
| Speech recognition with a **block-diffusion language model as the decoder**. A frozen audio |
| encoder and a frozen text diffusion LM are joined by a small trained adapter plus LoRA — about |
| 1 % of the stack is trained. |
|
|
| **4.60 % WER on LibriSpeech test-clean** (all 2,620 clips). |
|
|
| > **These are trained deltas, not a standalone model.** This repo holds the 5.25 M-parameter |
| > adapter and the 17.4 M-parameter LoRA. Running them also needs the two frozen base models |
| > below and the inference code from |
| > [ipritamdash/diffusion-asr-research](https://github.com/ipritamdash/diffusion-asr-research). |
|
|
| ## Why this exists |
|
|
| The audio encoder came from a working ASR system. That system's own decoder is thrown away here |
| and a text-only diffusion model put in its place, connected through roughly 5 M trained |
| parameters. |
|
|
| The question is whether a language model that has never seen audio can be conditioned on |
| acoustic evidence through a thin learned interface — and whether writing several tokens per |
| forward pass buys anything. |
|
|
| Compare it against other frozen-encoder + frozen-LM + adapter systems. Do not compare it against |
| end-to-end ASR trained on thousands of hours; it saw 30. |
|
|
| ## Results |
|
|
| Full LibriSpeech test-clean, 2,620 clips, corpus WER after Whisper English normalisation. |
|
|
| | Decoder | Config | WER | vs real time | |
| |---|---|---:|---:| |
| | **Reference loop** | block 2 | **4.60** | — | |
| | Reference loop | block 4 | 5.23 | — | |
| | **LMDeploy engine, swept config** | **block 4 · 4 steps · `sequential`** | **4.78** | 7.0× | |
| | LMDeploy engine, old default | block 2 · dynamic | 5.00 | 9.1× | |
| | LMDeploy engine, old default | block 4 · dynamic | 5.63 | 13.5× | |
|
|
| Speed is single-stream real-time multiple on an A100; batched serving measured far higher |
| (~800× on one L40S in an earlier benchmark). Under the engine's original confidence-ordered |
| commit rule it read a constant 0.40 worse than the reference loop; the swept `sequential` |
| commit order closes half of that (engine best 4.78 vs reference 4.60). The remaining +0.18 is |
| unexplained and stated rather than left out. |
|
|
| ### Error analysis |
|
|
| Computed over every clip of the 4.60 decode. |
|
|
| **Output health** — zero empty transcripts, zero repeat loops, zero length blow-ups, and no clip |
| hit the token budget. |
|
|
| **Composition** — of 53,029 reference words: 3.75 % substituted, 0.40 % deleted, 0.45 % |
| inserted. Four errors in five are a *wrong* word rather than a missing one; the model commits to |
| a plausible word from unclear evidence instead of failing to produce one. |
|
|
| **Substitutions** — 40 % are within a small character edit of the reference (*jailer* → |
| *jailor*, *holmes* → *homes*), which reads as blurred acoustic detail rather than a language |
| failure. |
|
|
| **Multi-piece words** — words the tokeniser splits into several subword pieces are 7.5 % of |
| words but carry 52 % of errors, at 13× the failure rate of single-piece words. This is the |
| dominant failure mode. |
|
|
| **Length and position** — the shortest quarter of clips is hardest (5.85 % against 4.37 % on the |
| longest), and quality is flat across an utterance. A fixed token budget and a block decoder could |
| have degraded on long input; neither did. |
|
|
| ## Files |
|
|
| | File | Contents | Params | |
| |---|---|---:| |
| | `adapter.pt` | audio adapter + `gain`, CTC logit scale, CTC blank projection | 5,250,049 (+2,050) | |
| | `lora.pt` | LoRA on the decoder, PEFT state dict, 392 tensors | 17,432,576 | |
|
|
| `adapter.pt` keys: `adapter`, `log_s`, `blank_w`, `blank_b`, `epoch`, `adapter_kind` (`conv`), |
| `adapter_stride` (`1`), `attn_mode` (`stair`), plus dev metrics. |
| `lora.pt` keys: `lora`, `epoch`, `dev_ce`, `objective`. |
|
|
| `log_s`, `blank_w` and `blank_b` belong to the CTC alignment loss used during training. They are |
| not used at decode time, and are kept so training can be resumed or audited. |
|
|
| ## Base models — frozen, download separately |
|
|
| | Role | Repo | Pinned revision | |
| |---|---|---| |
| | Decoder | `JetLM/SDAR-1.7B-Chat` | `97cf0dc9a7e85433f3f71a6bbac6fec9d381294f` | |
| | Encoder | `Qwen/Qwen3-ASR-0.6B` | `5eb144179a02acc5e5ba31e748d22b0cf3e303b0` | |
|
|
| Use these exact revisions — the weights here were trained against them. |
|
|
| > The decoder must be loaded through the project repo's patched loader |
| > (`qwen_sdar/sdar.py::load_patched_sdar`), which swaps CUDA-only flash-attn for SDPA and adds |
| > the embeddings-input path. Plain `from_pretrained` will not reproduce these numbers. |
|
|
| ## Architecture |
|
|
| ``` |
| audio 16 kHz → Qwen3-ASR-0.6B encoder → adapter → SDAR-1.7B decoder + LoRA → text |
| frozen, 1024-d @ 12.5 Hz trained frozen trained |
| ``` |
|
|
| **Adapter** — `Conv1d(1024→1024, k=3, stride=1)` → GELU → `Linear(1024→2048)` → |
| `LayerNorm(2048)` → ×`gain`. Stride 1 keeps the encoder's native 12.5 Hz; pooling to 6.25 Hz |
| measured worse across three seeds per arm, at roughly half the training cost. |
|
|
| `gain` initialises to the RMS of the decoder's token embeddings. A random projection emits |
| vectors 10–50× the size of real word vectors, and a decoder fed a prefix that loud learns to |
| ignore it; matching the embedding scale is what makes the audio legible at all. |
|
|
| **LoRA** — `r=16`, `alpha=32`, dropout 0, on `q,k,v,o,gate,up,down` in all 28 layers. LoRA scales |
| updates by `alpha/r`, so the effective learning rate on those weights is 2× nominal. |
|
|
| **Attention is block-causal at every stage**, training and inference alike. Each block sees the |
| audio, all earlier blocks, and itself. Training and serving therefore use the same convention, |
| and the model runs on a standard diffusion inference engine without retraining. |
|
|
| ## Training |
|
|
| 30 hours of LibriSpeech `train-clean-100`, speaker-disjoint from the held-out split. Three |
| stages, never joint: |
|
|
| | Stage | Trains | Loss | |
| |---|---|---| |
| | Grounding | adapter | CTC only, decoder detached | |
| | SFT | adapter | CTC + diffusion cross-entropy | |
| | LoRA | LoRA | diffusion cross-entropy only, adapter frozen | |
|
|
| Roughly 32 GPU-hours in total. |
|
|
| ## Decoding |
|
|
| Block-causal, blocks written left to right, earlier blocks frozen once committed. |
|
|
| | | | |
| |---|---| |
| | Block width | 4 (engine, swept) · 2 (reference loop) | |
| | Denoising steps per block | up to 4 | |
| | Commit rule | reference loop: confidence ≥ 0.95 with forced fallback; engine: **`sequential` unmasking** (position order), found best in a 50-arm sweep | |
| | Token budget | 160 | |
|
|
| A 50-arm sweep of every engine dial (widths 1–64, steps, thresholds, all three unmasking |
| strategies), confirmed on the full dev set, chose the engine config above; decoding wider than |
| the trained block-4 width collapses (width 8 already mass-produces empty output), so the |
| frontier honestly ends at width 4. |
|
|
| With the original confidence-ordered commits, width 2 (below the paper's swept range) measured |
| best; with `sequential` commits the trained width 4 wins — commit order, not width alone, was |
| carrying that result. |
|
|
| ## Usage |
|
|
| The encoder and decoder have conflicting dependencies, so inference runs in two steps in two |
| environments. Both need a CUDA GPU. |
|
|
| ```bash |
| git clone https://github.com/ipritamdash/diffusion-asr-research |
| cd diffusion-asr-research |
| |
| huggingface-cli download Tachyeon/diffusion-asr adapter.pt lora.pt --local-dir . |
| |
| # 1. encoder environment: audio -> features |
| pip install -r requirements-encoder.txt |
| python featurize.py --audio clip.wav --out features.pt |
| |
| # 2. decoder environment: features -> transcript |
| pip install -r requirements-decode.txt |
| python decode.py --features features.pt --adapter adapter.pt --lora lora.pt --block_length 2 |
| ``` |
|
|
| `decode.py` reads `adapter_stride` from the checkpoint and refuses to run if it is not recorded. |
| A stride-1 and a stride-2 adapter have identical tensor shapes, so guessing wrong loads cleanly |
| and silently emits a prefix of the wrong length. |
|
|
| For serving, merge the LoRA and run on LMDeploy's PyTorch diffusion engine; see `engine/` in the |
| project repo. |
|
|
| ## Scope |
|
|
| A research model for **English read speech**. LibriSpeech is clean, read audiobook material and |
| these numbers describe that domain. |
|
|
| ## Licence |
|
|
| The adapter and LoRA here are trained deltas. Any use depends on the two frozen base models, |
| which carry **their own licences** — check `JetLM/SDAR-1.7B-Chat` and `Qwen/Qwen3-ASR-0.6B` |
| before redistributing anything, and do not redistribute a merged checkpoint (which embeds the |
| base weights) without confirming their terms. |
|
|
