---
title: Conveyor-Belt Diffusion LLM
emoji: 🥖
colorFrom: green
colorTo: purple
sdk: docker
app_port: 7860
license: mit
tags:
  - diffusion-llm
  - rolling-diffusion
  - language-modeling
  - tinystories
  - pytorch
---

# Conveyor-Belt Diffusion LLM (64-Slot Dual-Channel Elastic Planning Oven)

A **28.88M-parameter hybrid Autoregressive + Diffusion Language Model** where text generation works like a **64-slot bakery conveyor belt** (inspired by *Rolling Diffusion*, *FIFO-Diffusion*, *Diffusion Forcing*, and *AR-Diffusion*):

- **Input / Output Contract at every tick $t$:**
  $$\text{Input: } \underbrace{[x_0, \dots, x_{t-1}]}_{\text{Frozen Causal KV Prefix}} + \underbrace{[d_0^{(t)}, d_1^{(t)}, \dots, d_{63}^{(t)}]}_{\text{64-Slot Conveyor \& Planning Belt}} \quad\longrightarrow\quad \text{Output: } \underbrace{x_t}_{\text{Popped Slot-0 Token}} + \underbrace{[d_0^{(t+1)}, \dots, d_{63}^{(t+1)}]}_{\text{Shifted \& Refined Belt}}$$

- **Live Demo & Pretrained Weights:**
  - **Hugging Face Live Space:** [`https://huggingface.co/spaces/sumitk/diffusion-conveyer-llm`](https://huggingface.co/spaces/sumitk/diffusion-conveyer-llm)
  - **Pretrained 24h Checkpoint (`best.pt`):** [`https://huggingface.co/sumitk/diffusion-conveyer-llm`](https://huggingface.co/sumitk/diffusion-conveyer-llm)

---

## Key Architectural Innovations

1. **Dual-Channel "Soft-Vocab + Gated-Thought" Dough (`model.py`):**
   - **Channel 1 (Lexical Anchor):** Expected vocabulary embedding $e_{\text{soft}, k} = p_k^\top E_{\text{vocab}}$ under a positional temperature schedule $\tau_k \in [0.5, 4.0]$ (soft and exploratory at the back of the belt, crisp at the front).
   - **Channel 2 (Continuous Thought Highway):** Gated residual thought state $m_k \in \mathbb{R}^d$ initialized with forget-gate bias $+1.0$, fused via $\alpha_k = k / 63$ so **$\alpha_0 = 0$ exactly at Slot 0** (guaranteeing zero distribution mismatch when $x_t$ commits into the causal KV cache).
2. **Two-Zone Human-Style Thinking & Speaking Belt:**
   - **Front 25% Articulation Stream (`Slots 0..15`):** Shifts left by 1 slot every tick ($s_k = 1.0$), polishing immediate grammar toward `Slot 0`.
   - **Mid-Belt Clause Expansion (`Slot 15`):** Synthesizes new intermediate phrase embeddings (`expand_head`) at the boundary between the paused planning zone and moving front conveyor.
   - **Back 75% Elastic Planning Horizon (`Slots 16..63`):** Spans a $2\times$ wider lookahead horizon with a learned shift/pause gate $s_k \in (0, 1)$, holding sentence-level goals while clauses form in the middle.
3. **Self-Healing Belt Recovery Training (`losses.py`):**
   - Trained with 16-step unrolled BPTT, full-sequence causal LM supervision (`prefix_lm_weight=1.0`), front-belt shift-consistency KL, DAgger-style scheduled sampling, and mid-unroll belt corruption so the belt actively repairs noisy or contradictory plans.

---

## 24-Hour Training Results on 2,000,000 Real TinyStories

Trained for **24.00 hours (`122,743 steps`)** on a single RTX 5070 Ti GPU across **2,000,000 real TinyStories** (0% synthetic templates, 4,096 compact `tiktoken` BPE vocabulary with byte fallback for zero `<UNK>`):

| Metric | Value | Notes |
| :--- | :---: | :--- |
| **Held-Out Real Val Slot-0 Perplexity** | **`4.21`** | Down from `104.49` at step 500 (10,000 unseen real stories) |
| **Held-Out Real Val Slot-0 Top-1 Acc** | **`64.39%`** | Exact next-subword match at Conveyor Slot 0 |
| **Causal Prefix Trunk PPL / Top-1 Acc** | **`3.49` / `66.78%`** | Comparable to Microsoft Research's `TinyStories-28M/33M` (`CE 1.25`) |
| **Zero-Shot StoryCloze Commonsense** | **`100.0% (16/16)`** | 2-way physical, emotional & entity-tracking pairs (`+2.58 nats` margin) |
| **Word-Order Sensitivity KL** | **`3.674 nats`** | `KL(P(slot0 \| prompt) \|\| P(slot0 \| shuffled_prompt))` |
| **64-Slot Crystallization (Slot 63 $\to$ 0)** | **`8.19` $\to$ `0.48 nats`** | `17x` uncertainty reduction as thoughts travel from back to front |

---

## Quickstart

### 1. Install
```bash
pip install -e .
```

### 2. Launch the Interactive 64-Slot Conveyor Web Simulator
```bash
PYTHONPATH=src python -m conveyer.web_sim --port 8766
```
*(If `runs/tinystories/best.pt` is not present locally, it will automatically download the pretrained 28.88M checkpoint from `huggingface.co/sumitk/diffusion-conveyer-llm`.)*

### 3. CLI Tick-by-Tick Generation & Belt Inspection
```bash
PYTHONPATH=src python -m conveyer.infer \
  --ckpt runs/tinystories/best.pt \
  --prompt "Once upon a time, Lily and a little bunny walked into the green forest to find the golden key near the locked treasure chest. Suddenly a big lion" \
  --tokens 80 --overlap 10
```

### 4. Run the Full Benchmark Suite
```bash
PYTHONPATH=src python -m conveyer.eval --ckpt runs/tinystories/best.pt
```

### 5. Train / Resume 24-Hour Real TinyStories Run
```bash
PYTHONPATH=src python -u -m conveyer.train --config configs/tinystories.yaml --resume auto
```