Marimo 0.6B (MLX)
Marimo is a research retrofit of 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 perkey: valuefact, 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
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.