qxlsz/Qwen3.5-0.8B-weight-space-programming-mlx

Qwen3.5-0.8B Weight-Space Programming (MLX)

A small model that stops getting stuck in thinking loops, running 1.7x faster on a Mac than PyTorch.

This is an MLX conversion of RikiyaT/Qwen3.5-0.8B-weight-space-programming, at the checkpoint's native bfloat16, with no quantization. It handles text and images, and takes about 1.6 GB on disk.

The source model is Qwen3.5-0.8B with two tensors edited directly, with no training, to reduce the base model's thinking loops on math problems. On 40 MATH-500 problems on MLX, it solved 21 against 5 for base and hit the 32,768-token limit on 8 against 35. Its 52.5% accuracy matches the 51.9% that the source repo's published PyTorch scores average on those same problems. On an Apple M5 Max it generates 217 tokens/s against 125 for PyTorch on the GPU. See Speed and Evaluation.

All credit for the model and the edit belongs to its authors: see the blog post and the source model card for the method and full results. This repo only changes the format.

Usage

pip install mlx-lm

from mlx_lm import load, generate
from mlx_lm.sample_utils import make_sampler, make_logits_processors

model, tokenizer = load("qxlsz/Qwen3.5-0.8B-weight-space-programming-mlx")

prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": "What is $2^{10} + 24$?\nPlease reason step by step, and put your final answer within \\boxed{}."}],
    tokenize=False, add_generation_prompt=True, enable_thinking=True)

print(generate(
    model, tokenizer, prompt, max_tokens=32768, verbose=True,
    sampler=make_sampler(temp=0.6, top_p=0.95, top_k=20),
    logits_processors=make_logits_processors(presence_penalty=1.5, presence_context_size=32768),
))

Set presence_context_size as shown: the mlx-lm default looks back over only 20 tokens, while the source model's published results penalize everything generated so far. Those settings are its thinking-mode protocol.

For images, use mlx-vlm, which loads this repo directly.

What was checked

scripts/verify_weights.py in this repo compares raw tensors against the source checkpoint at revision d1671a178705bcab7644806d23c535517f35723b:

Source tensors	Count
Byte-identical	393
Reordered to MLX axis order, values exact	19
Stored as x + 1.0, MLX's RMSNorm convention	61
mtp.* draft-model tensors, dropped	15

The two edited tensors match the sha256 values published on the source model card, byte for byte:

Tensor	sha256
language_model.model.layers.15.mlp.down_proj.weight	82e2d733ea64feaeb4af93a1e065641d3df98b494a8b7576ef2f52e90fe2ee53
language_model.model.layers.19.mlp.down_proj.weight	8ef3c61bc36af113ad58e89dbb0a4a14dcba5a45cb750c15ab515c82ddd12eb2

Compared with PyTorch at bfloat16 over 1,264 token positions from 4 prompts, the top prediction agrees at 99.37%, with mean KL divergence of 7.06e-4.

The tokenizer and chat template matched the source on a test prompt mixing math, CJK, emoji and special tokens.

Differences from a stock conversion and from the source

• linear_attn.norm.weight stays float32. --dtype bfloat16 would narrow these 18 tensors, as scripts/restore_norm_precision.py explains. Restoring them has no measurable speed cost.
• The 61 RMSNorm weights are slightly rounded. Storing x + 1.0 in bfloat16 rounds by up to 0.0156. Keeping them at float32 is closer to the source but ran about 4x slower, so they are left as is.

Speed on a Mac: MLX vs PyTorch

Same weights on one Apple M5 Max (36 GB), greedy decoding, each run in a fresh process, values are medians. PyTorch ran on the GPU through MPS with transformers at bfloat16.

Measure	MLX	PyTorch (MPS)	Change with MLX
Generation, 512 tokens	220 tok/s	129 tok/s	1.71x faster
Generation, 4,096 tokens	217 tok/s	125 tok/s	1.73x faster
Generation after a 3,772-token prompt	214 tok/s	123 tok/s	1.74x faster
Time to first token, 48-token prompt	34 ms	113 ms	3.3x faster
Time to first token, 3,772-token prompt	178 ms	556 ms	3.1x faster
Total time, 4,096 tokens	18.9 s	32.8 s	42% less
Model load	0.62 s	1.05 s	1.7x faster
Peak memory, 512 tokens	2.20 GB	2.96 GB	26% less
Peak memory, 3,772-token prompt	3.40 GB	2.99 GB	14% more

PyTorch on CPU reached 0.35 tok/s in one run capped at 5 minutes, which is not a practical way to use this model.

Limits of this comparison:

• It covers one machine and one model. Other Macs will differ.
• PyTorch used its reference code paths for this model's linear attention. causal_conv1d did not build on this Mac, and transformers still fell back after flash-linear-attention was installed.
• The machine was not fully idle, with a load average around 3 to 4, though backends alternated.
• Outputs are not token-for-token identical. Greedy runs first differed at token 180 on the short prompt and token 30 on the long one, consistent with floating-point differences between frameworks. Accuracy is the comparison that matters here, and it matches.

Raw data is in eval/bench.jsonl, from scripts/run_bench.py and scripts/bench_latency.py.

Evaluation on MLX

The edit's behaviour reproduces on MLX. Evenly spaced MATH-500 problems, one sample each, on an Apple M5 Max, using the source card's thinking-mode protocol: temperature 0.6, top_p 0.95, top_k 20, presence penalty 1.5 over all generated tokens, 32,768 max tokens, last \boxed{} graded with math-verify.

Model	Problems	Accuracy (95% CI)	Hit token limit (95% CI)	Mean tokens
Base Qwen3.5-0.8B, MLX	40	12.5% (5.5-26.1)	87.5% (73.9-94.5)	30,009
This model	40	52.5% (37.5-67.1)	20.0% (10.5-34.8)	19,310

Both ran the same 40 problems. This model alone solved 16 and base alone solved none, an exact McNemar p of 0.00003. The test was fixed in advance in docs/ANALYSIS_PLAN.md.

MLX does not change accuracy, and should not, since these are the same weights: on these problems the source repo's published PyTorch scores average 51.9%, a difference of +0.6% (95% CI -10.0% to +10.6%). Its full MATH-500 figures are 51.0% accuracy and 31.6% hitting the limit. This subset is small and its limit rate is lower, so use the source card's numbers for the model's overall performance.

Raw outputs are in eval/, and scripts/eval_math.py, scripts/grade.py and scripts/compare_paired.py reproduce them.

Why there is no 4-bit version

Plain 4-bit quantization did not preserve quality. On the same 40 problems it hit the limit on 27.5%, so loops stayed reduced, but it solved 12 against 21 for this model. This model alone was right on 10 problems and the 4-bit build alone on 1. Its raw outputs are in eval/edited-4bit.jsonl.

License

Apache 2.0, inherited from the source model.