README.md

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

# Marimo 0.6B (MLX)

Marimo is a research retrofit of [Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B)
into a hybrid autoregressive + discrete-diffusion chat model with two mechanisms trained in:

- **Adaptive thinking blocks.** The hidden `<think>` region is generated as a chain of
  MDLM-style denoised blocks. Block size is a per-thought decision: the model
  autoregressively emits a size token (`<sz32>`, `<sz64>`, `<sz128>`), the engine allocates
  that many masked positions and denoises them in parallel (16 steps per block by default),
  and the loop repeats until the model emits `</think>`. The visible answer is then decoded
  autoregressively.
- **Ledger memory.** Only the last 4 messages stay in context. Each turn's thinking blocks
  are kept as a note; when a turn falls out of the window its note is merged into a
  `Known so far:` ledger carried in the system turn (one entry per `key: value` fact,
  latest value wins). The prompt therefore stays near-constant in size regardless of
  conversation length.

This repository is the standalone MLX port: it depends only on `mlx`, `numpy` and
`tokenizers`, ships fp16 weights (`model.safetensors`, ~1.2 GB), and runs on Apple Silicon.

## Usage

```bash
pip install -r requirements.txt
python chat.py
```

`chat.py` is an interactive terminal REPL. Useful flags: `--show-thinking` prints the
denoised thought blocks after each turn, `--temperature` / `--top-p` control sampling,
`--keep-messages` resizes the visible window, `--system` replaces the default system
prompt. The default system prompt frames the model as a note-taking intake assistant,
which is the register most of the fine-tuning data uses.

Requires macOS on Apple Silicon (MLX). Maximum sequence length is 2048 tokens; the REPL
refuses turns that would overflow it.

## Training

Starting from Qwen3-0.6B, the model went through continued pretraining and then SFT on a
synthetic conversational corpus (~15k dialogues with per-turn thinking traces, size-token
supervision and note annotations), trained on a single RTX 3090. The tokenizer adds special
tokens for the mask, thought padding and the size menu; everything else is the Qwen3
tokenizer unchanged.

## Parity with the torch reference

The port is validated against the torch training stack before each release: mask builders
and the ChatML+ledger rendering match bit/byte-exactly; teacher-forced logits of the
shipped fp16 weights agree with the fp32 torch reference on top-1 for 99.97% of 13,156
probed rows (mean KL ≈ 4.6e-6); cached and uncached forwards agree within 1.6e-4.

On an M1 Pro at the production config, full turns take ~2.2–4.2 s and autoregressive
decoding runs at ~15–24 tok/s with a 500-token prefix.

## Evaluation snapshot

On an internal 40-turn "ledger needle" probe (10 planted facts recalled at the end,
scored on the value core), this model recalled 6/10 at a constant ~500-token prompt,
while the Qwen3-0.6B base with the full transcript in context recalled 1/10 — the base
degenerates into echoing the probe in long repetitive multi-turn chat. Single seed,
and the comparison bundles mechanism and training data, so read it as direction rather
than a benchmark result.

## Limitations

- 0.6B parameters: this is a mechanism study, not a general assistant. Expect factual
  errors and brittleness outside the conversational note-taking register it was tuned on.
- English-only training data.
- 2048-token context; long documents do not fit.
- The ledger only records what the model chose to note during thinking; facts it never
  noted are lost once the turn leaves the window.

## License

Apache-2.0. The weights derive from Qwen3-0.6B (Apache-2.0); `model.py` vendors a
minimal Qwen3 forward pass in MLX.