gdiamos/amx-reasoning-v1-instruct

A causal language model trained end to end on one CPU core --- a single Intel Emerald Rapids core, bf16 through AMX, OMP_NUM_THREADS=1. 3,315,552 active parameters per token, 7,492,448 stored.

The point of the project is not that a small model runs on a CPU. It is that the architecture is derived from a single-core roofline, that training is confined to the same core, and that at this scale the interesting behaviours show up much earlier in the token budget than we expected.

Paper: Outrageously Small Neural Networks: Emergent Basic Reasoning at 6,616 tok/sec on One Intel AMX Core, shipped here as paper.pdf. This checkpoint is the model behind its evaluation table, so the numbers below and the numbers in the paper are the same numbers.

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 goals did Neil Rackers make? A 34-yard -- reference 34-yard, exact.

Q How many years did Manipur raid the Upper Chindwin region? A 4 -- reference 46. The passage contains the dates; the subtraction does not happen. This is the single clearest limitation of the model.

SQuAD v2 --- span extraction works; adversarial unanswerables mostly do not.

Q In what year was the Royal Dutch Petroleum Company founded? A 1890 -- reference 1890, exact.

Q Where did researchers study chimps in heavily forested regions? (the passage says researchers studied gorillas) A Ouesso district of the Sangha Region -- reference The passage does not say. It matched the rest of the question to the passage and answered. It does refuse correctly elsewhere: asked "Which species avoided the polar areas?" of a passage that never says, it answers The passage does not say.

Dolly --- terse and often correct, which the metric punishes.

Q What is the Kentucky Derby Trophy? A four trophies -- reference is a full paragraph beginning "The Kentucky Derby Trophy is a set of four trophies that are awarded to...". The answer is correct and scores F1 0.05. Read the Dolly column with this in mind.

Q From the passage note down the name of the countries which have most voting power. List the results. A The World Bank -- reference U.S., Japan, China, Germany, U.K. Asked for a list, it returned the subject. Multi-item extraction is not there.

The vocabulary mask

generation.json lists 800 token ids that must be masked before the argmax. This is not a style preference, and it is the single largest correctness fix in this repo.

Negatives for the sampled-softmax loss are drawn from a unigram proposal over the corpus, so a token that never occurs in training is never drawn as a negative, never receives downward gradient, and keeps a logit near 0. Every trained token that happens to be wrong in context is pushed to about -7.9. The untrained rows therefore win the argmax precisely when the model is uncertain. Measured on this checkpoint over 20 DROP prompts, the argmax was an untrained row in 11 of them.

Before masking, two such tokens --- ' ballo' and 'Frequently' --- were 29% of all answers on DROP, and 39% of numeric-answer rows contained no digit at all.

example.py applies it. If you write your own decode loop, apply it too.

Training data

Every natural-language and code source above is a curated artefact built with the help of large models --- quality classification, rephrasing, model-assisted extraction, and in the case of the reasoning corpus, traces that are themselves generated output. Training a model this small on them is a form of distillation, with no teacher present at training time. This is worth stating plainly, because it means results at this scale depend on corpora that did not exist when models of this size were last studied seriously.

Validation loss

The run's own numbers on its fixed held-out set, as logged. These are a sampled loss --- a (1 + n_negatives)-way discrimination, not a full-vocabulary one --- except val_flat, which is full-vocabulary.

{
  "val_flat": {
    "last_3": [
      4.0861921310424805,
      4.071656227111816,
      4.041394233703613
    ],
    "mean": 4.06641419728597
  },
  "val_loss": {
    "last_3": [
      3.921875,
      3.90625,
      3.875
    ],
    "mean": 3.9010416666666665
  },
  "steps": 61034,
  "tokens": 249999360
}

Limitations

A research artifact, and the honest summary is that the failures are specific rather than diffuse.

It cannot calculate over a passage. It retrieves and compares reliably and fails on anything needing arithmetic, returning a plausible, well-formed, wrong number. See the DROP examples above. This is consistent with the paper's capability ladder, where two-digit addition sits at 47% while in-context induction is at 93%.

It answers a third of unanswerable questions anyway. Missed abstention is 62.9% on SQuAD v2's adversarial unanswerables, which differ from an answerable question by a single word. It over-abstains too, on 16.1% of answerable rows. Both directions are live.

It is extractive, not generative. Give it a passage. Without one it has very little to say, and multi-item answers are beyond it.

It has had no alignment, safety, or preference training of any kind, and will reproduce the biases and errors of its training corpus.

License and provenance

Apache 2.0, for the weights and for the source in m2r/; full text in LICENSE.

The training data is a mixture of public code, web, math and instruction corpora, with per-source token counts above. Those corpora carry their own terms, which the Apache licence on this model does not alter and does not extend to them.

Produced by tools/export_hf.py from run sft-qa-v2-stacked-20260906T163025 at step 61035.