feat(deep-linear): frozen-identity toy domain — the chunkwise-recon testbed

CLAUDE.md

@@ -174,9 +174,9 @@ and returns JAX-native as the #10 torch->jax adapter.
   (the LM composition root — `python -m param_decomp_lab.experiments.lm.run`), `config.py`
   (LM schema + LM build), `load_run.py` (open a finished JAX run), `data.py` /
   `prestage_tokenized.py` (offline tokenize → parquet shards), `jax_launch.py` (`pd-lm`).
-  The TMS and ResidualMLP domains live under `experiments/{tms,resid_mlp}/`
-  (`run.py` + `config.py` + `model.py`; `pd-tms` / `pd-resid-mlp`), calling the core
-  engine as a library.
+  The TMS, ResidualMLP and DeepLinear domains live under
+  `experiments/{tms,resid_mlp,deep_linear}/` (`run.py` + `config.py` + `model.py`;
+  `pd-tms` / `pd-resid-mlp` / `pd-deep-linear`), calling the core engine as a library.
 - `param_decomp_lab/{harvest,autointerp,clustering,investigate}/`
   — post-pipeline stages, each with its own CLAUDE.md.
 - `param_decomp_lab/postprocess/` — orchestrates the post-pipeline stages.
@@ -252,7 +252,7 @@ via `pd-lm`. Slow/plot eval is in-loop only (no CLI).
 | `python -m pretrain.train` | `pretrain/train.py` | The core in-house target-LM pretrainer |
 | `pd-lm` | `experiments/lm/launch.py` | Launch a decomposition trainer run; config-driven via `runtime.dp` (`dp=N` → snapshot + pinned launch config + sbatch across `N//8` nodes, per-node job-side venv; `dp=null` → inline) |
 | `pd-pretrain` | `experiments/lm/pretrain/launch.py` | Launch a pretrainer run; config-driven via `dp` (`dp=N` → sbatch; `dp=null` → inline) |
-| `pd-tms` / `pd-resid-mlp` | `experiments/{tms,resid_mlp}/run.py` | The CPU toy decomposition CLIs |
+| `pd-tms` / `pd-resid-mlp` / `pd-deep-linear` | `experiments/{tms,resid_mlp,deep_linear}/run.py` | The CPU toy decomposition CLIs |
 | `pd-harvest` | `harvest/scripts/run_slurm_cli.py` | Submit harvest SLURM job |
 | `pd-autointerp` | `autointerp/scripts/run_slurm_cli.py` | Submit autointerp SLURM job |
 | `pd-clustering` / `pd-cluster-merge` / `pd-cluster-distances` | `clustering/scripts/` | Clustering ensemble / merge / consensus distances |

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
+)
 
 
 # ---------------------------------------------------------------------------

param_decomp/experiments/invariance_check.py

@@ -120,6 +120,7 @@ def _run(steps: int, sharded: bool) -> list[dict[str, float]]:
             ppgd_cfg,
         ),
         lm.site_names,
+        100,
     )  # fmt: skip
     step = make_train_step(
         lm=lm,