| --- |
| license: apache-2.0 |
| language: |
| - en |
| tags: |
| - diffusion |
| - speculative-decoding |
| - rectified-flow |
| - dit |
| - qwen |
| - math-reasoning |
| - deltanet |
| datasets: |
| - AI-MO/NuminaMath-CoT |
| base_model: |
| - Qwen/Qwen3.5-9B |
| --- |
| |
| # Continuous Latent Speculative Decoding (CLSD) |
|
|
| **Architecture**: ~4.0B Hybrid Causal DiT (Rectified Flow) + 9B Frozen Verifier |
| **Key Innovation**: First hybrid DeltaNet/Attention causal diffusion transformer for parallel token generation |
| **Status**: Stage A converged, Stage C alignment in progress |
|
|
| --- |
|
|
| ## How it was envisaged |
|
|
| A machine for constructing candidate inner-machines that bias the verifier toward solution paths its default rollout would miss; learning a search policy that explores verifier-legible continuations and unlocks competence that ordinary autoregressive rollout fails to reach. |
|
|
| ## What is cool |
|
|
| The base verifier has fixed weights, but its inference process is not exhausted by ordinary left-to-right decoding. A learned continuous proposer can search for hidden-state trajectories and token paths that the verifier can recognize as correct, even if the verifier would rarely or never reach them under standard autoregressive rollout. |
| CLSD is a supervised-trained latent block proposer whose diffusion structure makes parallel search cheap enough to expose verifier-accessible solutions that AR decoding misses. |
| ## Thesis |
|
|
| Autoregressive language models are bottlenecked by sequential generation. CLSD deploys a hybrid causal Diffusion Transformer (DiT) -- a strided 12-layer slice of Qwen3.5-9B -- operating in the continuous embedding space of the same frozen Qwen3.5-9B verifier. Both models share the exact same 4096-dimensional manifold, the same tokenizer, and the same attention geometry. No projection bridges, no dimensional translation loss. |
|
|
| Qwen3.5-9B uses a hybrid architecture: 24 Gated DeltaNet (linear attention) layers + 8 standard quadratic attention layers in a repeating [3xDeltaNet, 1xAttention] pattern. The DiT preserves this hybrid structure and keeps **causal masking** -- DeltaNet linear recurrence is strictly causal by design. |
|
|
| The DiT drafts 32 candidate 128-token embedding sequences simultaneously in 2 Euler steps. The verifier evaluates them in a single batched forward pass. |
|
|
| > **Why causal diffusion works**: The conditioning vector C is injected via adaLN into every position simultaneously, providing global context regardless of attention mask. The causal constraint forces the DiT to learn autoregressive-like internal logic, which mirrors the frozen verifier expectations. |
|
|
| --- |
|
|
| ## Architecture |
|
|
| | Role | Model | Params | Dim | Layers | |
| |------|-------|--------|-----|--------| |
| | **Generator (DiT)** | Qwen3.5-9B strided slice | ~4.0B | 4096 | 12 (9 DeltaNet + 3 FullAttn) | |
| | **Verifier (frozen)** | Qwen3.5-9B (text tower) | 9B | 4096 | 32 | |
|
|
| ### The Strided Graft |
|
|
| ``` |
| Source layers: [0, 3, 6, 9, 12, 15, 18, 21, 24, 26, 28, 31] |
| Layer types: [D, A, D, D, D, A, D, D, D, D, D, A ] |
| |
| D = DeltaNet (linear_attention), A = full_attention |
| ``` |
|
|
| ### DiT Modifications |
| 1. **adaLN-Zero modulators** per block: nn.Linear(4096, 24576), zero-initialized |
| 2. **Timestep conditioning**: sinusoidal embedding + conditioning vector C |
| 3. **Learned local positional embedding**: nn.Parameter(zeros(1, 128, 4096)) |
| 4. Causal masking preserved from original Qwen weights |
|
|
| --- |
|
|
| ## Training Pipeline |
|
|
| ### Pre-Flight: Embedding Extraction |
|
|
| Target embeddings from **AI-MO/NuminaMath-CoT** (mathematical chain-of-thought): |
| - Tokenized with Qwen tokenizer, embeddings looked up via frozen embedding matrix |
| - Chunked into 128-token windows: [64, 128, 4096] safetensors shards |
| - **146,790 total chunks** across 2,294 files |
|
|
| ### Stage A: Rectified Flow (Velocity Regression) -- COMPLETE |
|
|
| The DiT learns the straight-line velocity field v = x1 - x0: |
|
|
| ``` |
| x_t = (1-t)*noise + t*target, t in [0,1] |
| L = ||v_pred - (target - noise)||^2 |
| ``` |
|
|
| | Parameter | Value | |
| |-----------|-------| |
| | Hardware | 1x NVIDIA B200 (183 GB) | |
| | Steps | 50,000 | |
| | Batch size | 32 | |
| | Optimizer | AdamW (lr=1e-4, cosine decay) | |
| | Wall-clock | 154.8 minutes | |
| | Final MSE | ~0.013 (converged by step 5K) | |
|
|
| ### Stage C: CE Alignment -- IN PROGRESS |
|
|
| Backpropagate through the frozen 9B verifier to teach the DiT semantic correctness: |
|
|
| ``` |
| noise -> DiT (2 Euler steps) -> draft_embeds |
| -> frozen Qwen 32 layers -> logits -> CE loss vs ground truth tokens |
| ``` |
|
|
| L_total = CE(logits, targets) + beta * MSE(drafts, true_embeddings) |
|
|
| Beta anneals from 0.1 to 0, gradually shifting from geometric to semantic alignment. |
|
|
| **Smoke test results** (50 steps, batch=1): |
| - CE dropped 12.8 -> 6.1: verifier starting to read DiT output |
| - Gradients flow correctly through frozen verifier |
|
|
| **Current run**: 2000 steps, batch=8, grad_accum=4 on B200 -- streaming to wandb |
| |
| --- |
| |
| ## Step 4: Live Inference (The Parallel Rollout) |
| |
| 1. User submits reasoning prompt |
| 2. 9B Verifier forward pass -> conditioning vector C + KV cache |
| 3. DiT generates **32 candidate 128-token branches** in 2 Euler steps |
| 4. 9B Verifier evaluates all 32 branches in one batched pass (shared prompt KV via PagedAttention) |
| 5. Score by mean log-probability across 128 positions |
| 6. **Causal Guillotine**: scan Top-1 left-to-right, truncate at first low-confidence position |
| 7. Qwen samples correct token, new C generated, loop repeats |
| |
| **Target latency**: <500ms per 128-token block |
| |
| --- |
| |
| ## Step 5: The Shadow Loop (Async RL -- Continuous Improvement) |
| |
| The Primary Node never stops drafting. A Shadow Node continuously improves the DiT: |
| |
| ``` |
| Primary Node --[Redis: 32 trajectories/cycle]--> Shadow Node |
| Shadow Node --[Weight sync every 1000 steps]--> Primary Node |
| ``` |
| |
| ### Objective Verification (Reward Signal) |
| |
| Feed Top-1 decoded tokens through: |
| - **Lean 4**: formal mathematical proof verification |
| - **Python sandbox**: code execution for correctness |
| |
| If verified -> reward the continuous vectors (positive signal) |
| If failed -> penalize (negative signal) |
| |
| This breaks the log-prob echo chamber. The DiT learns "alien intuition" -- solutions the 9B verifier would score as correct but would never stumble upon autoregressively. |
| |
| ### RL Objective |
| |
| Policy gradient from objective verification creates a reward signal independent of the verifier log-probs. The DiT explores the embedding space for novel solutions that: |
| 1. The verifier accepts (high log-prob) |
| 2. Actually solve the problem (Lean4/sandbox verification) |
| |
| This is an **infinite background process** -- the system improves continuously as long as compute is available. |
| |
| --- |
| |
| ## Repository Contents |
| |
| ``` |
| checkpoints/ |
| dit_stage_a_step_5000.pt # Early training |
| dit_stage_a_step_10000.pt # Mid training |
| dit_stage_a_step_30000.pt # Late training |
| dit_stage_a_final.pt # 50K steps, converged (MSE=0.013) |
| dit_stage_c_*.pt # CE alignment checkpoints (when available) |
| embeddings_sample/ # 50 representative embedding shards |
| batch_*.safetensors # Each: [64, 128, 4096] |
| ``` |
| |
| ### Loading a Checkpoint |
| |
| ```python |
| from clsd.grafted_dit import graft_dit_from_qwen, STRIDE_INDICES |
| from transformers import AutoModelForCausalLM |
| import torch |
|
|
| qwen = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-9B", dtype=torch.bfloat16) |
| dit, embed_tokens = graft_dit_from_qwen(qwen, slice_indices=STRIDE_INDICES) |
| state_dict = torch.load("checkpoints/dit_stage_a_final.pt", weights_only=True) |
| dit.load_state_dict(state_dict) |
| ``` |
| |
| --- |
| |
| ## Roadmap |
| |
| - [x] Pre-flight: embedding extraction (146K chunks from NuminaMath-CoT) |
| - [x] Step 1: Frankenstein graft (4.0B hybrid DiT from 9B) |
| - [x] Step 2: Stage A rectified flow (50K steps, converged) |
| - [x] Stage C smoke test (50 steps, pipeline validated) |
| - [ ] Step 3: Stage C full alignment (2000+ steps on B200) |
| - [ ] Step 4: Live inference with Causal Guillotine |
| - [ ] Step 5: Shadow Loop async RL with Lean4/sandbox verification |
| - [ ] Scale to 8x H200 cluster for production training |
| |
| ## Wandb |
| |
| - Stage A: [clsd-speedrun](https://wandb.ai/dalletest123/clsd-speedrun) |
| - Stage C smoke: [clsd-speedrun-smoke](https://wandb.ai/dalletest123/clsd-speedrun-smoke) |
| |
| ## License |
| |
| Apache 2.0 |
| |
