amx-reasoning

What it does

It answers questions about a passage you give it, in one or two words, and it stops. On held-out extractive QA it reaches 18.2% exact match and 23.2% F1, and on DROP specifically 25.0% EM against a 7.5% majority baseline.

That is a weak model by contemporary standards and a surprising one for its size. The paper's framing applies here too: what is interesting is not the absolute score but that a model with 3,315,552 active parameters, trained on one core, does passage-grounded retrieval at all --- and that its remaining failures are specific and nameable rather than general incompetence. It retrieves and compares; it cannot calculate.

Running it

AutoModelForCausalLM.from_pretrained will not work: the architecture is not one transformers knows --- chunked sliding-window attention interleaved with log-decay linear attention, and a tied readout. The model's own source ships here under m2r/, unmodified from the repository that trained it.

hf download gdiamos/amx-reasoning-v1-instruct --local-dir amx-reasoning-v1-instruct
cd amx-reasoning-v1-instruct && pip install -r requirements.txt && python example.py

example.py is the whole thing, and it is short. The two parts that are not optional:

from m2r.data.templates import EOT, render_prompt

# 1. THE PROMPT FORMAT the model was tuned on, imported rather than copied so
#    it cannot drift. It already emits BOS -- do not prepend one.
prompt = render_prompt([f"{{question}}\n\n{{passage}}"], thinking=False)
ids = tok.encode(prompt, add_special_tokens=False).ids

# 2. THE VOCABULARY MASK from generation.json. See "The vocabulary mask" below;
#    without it this model answers " ballo" to a fifth of DROP questions.
BAN = torch.tensor(json.loads(open("generation.json").read())["banned_token_ids"])
logits[BAN] = -1e30
next_id = int(logits.argmax())        # greedy: answers are one or two words

Decode greedily and stop on EOT. Sampling is the right call for the base model and the wrong one here --- this checkpoint was tuned to emit a short span and halt.

What is in this repo

Architecture

Attention is confined to document boundaries: a training window packs many documents, and without isolation sliding-window attention reaches into its neighbours while linear attention carries state across the whole window.

The shape is deliberate. One AMX core sustains roughly 2,231 GF/s of bf16 matrix multiply at these dimensions but pays a 1.4--1.5 microsecond floor per GEMM dispatch, so every design choice here is about issuing few large matrix multiplies rather than many small ones.

Lineage

The token count for this checkpoint is only its last stage. Three runs, each starting from the previous one's weights:

About 5.4B tokens in total, at roughly 1,481 tokens per active parameter --- far past compute-optimal, deliberately, because the target is inference cost rather than training cost. Applying the QA stage on top of the instruction-tuned model rather than directly on the base was worth +1.3 EM, +2.2 F1, and 8.7 points less over-abstention.

Evaluation

Held-out extractive QA, greedy decoding, generation.json mask applied. EM and F1 are over answerable rows only. Over-abstention is refusing when the answer was present; missed abstention is answering when it was not.

Examples

Real greedy outputs, two per source, chosen to show a success and the characteristic failure rather than a highlight reel.

DROP --- retrieval and comparison work; arithmetic does not.

Q Which field g