---
title: "5.2 test-time compute scaling experiment"
sidebar:
  order: 502
---

<!-- PRD metadata (moved from frontmatter for Blume schema compatibility) -->
**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](https://arxiv.org/abs/2408.03314)
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

| File | Change |
|---|---|
| `Sources/TinyGPT/ScanTesttime.swift` | new — harness |
| `Sources/TinyGPTModel/FlopsCount.swift` | new — closed-form FLOPs |
| `Sources/TinyGPT/TinyGPT.swift` | `case "scan-testtime"` |
| `docs/research/testtime-22m-curve.md` | new — paper-shaped artifact |
| `evals/testtime-curve-smoke.sh` | new — 3 budgets × tiny model × 5 prompts; assert curve increases monotonically with budget |
| `docs/PLAN.md` | flip 5.2 ⬜ → ✅ on ship |

## 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](https://arxiv.org/abs/2408.03314) — 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.