feat(train): stale-CI replay (S34) — repeat steps reuse the window-first CI envelope

param_decomp/CLAUDE.md

@@ -24,6 +24,13 @@ never silently diverge. Cite IDs (`S14`, `N1`, …) in commit messages and revie
 
 ## Architecture in one breath
 
+**Stale-CI replay** (SPEC S34, experimental, `data.replay_stale_ci`): under free-AT batch
+replay (`train_batch_replay > 1`), repeat steps reuse the window-first CI envelope as a
+constant — no taps / CI-fn forward / CI-fn backward on repeats, ci_fn updates only on
+window-first steps (unit-LR optimizer + in-step global-step LR schedule). Built by
+`train.make_stale_ci_train_steps` (fresh/repeat/compute_ci), dispatched in
+`run.py::do_step`. Off by default; flag-off emits the unchanged step.
+
 `lm.py` defines `DecomposedModel` — a `@runtime_checkable Protocol`: ordered `sites` +
 `leading_axes` + the methods `clean_output`, `read_activations`, `masked_output`,
 `masked_site_outputs`, `weight_deltas`, and a `recon_loss_fn` (LM: `kl_per_position`). The

param_decomp/SPEC.md

@@ -303,6 +303,7 @@ faithfulness, sources/PPGD, the squashings (S5/S6), the imp-min reduction, the g
 | S30 | `cfg.cadence.log_every` divides `cfg.eval.every` (`eval.every % log_every == 0`, asserted at `run.py:255`) so every eval step is also a train-log step. |
 | S31 | The two hidden-acts recon metrics (`CIHiddenActsReconLoss`, `StochasticHiddenActsReconLoss`) are STANDALONE OFFLINE EVAL metrics in JAX (`hidden_acts_eval.py`, wired into `jsp-slow-eval`) — **NOT** recon-grid training terms. `build_loss_terms` still refuses them as training losses (the parameterized recon loss stays KL-on-final-logits only, §2.3–2.5); their objective is per-ELEMENT MSE on each decomposed site's OUTPUT activations, which as a *training* loss is exactly the site-local recon the trainer treats as a conceptual no-no (LOSS_PARITY_DESIGN §4c). The port adds a fifth per-target seam `masked_site_outputs(vu, batch, masks, delta_masks, routes, live, has_delta) -> dict[site, (B,T,d_out)]` (`lm.py`), factored out of `masked_output` (the shared masked forward with a per-site `collect` — the masked per-site output is an intermediate of that forward, so no logic is duplicated). The clean (target) per-site output is the frozen `x @ W`, obtained from the same seam by routing FALSE everywhere (`_site_out`'s frozen branch). Per site, `MSE(masked_site_output, clean_site_output)` with `reduction="sum"` accumulated host-side as `(Σ sum_mse, Σ n_elements)` (token-weighted, exact under micro-batching), divided once at the end; log keys mirror torch exactly (`<ClassName>/<site>` + a combined `<ClassName>` = Σmse/Σn over all sites). `CIHiddenActsReconLoss` is the deterministic `lower_leaky` CI mask, no delta, one forward (tight torch parity); `StochasticHiddenActsReconLoss` draws `n_mask_samples` stochastic CI masks (`mask = ci + (1−ci)·s`) WITH weight deltas (the delta component is always built in JAX runs) — its draws are NOT seed-aligned to torch, so exact bitwise parity is impossible there (expected). Masked + clean run in COMPUTE_DT (bf16, matching the trained model, mirroring `load_run.py`); the MSE reduction is fp32. **AMENDED 2026-06-16** (Oli-approved): superseded the prior "keep-on-bridge / seam refused" decision — these now have a native JAX eval path; the `pd-offline-eval` torch bridge remains available for cross-framework parity checks. |
 | S33 | Fine-tune init (`ExperimentConfig.resume_provenance`, LM-only): a fresh run whose own `ckpts/` is EMPTY and whose `resume_provenance is not None` initializes from a PARENT checkpoint — it loads the parent's `ckpts/<parent_step>` onto the fresh reference `TrainState` and keeps ONLY the trained `components` (V/U) + `ci_fn`; the optimizer states, persistent sources, and `step` are the FRESH reference's (`step = 0`). Rationale: a fine-tune runs a NEW LR/p-anneal schedule computed over the new `cfg.steps` from 0, so carrying stale Adam momentum / a stale adversary would mis-scale the restart; faith warmup is also skipped (the parent's V/U is already faithful). The run records lineage via `resume_provenance` in `config.yaml` + `wandb.config`. The parent's decomposition STRUCTURE (sites names + C, ci-fn arch) must match the new config's — asserted from the parent's pinned `config.yaml` (`run.py::assert_finetune_structural_compat`) before the orbax restore; only LR / coeffs / eps / seq / batch / steps may change. This is distinct from same-config requeue-resume (S22): on a subsequent requeue the run's own `ckpts/` is non-empty, so `restore_latest` from its own dir wins and provenance is ignored. |
+| S34 | **Stale-CI replay** (config-gated, OFF by default: `DataConfig.replay_stale_ci`, requires `train_batch_replay > 1`; **ADDED 2026-07-06**, experimental — an S4-cadence amendment riding the free-AT batch replay): on the REPEAT steps of a replay window the CI envelope is the window-FIRST step's, reused as a CONSTANT — no taps read, no CI-fn forward or backward, no ci_fn/optimizer update (the ci_fn trains only on window-first steps, matching the 1-per-window unique-data rate). V/U, faithfulness, and both adversaries behave exactly as S12′–S14′ with `ci_lower` a constant (their gradients are unchanged by CI being live vs constant); the imp-min term is computed from the cached upper for logging and contributes zero gradient on repeats. The ci_fn LR follows the run's GLOBAL-step schedule via a unit-LR optimizer + in-step `scheduled_value_traced` (`run_state.unit_lr_ci_fn_optimizer`; optax's update count advances only on window-first steps, so a count-driven schedule would stretch by the replay factor); Adam bias correction correctly tracks the update count. A requeue-resume landing MID-window rebuilds the envelope from the current ci_fn (`StaleCITrainSteps.compute_ci`) — one ci_fn update fresher than the interrupted window's envelope, a tiny preemption-only trajectory wobble (weaker than S22's exactness; accepted for the experimental flag). Flag-off emits the pre-S34 step unchanged. |
 
 ## 6. Variation points
 

param_decomp/built_run.py

@@ -49,6 +49,11 @@ class DataConfig:
     """Consecutive train steps that share one token batch (free-AT-style replay: the
     persistent adversary's once-per-step ascent becomes batch-matched on the repeated
     steps). 1 = a fresh batch every step."""
+    replay_stale_ci: bool
+    """Stale-CI replay (SPEC S34): repeat steps of a replay window reuse the window-first
+    step's CI envelope as a constant — no taps / CI-fn forward / CI-fn backward on
+    repeats; the ci_fn updates only on window-first steps. Requires
+    `train_batch_replay > 1`."""
 
 
 @dataclass(frozen=True)

param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay16_staleci_400k.yaml

@@ -0,0 +1,179 @@
+# Stale-CI replay-scaling arm (SPEC S34): free-AT batch replay x16 with
+# replay_stale_ci — repeat steps reuse the window-first CI envelope as a constant
+# (no taps / CI-fn fwd/bwd on repeats; ci_fn updates only on window-first steps, LR
+# still on the global-step schedule). Derived from
+# pile_ppgd_bsc_nw0_lr0p01_replay16_400k.yaml; only run_name / replay_stale_ci differ,
+# so each arm A/Bs stale CI against its running replay-only twin. replay=1 has no
+# repeat steps (no stale-CI arm); the replay-only sweep's replay1 is the shared
+# unique-data baseline. `steps` counts optimizer steps (replay repeats ARE steps):
+# unique batches = 400k / replay.
+run_name: pile4l-ppgd-nw0-lr0p01-replay16-staleci-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/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet
+  dataset_name: parquet
+  eval_split: train
+  is_tokenized: true
+  max_seq_len: 512
+  revision: null
+  shuffle_each_epoch: true
+  train_batch_replay: 16
+  replay_stale_ci: 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: shared_across_batch
+    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: 1.0e-04
+      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: uniform_k_subset
+    type: StochasticReconSubsetLoss
+  - coeff: 0.5
+    n_warmup_steps: 0
+    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: per_batch_per_position
+    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

param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay1_staleci_base_400k.yaml

@@ -0,0 +1,174 @@
+# Code baseline for the stale-CI replay sweep: replay=1, replay_stale_ci OFF, run from
+# the feature/replay-stale-ci branch. Identical config to the running
+# pile_ppgd_bsc_nw0_lr0p01_replay1_400k.yaml arm (only run_name differs) — it pins that
+# the stale-CI refactor's flag-off step path reproduces the parent branch's trajectory
+# at scale before attributing any stale-CI arm differences to the algorithm change.
+run_name: pile4l-ppgd-nw0-lr0p01-replay1-staleci-base-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/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet
+  dataset_name: parquet
+  eval_split: train
+  is_tokenized: true
+  max_seq_len: 512
+  revision: null
+  shuffle_each_epoch: true
+  train_batch_replay: 1
+  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: shared_across_batch
+    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: 1.0e-04
+      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: uniform_k_subset
+    type: StochasticReconSubsetLoss
+  - coeff: 0.5
+    n_warmup_steps: 0
+    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: per_batch_per_position
+    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