README.md

---
license: apache-2.0
base_model: Qwen/Qwen3-0.6B
pipeline_tag: text-generation
language:
  - en
tags:
  - diffusion
  - block-diffusion
  - masked-diffusion
  - chat
  - memory
  - qwen3
---

# Marimo Diffusion 0.6B

A 0.6B chat model that **thinks in denoised blocks and remembers through its own notes**
instead of re-reading the conversation. It is a retrofit of
[Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B): continued pretraining converts the
autoregressive base into a hybrid block-diffusion model, and a supervised fine-tune teaches it
a note-taking chat format.

Three things make it different from a standard chat model of this size:

- **Adaptive thinking blocks.** Before answering, the model may open thinking blocks of 32/64/128
  tokens, chosen per thought via control tokens (`<szN>`), each denoised bidirectionally in 16
  steps. Trivial turns skip thinking entirely — that decision is trained, not prompted.
- **Ledger memory.** Only the last 4 messages are kept verbatim in the prefix. Everything older
  survives only as the model's own notes (`key: value`), merged so the latest value wins. A
  100-turn conversation fits in a ~500-token prefix that never grows.
- **Constant cost per turn.** ~2 s/turn on an RTX 3090 regardless of conversation length,
  because the prefix is constant by construction.

## ⚠️ This model does NOT run under standard runtimes

The weights are Qwen3 architecture, but generation requires the block-denoising sampler included
in this repository. **transformers `generate()`, llama.cpp, GGUF, Ollama and LM Studio will not
produce correct output** — their autoregressive decoding never matches the training objective.
Use the bundled code.

## Quickstart

```bash
git clone https://huggingface.co/goldenfox/marimo-diffusion
cd marimo-diffusion
pip install torch tokenizers numpy

# OpenAI-compatible server (any OpenAI-API chat client can connect)
PYTHONPATH=src python -m diffusion_lm.chat_server \
  --checkpoint marimo-diffusion-0.6b.pt \
  --tokenizer tokenizer-qwen3-adaptive.json \
  --port 7998
```

Then point any OpenAI-compatible client (Chatbox, Open WebUI, curl) at
`http://127.0.0.1:7998/v1` with model id `marimo-diffusion-0.6b`:

```bash
curl http://127.0.0.1:7998/v1/chat/completions -H 'Content-Type: application/json' -d '{
  "model": "marimo-diffusion-0.6b",
  "messages": [{"role": "user", "content": "my sister lands friday 6pm, flight AR1420"}]
}'
```

The response carries the model's notes in `reasoning_content` (same field DeepSeek uses), so
clients that render reasoning show them automatically. The server caches each turn's notes and
rebuilds the ledger across stateless requests.

For the interactive playground (streaming denoise view, per-turn data log):

```bash
pip install gradio
PYTHONPATH=src python -m diffusion_lm.reasoning_playground \
  --outputs-dir . --prefix marimo \
  --tokenizer tokenizer-qwen3-adaptive.json --port 7999
```

A CUDA GPU is recommended (any 6 GB+ card fits the bf16 weights). GPU memory: ~2.5 GB.

## Benchmark: ledger needle (memory across 100 turns)

The long-context needle test, adapted to what this architecture claims: a scripted 100-turn
conversation plants 15 facts, corrects 5 of them, and probes recall at distances of 3–96 turns.
Baselines get every reasonable advantage: full history in context, greedy decoding, an explicit
memory instruction, and native thinking mode where it exists. Identical scoring for all systems.
Full per-turn data and the interactive viewers are in [`bench/`](./bench).

| system | params | recall | test total (s) | s/turn | max prefix (tok) |
|---|---|---|---|---|---|
| Qwen2.5-1.5B-Instruct · full history | 1.5B | **9/10** | 80 | 0.8 | 2,881 |
| Qwen3-0.6B + thinking · full history | 0.6B | 7/10 | 1,320 | 13.2 | 4,450 |
| **Marimo Diffusion (ledger)** | **0.6B** | **6/10** | **220** | **2.2** | **625** |
| Qwen2.5-0.5B-Instruct · full history | 0.5B | 4/10 | 110 | 1.1 | 3,908 |
| SmolLM2-360M-Instruct · full history | 0.36B | 4/10 | 60 | 0.6 | 2,712 |
| TinyLlama-1.1B-Chat · full history | 1.1B | 3/10 | 120 | 1.2 | 5,126 |
| Qwen3-0.6B + thinking · 512-token budget | 0.6B | 2/10 | 1,337 | 13.4 | 508 |
| Qwen3-0.6B no thinking · full history | 0.6B | 1/10 | 130 | 1.3 | 3,722 |
| Qwen3-0.6B no thinking · 512-token budget | 0.6B | 1/10 | 160 | 1.6 | 512 |

