This repository was archived by the owner on Jul 17, 2026. It is now read-only.
feat(recon): fixed-k and scheduled-probability subset routing (S11, additive)#977
Open
claude-spd1 wants to merge 1 commit into
Open
feat(recon): fixed-k and scheduled-probability subset routing (S11, additive)#977claude-spd1 wants to merge 1 commit into
claude-spd1 wants to merge 1 commit into
Conversation
…ended, additive)
Two new routing configs for the subset recon terms:
- fixed_k_subset: per position, a uniform subset of exactly k live sites
(vs uniform_k's k ~ U{1..n}).
- scheduled_probability: static_probability with p evaluated per step from
a ScheduleConfig (increasing ramps use final_val_frac > 1; both endpoints
validated as probabilities).
SAMPLE_ROUTING gains the traced step: RoutingSampler is now
(key, leading_shape, step_f32) -> draws; only scheduled samplers read the
step (all others ignore it), and build_loss_terms/make_plan thread
total_steps down so the scheduled sampler can evaluate its schedule with
scheduled_value_traced. SPEC S11 + the SAMPLE_ROUTING binding-seam row
amended additively; R1 independence unchanged (draws stay independent
across sites, positions, forwards, steps — the schedule only moves p).
Motivation (board task run-some-random-architecture-ideas): test whether
the subset recon term wants more base-model supervision early — low
routing p ramped up over training — vs the static uniform-k baseline.
Crew-Address: task/run-some-random-architecture-ideas
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Two new subset-routing configs for the recon terms, plus the plumbing they need:
fixed_k_subset(k: PositiveInt) — per position, a uniformly-drawn subset of exactlyklive sites (vsuniform_k's per-positionk ~ U{1..n}).scheduled_probability(p: ScheduleConfig) —static_probabilitywithpevaluated at the traced step viascheduled_value_traced. Increasing ramps usefinal_val_frac > 1; a validator pins both schedule endpoints to [0, 1].Plumbing:
RoutingSampleris now(key, leading_shape, step_f32) -> draws— samplers already run inside the jitted step, so the traced step just threads in; only scheduled samplers read it.build_loss_terms/make_plan/subset_chunk_plangaintotal_stepsso the scheduled sampler can anchor its schedule (all call sites updated, incl. tests).SPEC: S11 and the SAMPLE_ROUTING binding-seam row amended additively (marked AMENDED 2026-07-12) — @oli-clive-griffin please eyeball per the one-rule. R1 independence is unchanged: draws stay independent across sites/positions/forwards/steps; the schedule only moves
p.Motivation and Context
Board task
run-some-random-architecture-ideas: Oli's scheduled-subset-recon idea (early training wants more base-model supervision → low routing p ramped up) plus his ask for more static k values. Wave 1 (static p dose-response + uniform_k baseline) is already running as wandb grouproutep-sched-1; this PR unblocks wave 2 (3 schedule arms + fixed-k {1,3,6,12,24}).Related Issue
Board task
run-some-random-architecture-ideas.How Has This Been Tested?
param_decomp/tests/test_routing.py: fixed-k routes exactly k per position (k ∈ {1,3,8}), k > n refused, uniform_k span sanity, static-p empirical mean ≈ p, scheduled-p empirical mean tracksget_scheduled_valueat steps {0, 50, 99} (incl. ramp reaching 1.0), endpoint validators refuse p > 1. Green at default device count andXLA_FLAGS=--xla_force_host_platform_device_count=4.make typeclean; equivalence goldens untouched (no numeric path changed — existing samplers gained an ignored arg).Does this PR introduce a breaking change?
No behavior change for existing configs (new arg is ignored by existing samplers; RNG key derivation unchanged, so existing-arm trajectories are bit-identical). API-level:
build_loss_terms/make_plan/subset_chunk_plan/RoutingSamplersignatures gainedtotal_steps/step_f32— all in-repo callers updated.Crew-Address: task/run-some-random-architecture-ideas
🤖 Generated with Claude Code