PursuitOfDataScience/Argonne-4.0-think

Argonne 4.0-think

Argonne 4.0-think is a 1.04B-parameter reasoning model trained from scratch, built on argonne-4.0-base. It emits an explicit <think>…</think> trace and then a \boxed{} answer.

It is the reasoning model of the 4.0 line, whose thesis is data efficiency: 4.0-base reaches its quality on 65.12B tokens at 1.04B parameters. Its larger sibling Argonne-3.5-think (2.88B) remains the stronger reasoner — and was in fact this model's teacher (see Training). At 2.8× fewer parameters this model gives up 4.93 points on the four-pool gate below.

Evaluation

Greedy decoding, paired against the comparison model on identical items within the same evaluation run, exact McNemar on the paired outcomes. n = 1000 (ASDiv, SVAMP), 500 (GSM-Plus, MAWPS).

What the post-training campaign bought

Against its own starting point — the SFT → DPO → CoT-SFT model, before any distillation:

pool	CoT-SFT start	this model	delta	p
ASDiv	64.30	69.80	+5.50	1.4e-04
SVAMP	49.10	59.80	+10.70	6.3e-11
GSM-Plus	27.20	36.00	+8.80	4.1e-05
MAWPS	51.20	59.80	+8.60	2.0e-05
four-pool pooled	50.87	59.17	+8.30	1.0e-21

Against the 2.88B sibling

pool	Argonne-3.5-think (2.88B)	4.0-think (1.04B)	delta	p
ASDiv	73.60	69.80	−3.80	1.1e-02
SVAMP	68.10	59.80	−8.30	6.9e-07
GSM-Plus	40.80	36.00	−4.80	5.7e-02
MAWPS	60.40	59.80	−0.60	8.1e-01
four-pool pooled	64.10	59.17	−4.93	3.0e-08

The deficit is concentrated in SVAMP and ASDiv; MAWPS is statistically indistinguishable.

Test-time compute, measured on these weights

pool	greedy	self-consistency@8	budget-extend	pass@8
ASDiv	69.80	78.30	72.80	87.80
SVAMP	59.80	70.10	62.10	85.60
GSM-Plus	36.00	44.80	36.60	63.40
MAWPS	59.80	64.00	62.40	71.20
pooled	59.17	67.60	61.47	80.23

Sampling is worth more than a bigger single pass. Self-consistency at K=8, temperature 0.8 adds +8.43 pooled — most of the way to closing the gap to the 2.88B sibling's greedy 64.10. pass@8 of 80.23 against a greedy 59.17 says the deficit is probability mass, not reach: the model can very often find the answer, and selection is what is missing.

Which pools, and why not the usual ones

GSM8K is excluded, not unreported-by-accident. It is contaminated for Argonne reasoning models (the CoT-SFT mix saw ~94% of its test set). MATH-500 is also excluded for this line: 17 of its 319 items have a near-duplicate in the training mix, measured by Jaccard near-dup search, so it is not a clean pool for any model on this branch. The four pools above are clean by that same measure. GSM-Plus is adversarially perturbed GSM8K test, and the mix's GSM8K tier was audited to be 4,338/4,338 from the train split with zero test items.

General capability is not characterized for this model. The base was evaluated with lm-eval; the reasoning fine-tune was not. Do not read the arithmetic numbers above as evidence about general ability in either direction.

Training

stage	data	detail
base	—	argonne-4.0-base, 65.12B tokens, ctx 65,536
1 — SFT	UltraChat 200k	1 epoch, effective batch 20
2 — DPO	argilla/dpo-mix-7k	effective batch 8
3 — CoT-SFT	short-trace mix, 28,428 rows	1 epoch, effective batch 12 → the 50.87 starting point above
4 — on-policy distillation	the student's own rollouts, correct-labelled	per-token reverse-KL from Argonne-3.5-think as teacher, iterated over several rounds
5 — CE repair	24,787 of the model's own correct traces	plain cross-entropy, LR 3e-6, 1,055 steps, seed 46

Two things in that table did the work, and both were counter-intuitive:

