marimo-0.6b-mlx

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 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

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.