aibuildai-inc/aibuildai-llm-posttrain-agent/checkpoint-selection

Checkpoint Selection Playbook

Pipeline contract (read first). In this pipeline each design trains exactly one model on the full training set in one attempt — there is NO k-fold / cross-validation, and NO multi-attempt search within a design. So only the within-run layers below apply: choose how to select/average checkpoints across the steps of that single training run. The cross-fold (Principle 09) and across-attempt / OOF-aggregation principles (10, 11) and any "per fold" wording are NOT used here — design-to-design variance comes from the Designer/Reviser exploring different designs, and combining across designs is the Aggregator's job (see the ensemble skill). Ignore fold/OOF/attempt framing in the principles when implementing.

A training run produces a sequence of parameter vectors over its steps. Which point in that sequence do you deploy?

The naive answer — pick the single step with the best validation score — is also the answer most prone to cherry-picking bias. Every selection over noisy validation introduces a positive bias of order σ·√(2 log N), where σ is the validation noise and N is the number of choices made.

This playbook catalogues the alternatives that mature competition pipelines actually use. They share one property: they aggregate instead of selecting a single point. Aggregation can be done in weight space (average parameter tensors into one model) or prediction space (predict with each checkpoint and average outputs). Both compose with downstream ensembling across designs.

The two within-run decision layers (the ones that apply here)

• Layer 1 — within-run, weight space. During the single training run, deploy a parameter average of multiple checkpoints rather than the single best-validation point. Inference cost stays at one forward pass.
• Layer 2 — within-run, prediction space. Keep K separate checkpoints from the one training run and average their predictions at inference. Inference cost scales linearly with K; gains are typically larger than weight-averaging when the K checkpoints come from different basins (e.g., snapshot ensemble cycles).

The two layers are orthogonal and compose: e.g. SWA within the run (Layer 1), then a multi-strategy prediction average across SWA+EMA+best (Layer 2). (A third "across attempts / across folds" layer exists in the principle files but does NOT apply in this single-model, single-attempt pipeline.)

When to consult this skill

• Designer / Reviser — when planning how a design's training will save and select checkpoints. Name an explicit checkpoint-selection rule in the plan (e.g., "best-val + EMA, predictions averaged at inference"). Leaving this implicit defaults to "max-spike cherry-pick", which is the unsafe choice on noisy validation.
• Coder — when implementing the chosen rule inside train.py for the single training run. Use Layer 1 / Layer 2 within-run rules only; there are no folds and no multiple attempts to be consistent across.
• Aggregator — for ensembling across designs, combine each design's raw_test outputs (averaging / weighted blend / rank / voting — see the ensemble skill). There is no OOF or attempt aggregation step inside one design.

Principle index

How to use

Read only the principles whose "When to consult" condition matches your current decision. Each principle is self-contained — pick one and you can use it as a design-level decision without reading the rest. Always read principle 12 (pitfalls) before committing to a plan.

Default stack for our scenarios

For small-to-medium datasets with noisy validation curves, the recommended default for the single training run is:

within-run:    best-val + EMA, predict with both, average (Principle 08)

This is the lowest-cherry-pick configuration achievable without modifying training infrastructure. Stronger configurations (SWA with cyclic LR + snapshot ensemble) are available via Principles 04 and 06 when training compute permits. (There is no across-attempt layer in this pipeline — one attempt per design.)

References (primary sources, verified)

• Izmailov, Podoprikhin, Garipov, Vetrov, Wilson 2018, "Averaging Weights Leads to Wider Optima and Better Generalization", arXiv:1803.05407 (SWA).
• PyTorch torch.optim.swa_utils documentation (SWA, EMA, batch-norm re-estimation).
• timm avg_checkpoints.py — production implementations of last-N and top-K weight averaging.
• ChemProp documentation, checkpoint averaging options (best, SWA, avg5, avg_many).
• Caruana, Niculescu-Mizil, Crew, Ksikes 2004, "Ensemble Selection from Libraries of Models", ICML.
• Huang, Li, Pleiss, Liu, Hopcroft, Weinberger 2017, "Snapshot Ensembles: Train 1, get M for free", arXiv:1704.00109.
• Kaggle ASL Fingerspelling 9th place writeup (last-N weight averaging in production).
• Kaggle TGS Salt 1st place writeup (snapshot ensemble × fold ensemble × architecture ensemble).

Cross-references

The ensemble playbook (aibuildai-builtin:ensemble) covers SWA, EMA, last-K averaging, and snapshot ensemble as ensemble techniques (Principles 1–4 of that skill). This skill reframes them as checkpoint-selection rules — same mechanism, different decision context. Consult the ensemble principle when you need to understand the underlying averaging mathematics; consult this skill's principle when you need to choose which rule a design's training script should follow.