paper-sao-single-rollout-1.5b

Recipe: recipes/papers/sao-single-rollout · Collection: Papers, replicated

Single-rollout policy gradient with a one-layer critic on GSM8K, where the recipe masks tokens whose importance ratio leaves (0.7, 6.0). value_head.pt is the critic; train_log.json is the per-step record.

Result

Arm	pass@1	95% CI	pass@4	Steps	GPU min
Base, no training	0.36	[0.30, 0.44]	0.59	0	0
Baseline (ratio applied, never masked)	0.46	[0.40, 0.53]	0.72	10 + 40	25.7
Recipe (wai.SAO(): masked outside (0.7, 6.0))	0.46	[0.40, 0.53]	0.71	10 + 40	20.1

Recipe vs baseline: -0.004 [-0.040, +0.031] over 120 paired tasks. The band masked 0.03% of tokens, so the two arms ran the same experiment. Both went 0.36 to 0.46 from the critic-baselined single-rollout gradient, not from the mask. The recipe now logs masked_share per step so the reachability check happens before the GPU hour.

Arms in this repo

The root holds the arm the recipe README's headline number reports. Every other arm is a subfolder named after it. checkpoints/ never ships.

folder	arm
.	recipe arm: wai.SAO() masked update, 2026-09-21 run

Load

from peft import PeftModel
from transformers import AutoModelForCausalLM

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
model = PeftModel.from_pretrained(base, "while-ai/paper-sao-single-rollout-1.5b")  # the headline arm

Reproduce

git clone https://github.com/whilehq/whileai-sdk && cd whileai-sdk/recipes/papers/sao-single-rollout
python recipe.py

The recipe README pins the seed, the library versions and the GPU, and its Checks table says what the eval verified. Read the Learned section before quoting a number from this card.