feat(recon): per-sample assignment for MergedStochasticPGDReconLoss (+ 4L 400k arm configs)

param_decomp/SPEC.md

@@ -155,7 +155,7 @@ def imp_min_terms(ci_upper, pnorm, reference_token_count):   # per-site grouping
 def stochastic_recon_loss(components, ci_lower, residual, clean_output):
     total, n_forwards = 0, 0
     for (live_sites, SAMPLE_ROUTING) in RECON_PLAN:
-        for routes in SAMPLE_ROUTING(key, [B,T]):        # fresh per step                 (R1,S11)
+        for routes in SAMPLE_ROUTING(key, [B,T], step):  # fresh per step                 (R1,S11)
             masks, delta_masks = make_masks(ci_lower_s, source_s ~ U[0,1]^[B,T,C+1])  ∀ s ∈ live_sites
             total += kl_per_position(masked_forward(residual, live_sites, masks, delta_masks, routes),
                                      clean_output)
@@ -280,7 +280,7 @@ faithfulness, sources/PPGD, the squashings (S5/S6), the imp-min reduction, the g
 | S8' | Imp-min contributes `imp_coeff·lp + freq_coeff·freq` with INDEPENDENT coefficients; the frequency normalizer `a' = reference_token_count` is explicit (batch-invariant at fixed firing rate), not the implicit `B·T`. `freq` is absent when no `frequency` is configured. `a' = B·T` recovers the old rolled `imp_coeff·Σ_c f·(1 + beta·log2(1 + B·T·f))` with `freq_coeff = imp_coeff·beta`. |
 | S9 | `pnorm(step)` follows its `ScheduleConfig` (canonical: linear `2.0 → 0.4` over the run, `final_val_frac=0.2`), evaluated by `schedule.scheduled_value_traced`; `eps` sits inside the power. The smooth-L0 `gamma(step)` (S9′) follows the same rule. Constant-`p` is `fn_type=constant`; warmup on `p`/`gamma` is refused (`build_loss_terms`). **AMENDED 2026-07-02** (#915): the former windowed linear anneal (`{p,gamma}_anneal_{start,end}_frac`) is retired — every config on record used the trivial `[0, 1]` window — and the anneal now shares the schedule's `decay_steps − 1` denominator (S20), reaching the final value AT `step = steps − 1` instead of one step past the run (a ≤O(1/steps) per-step trajectory change vs the old `step / steps`). |
 | S10′ | The recon objective is a static tuple of coefficiented loss TERMS (one per configured recon loss metric, in config order). Each term is a static plan of `(live_sites, SAMPLE_ROUTING, MASK_SOURCE)` entries; the term's loss = mean over ALL its forwards (every draw of every entry) of `kl_per_position`; the total adds `coeff · term` per term. Plan structures (live-sets, sampler identities, family sizes, strategy kinds) are fixed across steps. The §4 pseudocode shows the production two-term instantiation (`stochastic_recon_loss` + `adversarial_recon_loss`). Recon KL direction is pinned by S25; the mean-over-forwards ≡ accumulator identity by S26. |
-| S11 | `uniform_k_routing`, per position: `k ~ U{1..|live_sites|}` then a uniform `k`-subset of the live sites routes True; non-live sites are not live at all. Routing draws are fresh per step, sampled inside the step. |
+| S11 | `uniform_k_routing`, per position: `k ~ U{1..|live_sites|}` then a uniform `k`-subset of the live sites routes True; non-live sites are not live at all. Routing draws are fresh per step, sampled inside the step. **AMENDED 2026-07-12** (additive): three more routing samplers share the rule — `fixed_k_routing` (per position, a uniform subset of exactly `k` live sites, `1 ≤ k ≤ |live_sites|`), `static_probability_routing` (per position × live site, independent `Bernoulli(p)` — the torch `StaticProbabilityRouter`), and `scheduled_probability_routing` (`Bernoulli(p(step))`, `p` a `ScheduleConfig` evaluated by `scheduled_value_traced`; both schedule endpoints must be probabilities, an increasing ramp uses `final_val_frac > 1`). SAMPLE_ROUTING therefore takes the traced step: `(key, [B,T], step_f32)`; only scheduled samplers read it. |
 | S12′ | An adversarial term's loss forward consumes its sources as LEAVES (no ascent-graph history); gradient flows to components and (through `ci_lower`) to the CI fn — and, for persistent sources, to the leaves themselves (S14′). The PRODUCTION adversarial term masks ALL sites and routes everywhere; subset-routed adversarial terms route per their plan. |
 | S13′ | Per persistent term: source updates per training step = `n_warmup + 1`, all through THAT term's persistent SRC_STEP optimizer state; its source LR schedule advances once per training step. The source LR is `scheduled_value_traced` over the term's `ScheduleConfig` (constant-after-warmup only — `build_loss_terms` and the lab conversion both refuse decay). **AMENDED 2026-07-02** (#915): the former `warmup_pct==0` accepted seam (JAX clamped `warmup_steps = max(floor(...), 1)` → source LR `=0` at step 0) is retired — at `warmup_pct==0` JAX now matches torch's full LR at step 0. Production uses 2.5% warmup and is unaffected either way. |
 | S14′ | Each persistent term's final ascent gradient comes from the SAME graph as the main backward (pre-update components, live `ci_lower`), unscaled by THAT term's coeff. It is applied after backward; it must not use post-update params. |
@@ -310,7 +310,7 @@ faithfulness, sources/PPGD, the squashings (S5/S6), the imp-min reduction, the g
 |---|---|---|
 | `RECON_TERMS` | any static tuple of coefficiented terms, each a static plan of `(live_sites, SAMPLE_ROUTING, MASK_SOURCE)` entries: `subset_chunk_plan` · `per_site_plan` (the torch "layerwise" shape) · `all_sites_plan` · custom subset families (pairs, covers, …); built from the shared torch loss configs by `build_loss_terms` | ★ stochastic subset term + persistent-PGD term |
 | `MASK_SOURCE` | `stochastic` (fresh U[0,1]/Bernoulli per draw) · `constant(v)` (`v=0` CI-masked, `v=1` unmasked; no delta path) · `fresh_pgd(init, n_steps, step_size, scope)` · `persistent(state_key)` | ★ stochastic + persistent |
-| `SAMPLE_ROUTING` | `(key, [B,T]) → tuple of routing draws`, statically sized; draws may be jointly sampled (independent repeats, antithetic/complementary subsets, per-step covers) | ★ uniform_k_routing(·, 1) |
+| `SAMPLE_ROUTING` | `(key, [B,T], step_f32) → tuple of routing draws`, statically sized; draws may be jointly sampled (independent repeats, antithetic/complementary subsets, per-step covers); `uniform_k` · `fixed_k` · `static_probability` · `scheduled_probability(p(step))` · `all` (S11) | ★ uniform_k_routing(·, 1) |
 | `SRC_STEP` | `adam(β₁,β₂,ε)` with bias correction; `sign` (`sources += lr·sign(grad)`) | ★ adam(.5, .99, 1e-8) |
 | `PROJ` / `EFFECTIVE` | **clamp**: PROJ = clamp[0,1], EFFECTIVE = identity, init U[0,1] · **sigmoid**: PROJ = identity (unbounded raw), EFFECTIVE = sigmoid, init N(0,1) | ★ clamp |
 | `SCOPE` | persistent: `c (1,1)` · `sc (1,T)` · `nsc (n,T), n|B` · `bsc (B,T)` (jax: `sc` + `bsc` today) — fresh-PGD: `c` · `bc` · `bsc` | ★ sc |

param_decomp/configs.py

@@ -46,14 +46,46 @@ class StaticProbabilityRoutingConfig(BaseConfig):
     p: Probability
 
 
+class FixedKSubsetRoutingConfig(BaseConfig):
+    """Route each position to a uniformly-drawn subset of exactly `k` of the live sites."""
+
+    type: Literal["fixed_k_subset"] = "fixed_k_subset"
+    k: PositiveInt
+
+
+class ScheduledProbabilityRoutingConfig(BaseConfig):
+    """`static_probability` with `p` evaluated per step from a schedule. The schedule ends
+    at `start_val * final_val_frac`; an increasing ramp uses `final_val_frac > 1`. Both
+    endpoints must be valid probabilities."""
+
+    type: Literal["scheduled_probability"] = "scheduled_probability"
+    p: ScheduleConfig
+
+    @model_validator(mode="after")
+    def validate_probability_endpoints(self) -> Self:
+        end_val = self.p.start_val * self.p.final_val_frac
+        if not (self.p.start_val <= 1.0 and end_val <= 1.0):
+            raise ValueError(
+                f"p schedule endpoints must be probabilities: "
+                f"start={self.p.start_val}, end={end_val}"
+            )
+        return self
+
+
 class AllRoutingConfig(BaseConfig):
     """Route every position to every module (the `"all"` fast path)."""
 
     type: Literal["all"] = "all"
 
 
 # Discriminated union over the subset-routing configs (keyed by ``type``).
-SubsetRoutingType = UniformKSubsetRoutingConfig | StaticProbabilityRoutingConfig | AllRoutingConfig
+SubsetRoutingType = (
+    UniformKSubsetRoutingConfig
+    | StaticProbabilityRoutingConfig
+    | FixedKSubsetRoutingConfig
+    | ScheduledProbabilityRoutingConfig
+    | AllRoutingConfig
+)
 
 
 # ---------------------------------------------------------------------------
@@ -428,14 +460,21 @@ class MergedStochasticPGDReconLossConfig(LossMetricConfig):
     canonical 0.5/0.5 pair). Carries the persistent-adversary fields; one source bundle
     feeds this one term (SPEC S23) and the S14' final ascent flows through its backward.
     A mixed assignment is a legal point of the mask box, so this samples the SAME
-    feasible set as the two-term objective, under a different (joint) sampler."""
+    feasible set as the two-term objective, under a different (joint) sampler.
+
+    `assignment` picks the Bernoulli granularity: `per_position` draws one bit per
+    (batch element, position) — a sequence's masked forward sees both families, coupled
+    through attention; `per_sample` draws one bit per batch element, so every position
+    in a sequence takes the same family (no within-sample mixing — each sample's loss is
+    scored against a pure-family context)."""
 
     type: Literal["MergedStochasticPGDReconLoss"] = "MergedStochasticPGDReconLoss"
     optimizer: AdamPGDConfig
     scope: PersistentPGDSourceScope
     source_dtype: Literal["float32", "bfloat16"] = "float32"
     n_warmup_steps: NonNegativeInt = 0
     adv_fraction: float = Field(gt=0.0, lt=1.0)
+    assignment: Literal["per_position", "per_sample"] = "per_position"
     routing: SubsetRoutingType = Field(default_factory=UniformKSubsetRoutingConfig)
 
 

param_decomp/configs/pile4l_chunks_ab_1chunk_fk12_400k_cweast.yaml

@@ -0,0 +1,182 @@
+# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1
+# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp
+# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper
+# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64).
+# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k
+# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0
+# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the
+# going-forward penalty) - so only the CI-fn capacity changes vs round 1.
+# data_files points at the crew user's staging of the same HF dataset (the canonical
+# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results.
+# CHUNK SWEEP rescue arm: single all-sites pass, fixed_k routing k=12 — an absolute
+# routed-site budget per forward (deep-linear-identity finding: locality is a count,
+# not a fraction; uniform-k routes ~n/2 sites).
+run_name: pile4l-chunks-ab-1chunk-fk12-400k
+# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of
+# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot
+# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the
+# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile
+# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15.
+# Target: LlamaSimpleMLP 4L pile model
+# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope
+# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/
+# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml
+# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream:
+#   1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged
+#      parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 —
+#      the JAX trainer reads pre-tokenized parquet shards only).
+#   2. data.eval_split: val -> train (parquet loading exposes a single "train"
+#      split; the staged _val dir exists for ad-hoc use but the loaders don't
+#      address it).
+#   3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves
+#      nothing; the JAX leg checkpoints for requeue-resume + offline eval).
+# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is
+# upstream-identical.
+cadence:
+  keep_last_n_checkpoints: 2
+  save_every: 5000
+  train_log_every: 200
+data:
+  buffer_size: 1000
+  column_name: input_ids
+  data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.parquet
+  dataset_name: parquet
+  eval_split: train
+  is_tokenized: true
+  max_seq_len: 512
+  revision: null
+  shuffle_each_epoch: true
+  streaming: false
+  tokenizer_name: EleutherAI/gpt-neox-20b
+  train_split: train
+eval:
+  batch_size: 128
+  every: 1000
+  metrics:
+  - n_batches_accum: 1
+    type: CIHistograms
+  - ci_alive_threshold: 0.0
+    type: ComponentActivationDensity
+  - ci_alive_threshold: 0.0
+    groups:
+      layer_0:
+      - h.0.*
+      layer_1:
+      - h.1.*
+      layer_2:
+      - h.2.*
+      layer_3:
+      - h.3.*
+      total:
+      - '*'
+    type: CI_L0
+  - rounding_threshold: 0.0
+    type: CEandKLLosses
+  - type: CIMeanPerComponent
+  - coeff: null
+    type: StochasticHiddenActsReconLoss
+  - type: CIHiddenActsReconLoss
+  - coeff: null
+    init: random
+    mask_scope: c
+    n_steps: 20
+    name: PGDReconLoss_20step
+    step_size: 0.1
+    type: PGDReconLoss
+  n_steps: 1
+  slow_every: 10000
+  slow_on_first_step: true
+pd:
+  batch_size: 64
+  ci_config:
+    type: chunkwise_transformer
+    blocks_per_chunk: 1
+    d_model: 2048
+    n_blocks: 8
+    n_heads: 16
+    mlp_hidden: 8192
+  ci_fn_optimizer:
+    betas:
+    - 0.9
+    - 0.999
+    grad_clip_norm: null
+    lr_schedule:
+      final_val_frac: 0.1
+      fn_type: cosine
+      start_val: 5.0e-05
+      warmup_pct: 0.0
+    weight_decay: 0.0
+  components_optimizer:
+    betas:
+    - 0.9
+    - 0.999
+    grad_clip_norm: 0.01
+    lr_schedule:
+      final_val_frac: 0.1
+      fn_type: cosine
+      start_val: 5.0e-05
+      warmup_pct: 0.0
+    weight_decay: 0.0
+  decomposition_targets:
+  - C: 3072
+    module_pattern: h.*.mlp.c_fc
+  - C: 3584
+    module_pattern: h.*.mlp.down_proj
+  - C: 512
+    module_pattern: h.*.attn.q_proj
+  - C: 512
+    module_pattern: h.*.attn.k_proj
+  - C: 1024
+    module_pattern: h.*.attn.v_proj
+  - C: 1024
+    module_pattern: h.*.attn.o_proj
+  faithfulness_warmup_lr: 0.001
+  faithfulness_warmup_steps: 400
+  faithfulness_warmup_weight_decay: 0.0
+  loss_metrics:
+  - frequency:
+      coeff: 0.0001
+      reference_token_count: 65536
+    coeff: 0.0002
+    gamma:
+      start_val: 1.0
+      fn_type: linear
+      final_val_frac: 0.01
+    type: SmoothL0ImportanceMinimalityLoss
+  - coeff: 0.5
+    routing:
+      type: fixed_k_subset
+      k: 12
+    type: StochasticReconSubsetLoss
+  - coeff: 0.5
+    n_warmup_steps: 2
+    optimizer:
+      beta1: 0.5
+      beta2: 0.99
+      eps: 1.0e-08
+      lr_schedule:
+        final_val_frac: 1.0
+        fn_type: constant
+        start_val: 0.01
+        warmup_pct: 0.025
+      type: adam
+    scope:
+      type: bsc
+    type: PersistentPGDReconLoss
+  - coeff: 10000000.0
+    type: FaithfulnessLoss
+  seed: 0
+  steps: 400000
+runtime:
+  remat_recon_forwards: false
+  dp: 8
+target:
+  output_extract: 0
+  weights_dtype: bfloat16
+  spec:
+    kind: pretrained
+    model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP
+    run_path: goodfire/spd/runs/t-9d2b8f02
+wandb:
+  entity: null
+  project: param-decomp