Reading this honestly:

- **In its size class it leads**: every ≤0.6B baseline with the full transcript in context
  scores 4/10 or less; the ledger reaches 6/10 from a 7× smaller prefix.
- **Beating it costs something**: 2.5× the parameters (Qwen2.5-1.5B), or the same base model's
  thinking mode at **6× the latency** with an unbounded prefix — and 7 vs 6 on ten probes is
  within noise.
- **At the same 512-token budget the mechanism decides**: thinking mode drops to 2/10 while the
  ledger holds 6/10. Same weights, same reasoning mode, same token budget — the only difference
  is a sliding window versus the model's own notes.
- **Failure profiles are complementary.** The thinking baseline re-reads verbatim, so it never
  suffers a corrupted note; the ledger never suffers long-context attention loss (it recalled
  facts at distance 59–60 that the thinking baseline missed with the text in front of it).
- **Caveats**: one seed, ten probes; and the comparison measures the mechanism *and* its
  training together — this model was trained on this conversational register, the baselines
  were not.

## Training

- Base: Qwen3-0.6B. Continued pretraining converts AR → hybrid block diffusion (answer region
  stays autoregressive; thinking blocks are masked-denoised bidirectionally).
- SFT: 174k examples from ~38k conversations — synthetic memory-task dialogues (recap,
  correction, distant-combination, each ending in a consolidating close), passage-grounded QA
  with the source dataset's reference answer as an exact quality gate, an abstention slice, and
  persona-grounded dialogues. Sequence length 512 (median example: 194 tokens); the checkpoint
  is served at 2,048 (RoPE, no learned positions).
- `steps_per_block 16` is the measured optimum for this checkpoint: best numeric fidelity at
  half the latency of 32; below 8 both prose and numbers degrade.

### Training data provenance

| source | role | license |
|---|---|---|
| synthetic dialogues (DeepSeek v4-flash generated, machine-audited) | chat + memory tasks | — |
| [stanfordnlp/coqa](https://huggingface.co/datasets/stanfordnlp/coqa) | multi-turn grounded QA | other (mixed provenance) |
| [rajpurkar/squad_v2](https://huggingface.co/datasets/rajpurkar/squad_v2) | abstention | cc-by-sa-4.0 |
| [dgslibisey/MuSiQue](https://huggingface.co/datasets/dgslibisey/MuSiQue) | multi-hop reasoning | undeclared on mirror |
| [ucinlp/drop](https://huggingface.co/datasets/ucinlp/drop) | arithmetic over passages | cc-by-sa-4.0 |
| [nayohan/multi_session_chat](https://huggingface.co/datasets/nayohan/multi_session_chat) | human-written persona facts | undeclared on mirror |

## Limitations

- **0.6B knowledge ceiling.** It confabulates on open-domain facts like any model this size;
  the training includes an abstention slice ("the passage doesn't say") but it is not a fix.
- **Note-taking can corrupt compound values** (an alphanumeric like `harbor858` was once noted
  as `8858` and then faithfully recalled wrong). What enters the ledger wrong stays wrong.
- **Ledger interference**: with 50+ accumulated entries, similar-typed values (several money
  amounts) can cross-contaminate. Training saw ~24 entries max.
- **No code in training data.** Reasoning about pasted code runs on the base model's residual
  ability.
- SFT never saw examples past 512 tokens; behaviour between 512 and 2,048 rides on the
  continued pretraining.
- English only.

## License

Apache 2.0, inheriting the Qwen3-0.6B base license. Training data licenses are listed above;
CoQA carries mixed-provenance terms and two mirrors declare no license — review them if you
redistribute derived data.