5.2 test-time compute scaling experiment

status: not-started
owner: unassigned
created: 2026-06-13
parent_plan: docs/PLAN.md §3 Tier 5 (5.2)
related_prds: 5.1-reasoning-on-22M.md (sibling — reasoning training; this PRD measures test-time not training-time)

PRD — Test-time compute scaling curve at 22M (Snell-methodology)

Goal

Run the Snell et al. 2024 methodology end-to-end at 22M: vary the test-time compute budget (N samples for Best-of-N, M tokens in CoT, K depth in tree search) and measure the quality-vs-FLOPs curve on a fixed eval (GSM8K suffices for V1). Produce a 22M data point for the Snell curve, which currently anchors on much larger models.

Tier-5 research arc. The most cleanly publishable arc in Tier 5 per PLAN.md — the methodology is established, the platform has all the test-time-compute primitives shipped (BoN ✅, MEMIT ✅, streaming-LLM ✅, KV cache ✅).

Why now

• posttrainllm bon --scan is already shipped — that’s half the work done. Plus the corpus-PPL verifier (also shipped) gives us the scoring signal Snell uses.
• 22M is a known posttrainllm scale; flagship-huge-v5 (221M) and the smaller pretrained configs let us actually do a cross-scale Snell curve as a stretch, not just one point.
• All of 5.2 reuses shipped infra; the only new code is the scaling-curve harness and the FLOPs accountant.

Scope — in

• New harness: posttrainllm scan-testtime <model> --task gsm8k --strategies bon,cot,tree --budgets 1,4,16,64,256 --out testtime-curve.jsonl
• Strategy implementations (all shipped or near-shipped):
  • BoN: posttrainllm bon already exists; wrap.
  • CoT-by-token-budget: run posttrainllm sample with --max-new N over a fixed prompt template; score by exact-match on the final answer.
  • Tree (lookahead): the experimental tree path mentioned in docs/decision_log.md — V1 may defer this strategy if not cheaply usable; if so, just BoN + CoT.
• FLOPs accountant: Sources/TinyGPTModel/FlopsCount.swift — closed-form FLOPs given (model config, prompt_tokens, gen_tokens, N_samples). Already roughly derivable from the bench code.
• Output: a JSONL row per (strategy, budget) with flops, quality, wall_seconds, passed/total.
• Plot: docs/research/testtime-22m-curve.md — quality-vs-FLOPs log-log plot + the cross-strategy comparison.

Scope — out

• Multi-task aggregate (a single “test-time quality” metric averaged over multiple evals). V1 = GSM8K only. Multi-task is a follow-up.
• Process reward models for tree decoding scoring — outcome rewards only.
• Cross-architecture scaling. Snell’s curve is one architecture varying scale; V1 is one architecture (posttrainllm flagship) at one size (22M). Two more sizes (50M, 221M) for the cross-size curve is the stretch.

Files to touch

Don’t touch

• posttrainllm bon — reuse, don’t fork.
• posttrainllm sample — V1 calls it as a library function (not subprocess) when in-process; subprocess fallback for clean isolation when FLOPs accounting needs the wall-time signal.

Acceptance criteria

• posttrainllm scan-testtime flagship-22m.tinygpt --task gsm8k --strategies bon,cot --budgets 1,4,16,64 --out /tmp/curve.jsonl produces a JSONL with 2 × 4 = 8 rows.
• Each row has consistent quality, flops, wall_seconds.
• Quality at budget=64 ≥ quality at budget=1 by ≥ 3pp (the experiment isn’t degenerate — more compute should help, at least marginally).
• The plot in testtime-22m-curve.md matches Snell’s log-log shape qualitatively (rising curve with diminishing returns).
• FLOPs accounting matches the bench-measured tok/s × tokens number within 20% (rough sanity).

Reference patterns

• Sources/TinyGPT/Bon.swift — the existing best-of-N scanner. Same JSONL shape; same scoring; just one of N strategies in the new harness.
• Snell et al. 2024 — the methodology. Cite. Their key plots are quality-vs-log-FLOPs; match that shape.

Open questions

• Whether to include the tree-decoding strategy in V1. Recommendation: ship BoN + CoT first; tree if the implementation is ≤1 day extra. Document the omission either way.
• Whether to do the cross-size stretch (22M + 50M + 221M) in this PRD or as a follow-up. Recommendation: stretch — same code path, just 3× the runs. Schedule one weekend.