🔬Modal Research Exchange

i'm curious about this recent post on group-size-1 RL: <https://www.trajectory.ai/field-notes/you-only-rollout-once|trajectory.ai/field-notes/you-only-rollout-once>

• can we reproduce their results on any public benchmarks, including how GS1 compares to FlashREINFORCE and GRPO?
• how does this compare (in terms of final task performance as well as training time) to just training an actual critic?
• they don't release their exact choice of kernel function, but claim that you can easily hillclimb to find a recipe yourself. how easy is it to actually find one that is efficient and performs reasonably well compared to GRPO as implied in their results? does this choice of kernel transfer across a diverse set of benchmark tasks, or do you end up having to find a new one per benchmark? (as a starting point they do provide a list of example kernels in their post, some of which have a close enough effective batch size to GS1)
    ◦ i'd like to see how high of an effective batch size we can hillclimb to offline using the technique they mentioned. _this part can be done offline without using spindle_

[Kevin Li] <@U0CE8HJ4BGX> wya

[Kevin Li] :party_try_not_to_cry:

{
  "adjacent_experiment": "",
  "reproducibility_reason": "A GS1 kernel-baseline advantage is client-side code in the cookbook RL loop with the stock importance-sampling loss and AdamW, so it can be compared with GRPO on GSM8K on the shared Spindle.",
  "reproducible_on_spindle": true,
  "sources": [
    {
      "brief": null,
      "error": "digest failed: ask_json failed after 2 attempts: Unterminated string starting at: line 1 column 585 (char 584)",
      "full_text": true,
      "kind": "page",
      "title": "You Only Rollout Once \u2014 Trajectory",
      "url": "https://www.trajectory.ai/field-notes/you-only-rollout-once"
    }
  ],
  "spec": {
    "activation_gate": {
      "applicable": true,
      "checks": [
        {
          "measures": "rollouts per distinct prompt in logged treatment batches",
          "metric": "executed_group_size_treatment",
          "one_of": [
            "1"
          ]
        },
        {
          "measures": "rollouts per distinct prompt in logged baseline batches",
          "metric": "executed_group_size_baseline",
          "one_of": [
            "8"
          ]
        },
        {
          "max": 0.9,
          "measures": "mean training reward over the first 2 treatment steps",
          "metric": "pilot_mean_reward",
          "min": 0.1
        },
        {
          "measures": "10th percentile over rollouts of (sum w)^2/sum w^2 for the kernel weights",
          "metric": "kernel_effective_neighbors_p10",
          "min": 8.0
        },
        {
          "measures": "fraction of treatment rollouts with |advantage|>1e-6",
          "metric": "frac_nonzero_advantage_treatment",
          "min": 0.5
        },
        {
          "measures": "Pearson correlation of b_i with a held-out 8-sample per-prompt mean reward, over 64 pilot prompts",
          "metric": "baseline_reward_corr",
          "min": 0.1
        }
      ],
      "reason": "A GS1 baseline only differs from a constant baseline if the kernel draws on several neighbors with varied rewards, and the prompts give a non-degenerate reward. The executed group size must actually be 1 for the treatment and 8 for the baseline."
    },
    "arms": [
      {
        "advantage": "group_mean",
        "group_size": 8,
        "groups_per_batch": 8,
        "lr": 4e-05,
        "name": "grpo_gs8",
        "seed": 0
      },
      {
        "advantage": "kernel_cross_prompt_baseline",
        "group_size": 1,
        "groups_per_batch": 64,
        "lr": 4e-05,
        "name": "gs1_kernel",
        "seed": 0
      }
    ],
    "base_model": "Qwen/Qwen3.5-2B",
    "baseline": "GRPO via the same client loop with group_size=8, groups_per_batch=8 and group-mean advantages. It uses the same model, LoRA rank, LR, steps, rollouts per step, max_tokens, renderer, reward, seed and GSM8K eval.",
    "batch_size": 64,
    "budget": {
      "gpu_hours": 2.0,
      "max_wall_minutes": 120
    },
    "command": "Custom client loop (python rx_gs1.py --arm {grpo_gs8|gs1_kernel}); the only custom part is the compute_advantages function.\n\nSetup:\n  import os, tinker\n  sc = tinker.ServiceClient(base_url=os.environ['TINKER_BASE_URL'])\n  tc = sc.create_lora_training_client(base_model='Qwen/Qwen3.5-2B', rank=32, seed=0)\n  Use the tinker_cookbook math_rl gsm8k env builder for prompts, renderer and reward.\n\nEach of 40 steps:\n  1. sampler = tc.save_weights_and_get_sampling_client(...).\n  2. Draw prompts from a fixed seed-0 shuffled order: GRPO takes 8 prompts x 8 samples, GS1 takes 64 prompts x 1 sample. max_tokens=512, temperature=1.0.\n  3. Grade with the env reward.\n  4. Compute advantages: GRPO uses r minus group mean. GS1 uses r_i minus a Gaussian-kernel-weighted mean over a 256-rollout window excluding i, with phi = (log prompt length, sampler mean prompt-token logprob) and bandwidth 0.5.\n  5. Build Datums with sampling logprobs and advantages.\n  6. tc.forward_backward(datums, loss_fn='importance_sampling').\n  7. tc.optim_step(tinker.AdamParams(learning_rate=4e-5)).\n\nAfter training: run save_weights_for_sampler, then greedy-decode all 1319 GSM8K test items for each arm. Write per-item correctness to /tmp/rx-gs1/{arm}/eval.jsonl, and score the paired Hoeffding rule from those files.\n\nActivation-failure plan:\n- (a) Offline first: unit-test compute_advantages in float64 on synthetic rewards. Check forced activation (heterogeneous phi gives nonconstant b_i) and forced nonactivation (huge bandwidth gives b_i equal to the window mean, i.e. a batch-mean REINFORCE baseline). Run gradient checks of the new loop with the treatment disabled against stock cookbook GRPO advantages on one fixed batch.\n- (b) A 2-step treatment pilot logs the gate metrics.\n- (c) If kernel_effective_neighbors fails, the single pre-registered calibration is bandwidth in {0.5, 1.0}, chosen by the gate on pilot training prompts only, never on GSM8K test. The cap is 2 pilot steps and 10 wall minutes.\n- (d) Keep the original failed pilot logs. Any other setting change requires a recorded revised spec before execution. Gates are never relaxed.\n\nNo full run happens while correctness_preflight is unresolved.",
    "correctness_preflight": {
      "paths": [
        {
          "baseline_revision": "UNRESOLVED: pin the tinker-cookbook commit SHA whose rl loop / math_rl group-mean advantage is the stock baseline",
          "id": "gs1_kernel_advantage_loop",
          "kind": "loss",
          "reference_revision": "UNRESOLVED: pin a commit of an independently written float64 numpy reference of the kernel baseline (or the Trajectory reference if released)",
          "tolerances": {
            "baseline_gradient": {
              "atol": 1e-05,
              "min_cosine": 0.99999,
              "rtol": 0.0001
            },
            "baseline_loss": {
              "atol": 1e-06,
              "rtol": 1e-05
            },
            "reference_gradient": {
              "atol": 1e-05,
              "min_cosine": 0.99999,
              "rtol": 0.0001
            },
            "reference_loss": {
              "atol": 1e-06,
              "rtol": 1e-05
            }
          }
        }
      ],
      "required": true,
      "status": "unresolved: the source was not fetched and its kernel is unreleased, so there is no trusted treatment reference; pin revisions and an independent reference before any full sweep"
    },
    "dataset": {
      "hf_path": "openai/gsm8k",
      "n_examples": 2560,
      "name": "gsm8k",
      "split": "main/train"
    },
    "decision_rule": {
      "alpha": 0.05,
      "criteria": [
        {
          "bounds": [
            0,
            1
          ],
          "id": "primary",
          "margin": -0.03,
          "metric": "accuracy",
          "weights": {
            "baseline": -1,
            "treatment": 1
          }
        }
      ],
      "estimand": "Paired difference in GSM8K test exact-match accuracy (GS1 kernel minus GRPO) for these fixed seed-0 final checkpoints; larger is better",
      "method": "paired_hoeffding",
      "seeds": [
        0
      ],
      "unit": "fixed_checkpoint_items",
      "version": 1
    },
    "eval": {
      "aggregation": null,
      "benchmark": "GSM8K test (openai/gsm8k main/test)",
      "metric": "exact-match accuracy, greedy, final checkpoint",
      "n_examples": 1319
    },
    "experiment_features": {
      "adaptive_selection": false,
      "checkpoint_restore": false,
      "custom_loss": true,
      "custom_optimizer": false,
      "generated_training_data": false,
      "heldout_loss": false,
      "learned_generator": false,
      "paired_comparison": true,
      "paired_training_arms": true,
      "warmup_adaptation": false
    },
    "feasibility": {
      "config": null,
      "effective_config": {
        "arms": [
          "grpo_gs8",
          "gs1_kernel"
        ],
        "base_model": "Qwen/Qwen3.5-2B",
        "batch_size": 64,
        "max_seq_len": 1024,
        "max_tokens": 512,
        "n_eval": 1319,
        "steps": 40
      },
      "generation": null,
      "measurements": null,
      "primary_endpoint": "gsm8k_test_exact_match_accuracy_final_checkpoint",
      "recovery": null,
      "source": null
    },
    "feasibility_check": {
      "generation_applicability": {
        "check": "generation_audit",
        "enforce": false,
        "features": {
          "generated_training_data": false
        },
        "reasons": [
          "spec.generation_audit opts out and no feature it covers is true"
        ],
        "status": "not_applicable"
      },
      "generation_counts": null,
      "gpu_hours": null,
      "known_gpu_hours": 0,
      "phases": {},
      "problems": [
        "Effective configuration not recorded",
        "Pilot evidence source not recorded",
        "Representative training timing not recorded",
        "Representative setup timing not recorded",
        "Representative evaluation timing not recorded",
        "Positive recovery reserve not recorded"
      ],
      "status": "unknown",
      "wall_seconds": null
    },
    "feasibility_required": true,
    "generation_audit": {
      "answer_comparison": "numeric_exact",
      "max_eval_overlaps": 0,
      "max_label_errors": 0,
      "min_sample_size": 100,
      "required": false
    },
    "hypothesis": "At 64 rollouts per step for 40 steps on Qwen3.5-2B, GS1 RL with a kernel-weighted cross-prompt baseline is noninferior to GRPO (group_size=8) on GSM8K test exact-match accuracy. The noninferiority margin is 3 points.",
    "lora_rank": 32,
    "max_seq_len": 1024,
    "method": "Both arms train a rank-32 LoRA on Qwen/Qwen3.5-2B from the same base. They use one client-side RL loop built on the tinker-cookbook rl loop, with the stock importance_sampling loss and AdamW. The prompts come from GSM8K train with the stock math_rl renderer and the verifiable exact-answer reward. Each step uses 64 rollouts, max_tokens=512 and temperature 1.0.\n\nThe arms differ only in the advantage estimator:\n- GRPO baseline: 8 prompts x 8 rollouts. Advantage = reward minus group mean, matching the cookbook's default.\n- GS1 treatment: 64 distinct prompts x 1 rollout. Advantage = r_i minus b_i. Here b_i is a kernel-weighted mean of rewards from other prompts in a sliding window of the last 256 rollouts. The window covers the current batch excluding i, plus the previous steps.\n\nThe kernel is fixed before running: a Gaussian kernel exp(-||phi_i - phi_j||^2 / (2*0.5^2)) on a 2-d standardized prompt feature phi = (log prompt-token length, sampling-policy mean token logprob of the prompt). The prompt logprob is computed by the sampling client before rollout. If kernel mass is below 1e-3, the estimator falls back to the window mean.\n\nThe Trajectory post could not be fetched and does not release its kernel, so this kernel is our own pre-registered choice and not the paper's exact method. Kernel hillclimbing, FlashREINFORCE, the trained-critic comparison and cross-task transfer are out of scope for this first node. Each is a follow-up spec.\n\nThe shared Spindle is enough; no fork is needed. Evaluation is greedy decoding on the full GSM8K test split (1319 items), scored by exact match, for the final checkpoint of each arm.",
    "optimizer": {
      "lr": 4e-05,
      "name": "adam"
    },
    "preflight": {
      "probe": "For each arm, create the training client with adam, rank=32 and seed=0. Run one forward_backward with loss_fn='importance_sampling' on one sampled GSM8K example carrying that arm's computed advantage, then one optim_step(tinker.AdamParams(learning_rate=4e-5)). Confirm both succeed and the loss is finite. Then run the offline advantage unit tests and the 2-step activation pilot.",
      "riskiest_assumption": "The custom GS1 advantage path is correct and engages. Specifically, the new loop reproduces stock GRPO advantages when the treatment is disabled, and the kernel baseline gives a nonconstant, reward-predictive b_i. Otherwise a GS1-vs-GRPO gap reflects bugs or a degenerate batch-mean baseline, not the method.",
      "tuning": "Both arms use the same optimizer (AdamW) and a fixed LR of 4e-5 taken from cookbook math_rl defaults for small Qwen models. There is no per-arm grid, because the hypothesis concerns the advantage estimator at a fixed recipe. An LR-sensitivity follow-up would give each arm a 3-value grid selected on a GSM8K train holdout, with the grid-edge rule applied."
    },
    "primary_endpoint": "gsm8k_test_exact_match_accuracy_final_checkpoint",
    "restoration_preflight": {
      "applicable": true,
      "fixed_batch_required": true,
      "groups": [
        {
          "arms": [
            {
              "name": "grpo_gs8",
              "optimizer": "adam",
              "optimizer_config": {
                "beta1": 0.9,
                "beta2": 0.95,
                "eps": 1e-08,
                "lora_rank": 32,
                "lr": 4e-05,
                "seed": 0
              }
            },
            {
              "name": "gs1_kernel",
              "optimizer": "adam",
              "optimizer_config": {
                "beta1": 0.9,
                "beta2": 0.95,
                "eps": 1e-08,
                "lora_rank": 32,
                "lr": 4e-05,
                "seed": 0
              }
            }
          ]
        }
      ],
      "reason": "Both trained arms start from the same base LoRA initialization within seed 0."
    },
    "spec_preflight": {
      "findings": [],
      "status": "passed"
    },
    "steps": 40,
    "success_criterion": "GS1 minus GRPO GSM8K test exact-match accuracy, as a lower 95% Hoeffding bound over 1319 paired items, must exceed -0.03 (noninferiority margin of 3 points). If the interval straddles -0.03, the pilot is inconclusive and is not counted as a win.",
    "title": "GS1 kernel-baseline RL vs GRPO at an equal rollout budget on GSM8K",
    "training_type": "rl"
  },
  "summary": "Group-size-1 RL takes one rollout per prompt and builds its baseline from other prompts' rewards; we compare GSM8K accuracy with GRPO at the same rollout count.",
  "tags": [
    "rl",
    "paper",
    "grpo",
    "lora",
    "advantage-estimation"
  ],
  "title": "Does group-size-1 RL with a kernel baseline match GRPO on GSM8K?",
  "value_estimate": 4,
  "value_rationale": "If one rollout per prompt matches GRPO, Spindle RL users could cut sampling cost roughly by the group size, which they would act on. Triage only saw the link because the Trajectory post could not be fetched, so the kernel here is our own pre-registered stand-in."
}