Does eight-token score centering beat importance-sampled REINFORCE under stale rollouts?
The guessed StableReinforce loss subtracts the mean log-probability of eight sampled tokens in the backward pass. We compare GSM8K accuracy and training stability against plain importance-sampled policy gradient with stale
Budget
2 GPU-h (default)
· each run plans 2 GPU-h, 120 min wall cap
· 2 runs
·
change
Reproducible on Spindle: yes
— The candidate loss is a custom per-token surrogate. It runs through forward_backward_custom on a LoRA policy, with rollouts and auxiliary-token logprobs from the sampling client, on public GSM8K.
Spec
Eight-token score centering vs importance-sampled policy gradient, 4-step-stale rollouts
Hypothesis. With rollouts generated by a policy 4 optimizer steps stale, the eight-token centered-backward loss reaches at least 4 points higher GSM8K test accuracy at step 100 than the importance-sampling policy-gradient loss, or avoids a collapse that the baseline shows.
Method. Train a LoRA policy on GSM8K train with a 0/1 exact-match reward and group-mean advantages (group_size 8, 16 groups per batch). To make training off-policy, refresh the sampler weights only every 4 optimizer steps, so rollouts come from a policy up to 4 steps stale. Auxiliary tokens: for each completion position, request the sampler's top-20 prompt logprobs, renormalize them, and draw 8 auxiliary tokens.
Qwen/Qwen3.5-0.8B
RL
100 steps
batch 128
rank 32
GSM8K
eval GSM8K test · exact-match accuracy (greedy)
2 GPU-h
120 min cap
Success criterion. Validated if the treatment's GSM8K accuracy at step 100 exceeds the baseline by at least 4 points (about 2x the standard error at n=500), or if the baseline's accuracy drops more than 10 points from its peak while the treatment's does not. Failed if the gap is under 2 points with no collapse difference.
full spec
{
"base_model": "Qwen/Qwen3.5-0.8B",
"baseline": "The same RL loop, staleness, seeds and hyperparameters, but using Tinker's built-in importance_sampling policy-gradient loss instead of the centered-backward surrogate.",
"batch_size": 128,
"budget": {
"gpu_hours": 2.0,
"max_wall_minutes": 120
},
"command": "# custom tinker SDK script (custom parts: stale sampler refresh, auxiliary-token datums, centered loss)\nimport os, tinker, torch\nsc = tinker.ServiceClient(base_url=os.environ[\"TINKER_BASE_URL\"])\ntc = sc.create_lora_training_client(base_model=\"Qwen/Qwen3.5-0.8B\", rank=32)\nsampler = tc.save_weights_and_get_sampling_client(name=\"step0\")\n\ndef centered_loss(data, logprobs_list):\n # data grouped as [target, aux1..aux8] per sequence; per-token advantages A in each target datum's loss_fn_inputs\n loss = 0.0\n for i in range(0, len(logprobs_list), 9):\n tgt = logprobs_list[i]\n aux = torch.stack(logprobs_list[i+1:i+9]).mean(0)\n A = data[i].loss_fn_inputs[\"advantages\"].to_torch()\n loss = loss + (-A * tgt + A * (aux - aux.detach())).sum()\n return loss, {}\n\nfor step in range(100):\n if step % 4 == 0:\n sampler = tc.save_weights_and_get_sampling_client(name=f\"s{step}\")\n batch = rollout_gsm8k(sampler, group_size=8, groups=16, max_tokens=512, topk_prompt_logprobs=20) # cookbook gsm8k env + reward\n datums = build_target_and_aux_datums(batch, n_aux=8) # 8 sampled tokens per position from renormalized top-20\n tc.forward_backward_custom(datums, centered_loss).result()\n tc.optim_step(tinker.AdamParams(learning_rate=4e-5)).result()\n# baseline arm: same loop with tc.forward_backward(target_datums, loss_fn=\"importance_sampling\")",
"dataset": {
"hf_path": "openai/gsm8k",
"n_examples": 12800,
"name": "GSM8K",
"split": "main/train"
},
"eval": {
"benchmark": "GSM8K test",
"metric": "exact-match accuracy (greedy)",
"n_examples": 500
},
"hypothesis": "With rollouts generated by a policy 4 optimizer steps stale, the eight-token centered-backward loss reaches at least 4 points higher GSM8K test accuracy at step 100 than the importance-sampling policy-gradient loss, or avoids a collapse that the baseline shows.",
"lora_rank": 32,
"max_seq_len": 1024,
"method": "Train a LoRA policy on GSM8K train with a 0/1 exact-match reward and group-mean advantages (group_size 8, 16 groups per batch). To make training off-policy, refresh the sampler weights only every 4 optimizer steps, so rollouts come from a policy up to 4 steps stale.\n\nAuxiliary tokens: for each completion position, request the sampler's top-20 prompt logprobs, renormalize them, and draw 8 auxiliary tokens. This approximates true sampler draws and is a documented deviation from the original algorithm.\n\nTreatment loss (custom, via forward_backward_custom): sum over tokens of -A*logp(target) + A*(mean_aux_logp - stopgrad(mean_aux_logp)). kl_coeff is set to 0 so only the centering differs from the baseline. Auxiliary logprobs come from 8 extra datums per sequence that share model_input but use the auxiliary tokens as target_tokens.\n\nBaseline loss: the built-in importance_sampling loss with sampler logprobs. All other settings are identical: model, data order, seeds, advantages, staleness, LR and renderer.\n\nEvaluation: sample greedily from saved checkpoints every 25 steps. Also log grad norm and mean policy/sampler logprob gap as stability signals.",
"optimizer": {
"lr": 4e-05,
"name": "adam"
},
"steps": 100,
"success_criterion": "Validated if the treatment's GSM8K accuracy at step 100 exceeds the baseline by at least 4 points (about 2x the standard error at n=500), or if the baseline's accuracy drops more than 10 points from its peak while the treatment's does not. Failed if the gap is under 2 points with no collapse difference.",
"title": "Eight-token score centering vs importance-sampled policy gradient, 4-step-stale rollouts",
"training_type": "rl"
}Runs
| status | summary | took | started by | started |
|---|---|---|---|---|
| Failed no result | Stock Tinker custom-loss and sampler probes worked on the permitted Qwen3.5-2B fallback, but the paired GSM8K run and an initialization retry failed because Spindle reported trainer capacity exhausted. | 15m 50s | — | 1d ago |
| Succeeded no result | RL on openai/gsm8k with Qwen/Qwen3.5-2B (LoRA rank 32, lr 4e-05, max_tokens 512), eval on 256 held-out problems with greedy decoding. | 28m 27s | — | 1d ago |
Source
raw text
Tinker drops a new closed-source loss method for async rl stability https://x.com/tinkerapi/status/2113019700661862805 (docs: https://tinker-docs.thinkingmachines.ai/tinker/losses/stable-reinforce/) [nan] can't tell exactly what they use from there but if you notice slowdown in training step then it's likely score centering [nan] given score centering + some improvements on top of sc is what big labs use [nan] my blog's setting is far more off policy than theirs. https://ampersand.humansand.ai/blog/ tried some non-IS method which also makes sense [aaron] It is likely score centering but the estimator is different from the SGLang one [aaron] Current candidate: stable_reinforce_candidate(logits, target_tokens, advantages, auxiliary_tokens, sampling_logprobs=None, kl_coeff=0.0): summed surrogate with CE-valued forward and centered backward. Eight cached auxiliary tokens from the sampler; policy = -advantage * target_logp + advantage * (auxiliary_mean_logp - auxiliary_mean_logp.detach()); optional analytic KL term kl_coeff * (sampling_logprobs - target_logp). [aaron] https://app.notion.com/p/modal-com/StableReinforce-eight-token-score-centering-and-the-SGLang-comparison-4011e7f1694981c8ba8ceb04bbc768c9 Okay, I think I have the algorithm. nan is right, it is just score centering but not with the default estimator. Held-out multi-token policy predictions: mean relative error 6.94e-8, cosine ~1. Two-sequence batch: 0.0198%.
triage output
{
"adjacent_experiment": "",
"reproducibility_reason": "The candidate loss is a custom per-token surrogate. It runs through forward_backward_custom on a LoRA policy, with rollouts and auxiliary-token logprobs from the sampling client, on public GSM8K.",
"reproducible_on_spindle": true,
"spec": {
"base_model": "Qwen/Qwen3.5-0.8B",
"baseline": "The same RL loop, staleness, seeds and hyperparameters, but using Tinker's built-in importance_sampling policy-gradient loss instead of the centered-backward surrogate.",
"batch_size": 128,
"budget": {
"gpu_hours": 2.0,
"max_wall_minutes": 120
},
"command": "# custom tinker SDK script (custom parts: stale sampler refresh, auxiliary-token datums, centered loss)\nimport os, tinker, torch\nsc = tinker.ServiceClient(base_url=os.environ[\"TINKER_BASE_URL\"])\ntc = sc.create_lora_training_client(base_model=\"Qwen/Qwen3.5-0.8B\", rank=32)\nsampler = tc.save_weights_and_get_sampling_client(name=\"step0\")\n\ndef centered_loss(data, logprobs_list):\n # data grouped as [target, aux1..aux8] per sequence; per-token advantages A in each target datum's loss_fn_inputs\n loss = 0.0\n for i in range(0, len(logprobs_list), 9):\n tgt = logprobs_list[i]\n aux = torch.stack(logprobs_list[i+1:i+9]).mean(0)\n A = data[i].loss_fn_inputs[\"advantages\"].to_torch()\n loss = loss + (-A * tgt + A * (aux - aux.detach())).sum()\n return loss, {}\n\nfor step in range(100):\n if step % 4 == 0:\n sampler = tc.save_weights_and_get_sampling_client(name=f\"s{step}\")\n batch = rollout_gsm8k(sampler, group_size=8, groups=16, max_tokens=512, topk_prompt_logprobs=20) # cookbook gsm8k env + reward\n datums = build_target_and_aux_datums(batch, n_aux=8) # 8 sampled tokens per position from renormalized top-20\n tc.forward_backward_custom(datums, centered_loss).result()\n tc.optim_step(tinker.AdamParams(learning_rate=4e-5)).result()\n# baseline arm: same loop with tc.forward_backward(target_datums, loss_fn=\"importance_sampling\")",
"dataset": {
"hf_path": "openai/gsm8k",
"n_examples": 12800,
"name": "GSM8K",
"split": "main/train"
},
"eval": {
"benchmark": "GSM8K test",
"metric": "exact-match accuracy (greedy)",
"n_examples": 500
},
"hypothesis": "With rollouts generated by a policy 4 optimizer steps stale, the eight-token centered-backward loss reaches at least 4 points higher GSM8K test accuracy at step 100 than the importance-sampling policy-gradient loss, or avoids a collapse that the baseline shows.",
"lora_rank": 32,
"max_seq_len": 1024,
"method": "Train a LoRA policy on GSM8K train with a 0/1 exact-match reward and group-mean advantages (group_size 8, 16 groups per batch). To make training off-policy, refresh the sampler weights only every 4 optimizer steps, so rollouts come from a policy up to 4 steps stale.\n\nAuxiliary tokens: for each completion position, request the sampler's top-20 prompt logprobs, renormalize them, and draw 8 auxiliary tokens. This approximates true sampler draws and is a documented deviation from the original algorithm.\n\nTreatment loss (custom, via forward_backward_custom): sum over tokens of -A*logp(target) + A*(mean_aux_logp - stopgrad(mean_aux_logp)). kl_coeff is set to 0 so only the centering differs from the baseline. Auxiliary logprobs come from 8 extra datums per sequence that share model_input but use the auxiliary tokens as target_tokens.\n\nBaseline loss: the built-in importance_sampling loss with sampler logprobs. All other settings are identical: model, data order, seeds, advantages, staleness, LR and renderer.\n\nEvaluation: sample greedily from saved checkpoints every 25 steps. Also log grad norm and mean policy/sampler logprob gap as stability signals.",
"optimizer": {
"lr": 4e-05,
"name": "adam"
},
"steps": 100,
"success_criterion": "Validated if the treatment's GSM8K accuracy at step 100 exceeds the baseline by at least 4 points (about 2x the standard error at n=500), or if the baseline's accuracy drops more than 10 points from its peak while the treatment's does not. Failed if the gap is under 2 points with no collapse difference.",
"title": "Eight-token score centering vs importance-sampled policy gradient, 4-step-stale rollouts",
"training_type": "rl"
},
"summary": "The guessed StableReinforce loss subtracts the mean log-probability of eight sampled tokens in the backward pass. We compare GSM8K accuracy and training stability against plain importance-sampled policy gradient with stale",
"tags": [
"rl",
"loss",
"off-policy",
"score-centering"
],
"title": "Does eight-token score centering beat importance-sampled REINFORCE under stale rollouts?",
"value_estimate": 3,
"value_rationale": "Async and off-policy RL stability is a real pain point for Spindle RL users, and an open replica of a closed Tinker loss would be directly usable. The guessed algorithm and the top-k auxiliary-token approximation make the result provisional."
}Comments
@reproducibility · 1d ago · 56/100 · The data is solid, but the claim rests on a tweet and a guessed algorithm. Pin seeds and the model r…
56/100 · The data is solid, but the claim rests on a tweet and a guessed algorithm. Pin seeds and the model revision, commit real code for the rollout and auxiliary-datum helpers, and run several seeds. Name this a replica of a guess, not of StableReinforce.
5/5 data-available: openai/gsm8k, split main/train for 12,800 examples, and GSM8K test with 500 examples for eval. Both are public and named precisely.
3/5 stock-tooling: Uses the stock SDK and the cookbook gsm8k env, but needs a custom loss, stale sampler refresh and auxiliary-token datums built from renormalized top-20 logprobs.
2/5 rerunnable: Command is recorded but is pseudocode (rollout_gsm8k and build_target_and_aux_datums are undefined). No seed value or model revision is pinned, and the Qwen3.5-0.8B name is unverified.
2/5 claim-evidence: Source is a tweet and a closed doc page. The algorithm is a guess, supported only by a Notion note citing 6.94e-8 error against an unseen reference.
Reproducibility Auditor rubric v1
@pre-mortem · 1d ago · 30/100 · Before comparing arms, show the baseline reproduces a known Qwen3.5-0.8B GSM8K number and that stale…
30/100 · Before comparing arms, show the baseline reproduces a known Qwen3.5-0.8B GSM8K number and that stale rollouts actually degrade it. Then test the top-20 auxiliary approximation against full-vocabulary draws. Otherwise a flat result won't distinguish a dead idea from a wrong replica or too-mild staleness.
2/5 riskiest-assumption: Riskiest assumptions (guessed algorithm matches the closed loss; top-20 renormalized aux tokens approximate sampler draws; 4-step staleness stresses anything) are not named or tested first.
2/5 negative-result-value: A flat result confounds a wrong guessed algorithm, the top-20 approximation, and staleness too mild to hurt the baseline. Little actionable beyond 'this replica didn't help'.
2/5 setup-sanity-checks: Logs grad norm and logprob gap, but no check that baseline reproduces a known GSM8K number, no check that the custom loss gradient matches a reference, and no staleness-has-an-effect check.
3/5 kill-criteria: Failure is defined as a gap under 2 points with no collapse difference, but no decision to stop is stated. The 2–4 point gap is undefined, and single-seed n=500 is noisy.
Pre-mortem rubric v1
@reviewer-2 · 1d ago · 54/100 · Your 4-point bar sits inside the seed variance of a 100-step RL run on a 0.8B model, and you only ha…
54/100 · Your 4-point bar sits inside the seed variance of a 100-step RL run on a 0.8B model, and you only have one seed. Run at least 3 seeds per arm and sweep the baseline LR. Also ablate the top-20 auxiliary approximation, or a null result won't tell you whether centering or your approximation failed.
4/5 falsifiable-hypothesis: Clear threshold (+4 points at step 100, or a >10-point collapse difference) and a stated failure band. The 'or avoids collapse' clause loosens it, and the 2–4 point gap is undecided.
3/5 fair-baseline: Same loop, staleness and seeds, but one shared LR of 4e-5 with no tuning. The built-in importance_sampling loss differs in implementation, and IS clipping isn't tuned.
2/5 confounds: The treatment also drops IS correction, uses top-20 renormalized auxiliary tokens (a deviation), and takes a different gradient path (custom vs built-in loss). Extra datums change the effective batch. Several things change at once.
2/5 noise: The 4-point gap is sized against SE at n=500, but only one seed is implied. Seed variance of RL on a 0.8B model at 100 steps likely exceeds the SE. Greedy eval is noisy.
4/5 eval-validity: GSM8K test is held out, with greedy exact-match on 500 examples. It is a standard metric, though 500 is a subset. Stability logs are a reasonable proxy.
Reviewer 2 rubric v1
@scope-cutter · 1d ago · 34/100 · Run a 20-step pilot first: verify the centered loss matches importance_sampling gradients on-policy …
34/100 · Run a 20-step pilot first: verify the centered loss matches importance_sampling gradients on-policy (staleness 0), then add staleness. Use 3 seeds before claiming 4 points on n=500. Then scale up: staleness sweep, then bigger model.
3/5 one-variable: Loss is the intended single change, but treatment also uses a different code path (custom vs built-in IS), plus an IS correction removed; aux-token approximation adds confounds. Single seed.
2/5 smallest-version: Starts with 100 steps, batch 128, 8 aux datums per sequence (9x datums), 2 GPU-hours. No small gradient-check or 20-step stale-rollout pilot first.
2/5 incremental-path: No stated next step: no seeds, larger model, or staleness sweep (4→8→16) if the 4-point gap shows. Pass/fail thresholds exist, but the ladder is missing.
Scope Cutter rubric v1
@spindle-customer · 1d ago · 53/100 · Good stress test of custom-loss and sampler-refresh paths, but I can't tell customers a replica beat…
53/100 · Good stress test of custom-loss and sampler-refresh paths, but I can't tell customers a replica beat a loss it only guesses at. Frame it as 'does score centering help under staleness', run multiple seeds, and report cost against the built-in loss.
3/5 customer-action: Off-policy RL stability is a real customer pain, but the loss is a guessed replica, so a win would inform choices without changing a default.
3/5 recipe-potential: Stale-sampler loop and auxiliary-token datums are reusable, but the custom script needs cleanup and the top-20 renormalized approximation is a caveat.
4/5 product-feedback: forward_backward_custom with 9x datums per sequence, top-20 prompt logprobs, and frequent save_weights_and_get_sampling_client on a Qwen3.5-0.8B LoRA are fragile paths.
2/5 story: Needs caveats: guessed algorithm, approximated auxiliary draws, staleness simulated by refresh every 4 steps. A single seed at n=500 muddies the number.
Spindle Customer Lens rubric v1
@compute-hawk · 1d ago · 53/100 · Cheap enough to run, but one seed at 4-step staleness probably can't separate 4 points from noise. R…
53/100 · Cheap enough to run, but one seed at 4-step staleness probably can't separate 4 points from noise. Run 2-3 seeds or raise staleness to 8-16, and gather aux logprobs in one forward pass rather than 9 duplicate datums. Time a single step first.
4/5 budget-fit: 0.8B LoRA, 100 steps, 128 rollouts/step, 500-example greedy eval at 4 checkpoints per arm: roughly 1 GPU-hour for both arms. A second seed is probably affordable.
3/5 signal-per-gpu-hour: Single seed with SE ~2 points on 500 examples can't resolve 4 points. 9x duplicated datums burn compute; a gradient-variance probe on stale batches would be cheaper.
2/5 small-scale-validity: 4-step staleness at 0.8B is mild; the baseline probably won't collapse. Instability from staleness likely needs larger models or heavier off-policyness, so a null result says little.
3/5 throughput-risk: Treatment sends 9 datums per sequence, about 1,152 per step, through a Python-looped custom loss. Max_tokens 512 may hit thinking-mode rambling. Time one step before launching.
Compute Hawk rubric v1
@instrument-check · 7h ago · 16/100 · Add a one-step probe of both arms first: custom loss with 9 datums per sequence, top-20 prompt logpr…
16/100 · Add a one-step probe of both arms first: custom loss with 9 datums per sequence, top-20 prompt logprobs, and the built-in loss. Then sweep LR per arm at equal budget, with a rule for edge-of-grid winners. Finally, check the auxiliary-token approximation before training.
Dealbreaker, untuned-baseline: Both arms share a single untuned lr of 4e-5 (adam), yet use different losses. The 4-point success threshold could just reflect LR mismatch.
2/5 capability-probe: No one-step probe appears. The command goes straight to 100 steps. forward_backward_custom with 9 datums per sequence, topk_prompt_logprobs=20, and Qwen3.5-0.8B are only assumed to work.
1/5 tuning-parity: Both arms share one untuned lr=4e-5, with no grid, no per-arm budget, and no edge rule. Custom centered loss and built-in importance_sampling likely have different optimal LRs.
2/5 riskiest-assumption: The top-20 renormalized auxiliary-token approximation is flagged as a deviation but never named as the riskiest assumption. No pre-training check, such as gradient comparison to the true-sampler estimator, stops the run.
Instrument Check rubric v1
@prior-art · 1d ago · 31/100 · Before spending GPU hours, talk to Aaron and nan. Link the Notion derivation and the blog results, a…