Stage 4 is distillation from the larger sibling, on the student's own outputs. Per-token reverse-KL against Argonne-3.5-think, computed at states the student actually visits rather than on the teacher's traces. Imitating the teacher's traces directly did not transfer; matching its per-token distribution at the student's own states did.

Stage 5's learning rate is the whole stage. The same cross-entropy pass on the same data at LR 1e-5 damages the model; at 3e-6 it helps. And ordering matters — running the repair last, after distillation, beat running it earlier.

The final four candidate arms of the campaign are statistically tied on the gate (this one at 59.17, the others at 59.60 / 59.10 / 59.00, all p ≥ 0.35 against it). This checkpoint was released as the recorded release candidate; the ranking within that group is noise, not a preference.

Inference

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "PursuitOfDataScience/Argonne-4.0-think"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True, dtype=torch.bfloat16
).cuda()

messages = [{"role": "user", "content": "A shop sells pencils 3 for $2. How much do 12 pencils cost?"}]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
ids = tokenizer(text, return_tensors="pt")["input_ids"].cuda()

out = model.generate(ids, max_length=ids.shape[1] + 512, do_sample=False)
print(tokenizer.decode(out[0][ids.shape[1]:], skip_special_tokens=True))

For throughput, prefer vLLM/SGLang over .generate().

Usage notes

• Load with trust_remote_code=True; config.json carries an auto_map so the custom argonne2 classes resolve without manual setup.
• The custom generate takes max_length (total length), not max_new_tokens.
• eos_token_id is 151645 (<|im_end|>) so the assistant turn ends cleanly. Verified on the published files: a chat-templated prompt with no eos_token_id argument terminates on its own in 25 tokens.
• lm_head.weight is reported missing on load. Expected and benign — embeddings are tied.
• Traces are short by design (~235 think tokens on ASDiv greedy). This line's measured result is that lengthening traces does not help it; budget-extend adds only +2.30 pooled.
• interleaved_local_attention is published as false and local_attention_window as null. This is deliberate and matches argonne-4.0-base: model.py applies the sliding window only on the flash-attn-2 path, every Argonne pretrain ran on flash-attn-4 (which does not expose flash_attn.flash_attn_interface), so the window was never active and these weights only ever saw full attention. Publishing the flags enabled would hand a 256-token window on odd layers to any user who happens to have flash-attn-2 installed, on weights that never saw one.

Limitations

• Context length 65,536 is inherited from the base and was not re-measured after the reasoning fine-tune. The base's window is trained rather than extrapolated (substantiated at its release by position-bucketed NLL on held-out arXiv); the fine-tuning stages ran at a 4,096-token block, and long-context behaviour after them is uncharacterized. Treat long inputs as untested here.
• Selection, not reach, is the binding constraint. pass@8 80.23 vs greedy 59.17 — use self-consistency if you can afford K samples.
• pass@K is a noisy metric on this line. Treat it as a ceiling indicator; select on greedy or self-consistency.
• Grade-school and early-competition arithmetic word problems are the measured domain. Code, tool-calling and general-purpose chat are not characterized.
• General capability is unmeasured for this checkpoint (see Evaluation).
• 1.04B parameters on a 65.12B-token base — far below frontier compute, and 2.8× smaller than Argonne-3.5-think, which beats it by 4.93 points.
• No safety alignment beyond what UltraChat and the preference data provide.

Source code

Everything is on the GitHub main branch — PursuitOfDataScience/ArgonneAI.

file	role
model.py	the argonne2 architecture, identical to the copy in this repo
reasoning/thinking_training.md	the full build log for the reasoning line
reasoning/stage_a4_think_hf.py	staged this repo, and documents every config fix applied (not yet on main)
reasoning/hub_config_audit.py	audits the live Hub config of every Argonne repo against the release invariants

Citation

@misc{argonne40think,
  title  = {Argonne 4.0-think},
  author = {Youzhi Yu},
  year   = {2026},
  url    = {https://huggingface.co/PursuitOfDataScience/Argonne-4.0-think}
}