Consistent tPD metric names, and core evals on both streams

param_decomp/core/configs.py

@@ -233,6 +233,7 @@ class SmoothL0ImportanceMinimalityLossConfig(LossMetricConfig):
 
 
 class CIMaskedReconLossConfig(LossMetricConfig, HiddenActsReconstructionMixin):
+    slow: ClassVar[bool] = False
     type: Literal["CIMaskedReconLoss"] = "CIMaskedReconLoss"
 
 
@@ -467,6 +468,15 @@ class ComponentActivationDensityConfig(BaseConfig):
     ci_alive_threshold: float = 0.0
 
 
+class WeightMagnitudeConfig(BaseConfig):
+    """Per-site `‖V_c‖·‖U_c‖` scatter in descending magnitude order, log y.
+
+    Reads the trained V/U alone — no forward pass and no eval batch."""
+
+    slow: ClassVar[bool] = True
+    type: Literal["WeightMagnitude"] = "WeightMagnitude"
+
+
 class IdentityCITargetSpec(BaseConfig):
     """A layer expected to produce an Identity CI pattern over `n_features` features."""
 
@@ -657,6 +667,7 @@ class UnmaskedNoDeltaReconLossConfig(LossMetricConfig):
     exists here, and it is non-target-only — the plain and target-pass unions have no
     member for it."""
 
+    slow: ClassVar[bool] = False
     type: Literal["UnmaskedNoDeltaReconLoss"] = "UnmaskedNoDeltaReconLoss"
 
 

param_decomp/core/run.py

@@ -333,7 +333,10 @@ def log(self, step: int, record: "LogRecord") -> None:
         self._last_committed_step = step
         record = {
             _METRIC_KEYS.get(
-                k, f"train/{k}" if k.startswith(("grad_norms/", "loss/", "schedules/")) else k
+                k,
+                f"train/{k}"
+                if k.startswith(("grad_norms/", "loss/", "schedules/", "nontarget_data/"))
+                else k,
             ): v
             for k, v in record.items()
         }  # keys already starting "train/" or "eval/" pass through verbatim

param_decomp/core/slow_eval.py

@@ -68,6 +68,7 @@
     lower_leaky_hard_sigmoid,
     upper_leaky_hard_sigmoid,
 )
+from param_decomp.core.components import ComponentStacks
 from param_decomp.core.configs import (
     DenseCITargetSpec,
     IdentityCIErrorConfig,
@@ -198,6 +199,11 @@ def slow_eval_step(
     return filter_jit(slow_eval_step, compiler_options=compiler_options)
 
 
+def _raw_sample(chunks: dict[str, list[np.ndarray]], site: str) -> np.ndarray:
+    """A site's kept raw values, or empty when `n_batches_accum` kept none."""
+    return np.concatenate(chunks[site]) if site in chunks else np.empty(0, np.float32)
+
+
 def accumulate_site_reductions(
     slow_eval_step: SlowEvalStep,
     model: DecomposedModel,
@@ -246,8 +252,8 @@ def accumulate_site_reductions(
             density_counts=density[site],
             ci_sums=sums[site],
             n_positions=total_positions,
-            lower_sample=np.concatenate(lower_chunks[site]),
-            preactivations_sample=np.concatenate(preactivations_chunks[site]),
+            lower_sample=_raw_sample(lower_chunks, site),
+            preactivations_sample=_raw_sample(preactivations_chunks, site),
             density_hist=hist.get(site),
         )
         for site in density
@@ -578,6 +584,89 @@ def _plot_ci_matrices(matrices: dict[str, np.ndarray], colormap: str, title_pref
     return _render_figure(fig)
 
 
+def _component_weight_magnitudes(components: ComponentStacks) -> dict[str, Array]:
+    return {
+        name: jnp.linalg.norm(sc.V.astype(jnp.float32), axis=0)
+        * jnp.linalg.norm(sc.U.astype(jnp.float32), axis=1)
+        for name, sc in components.sites_items()
+    }
+
+
+def weight_magnitudes(components: ComponentStacks) -> dict[str, np.ndarray]:
+    """Per-site `‖V_c‖·‖U_c‖` as host `(C,)` vectors. The norms reduce ON DEVICE, so only
+    C floats per site cross the boundary — never the V/U matrices themselves."""
+    return {
+        name: np.asarray(value) for name, value in _component_weight_magnitudes(components).items()
+    }
+
+
+def mean_cis(reductions: dict[str, SiteReduction]) -> dict[str, np.ndarray]:
+    """Per-site token-weighted mean CI."""
+    assert all(r.n_positions > 0 for r in reductions.values())
+    return {site: r.ci_sums / r.n_positions for site, r in reductions.items()}
+
+
+def plot_weight_magnitudes(magnitudes: dict[str, np.ndarray]) -> bytes:
+    """Per-site `‖V_c‖·‖U_c‖` in descending magnitude order, log y. x is a component's rank
+    within its site, NOT its component id."""
+    n_rows, n_cols = _grid_dims(len(magnitudes))
+    fig = Figure(figsize=(8 * n_cols, 3 * n_rows))
+    axs = fig.subplots(n_rows, n_cols, squeeze=False)
+    flat_axes = axs.T.ravel()
+    for ax in flat_axes[len(magnitudes) :]:
+        ax.set_visible(False)
+    for ax, (name, values) in zip(flat_axes, magnitudes.items(), strict=False):
+        ax.scatter(range(len(values)), np.sort(values)[::-1], marker="x", s=10)
+        ax.set_yscale("log")
+        ax.set_xlabel("Component (descending ‖V‖·‖U‖)")
+        ax.set_ylabel("‖V‖·‖U‖")
+        ax.set_title(name, fontsize=10)
+    fig.tight_layout()
+    return _render_figure(fig)
+
+
+def plot_mean_component_cis_two_streams(
+    target_mean_cis: dict[str, np.ndarray],
+    nontarget_mean_cis: dict[str, np.ndarray],
+) -> tuple[bytes, bytes]:
+    """Both streams' mean CI on one axis per site, ordered by descending TARGET mean.
+
+    The nontarget series takes the same permutation rather than its own, so a component's
+    two series line up vertically."""
+    assert target_mean_cis.keys() == nontarget_mean_cis.keys(), (
+        sorted(target_mean_cis),
+        sorted(nontarget_mean_cis),
+    )
+    n_rows, n_cols = _grid_dims(len(target_mean_cis))
+    ordered = {
+        name: (target[order], nontarget_mean_cis[name][order])
+        for name, target in target_mean_cis.items()
+        for order in [np.argsort(target)[::-1]]
+    }
+    images: list[bytes] = []
+    for log_y in (False, True):
+        fig = Figure(figsize=(8 * n_cols, 3 * n_rows))
+        axs = fig.subplots(n_rows, n_cols, squeeze=False)
+        flat_axes = axs.T.ravel()
+        for ax in flat_axes[len(ordered) :]:
+            ax.set_visible(False)
+        for ax, (name, (target, nontarget)) in zip(flat_axes, ordered.items(), strict=False):
+            x = np.arange(len(target))
+            if log_y:
+                ax.set_yscale("log")
+            ax.fill_between(x, target, step="mid", color="#1f77b4", label="target")
+            ax.fill_between(
+                x, nontarget, step="mid", color="#d62728", label="non-target", alpha=0.6
+            )
+            ax.set_xlabel("Component (sorted by target mean CI)")
+            ax.set_ylabel("mean CI")
+            ax.set_title(name, fontsize=10)
+            ax.legend(fontsize=7)
+        fig.tight_layout()
+        images.append(_render_figure(fig))
+    return images[0], images[1]
+
+
 def plot_permuted_ci_heatmaps(
     position_ci: dict[str, PositionCI], permutation: dict[str, "Literal['identity', 'dense']"]
 ) -> tuple[bytes, bytes]:

param_decomp/core/tests/test_eval_averaging_parity.py

@@ -132,12 +132,14 @@ def step(
         jnp.array([0, 0], dtype=jnp.uint32),
         train_steps=0,
         eval_steps=2,
+        stream="nontarget",
     )
     context = LMEvalContext(
         state=_state_stub(),  # pyright: ignore[reportArgumentType]
         now_step=0,
         pass_index=0,
         batches=(jnp.asarray(1.0), jnp.asarray(3.0)),
+        target_batches=None,
     )
     assert operation.run(context)["eval/loss/probe/hidden_acts_reconstruction"] == 2.0
 

param_decomp/core/train.py

@@ -1185,14 +1185,18 @@ def make_targeted_train_step[PreparedT](
             for term in objective.target.recon
             if term.hidden_acts_reconstruction is not None
         },
-        "nontarget/impmin": objective.nontarget.impmin_coeff,
-        **{f"nontarget/{term.name}": term.coeff for term in nt_terms},
     }
     if objective.target.imp.cfg.frequency is not None:
         coeff_schedules[f"{objective.target.imp.name}/frequency"] = (
             objective.target.imp.cfg.frequency.coeff
         )
 
+    imp_name = objective.target.imp.name
+    nontarget_coeff_schedules: dict[str, LossCoeff] = {
+        imp_name: objective.nontarget.impmin_coeff,
+        **{term.name: term.coeff for term in nt_terms},
+    }
+
     def nontarget_draw_loss(
         model: DecomposedModel[PreparedT],
         prepared_weights: PreparedT,
@@ -1360,8 +1364,8 @@ def loss_fn(
             nt_imp_lp, nt_imp_freq = imp_min_terms(nt_ci.upper, atoms.imp_min, imp_min_param)
             nt_total = nt_imp_coeff * nt_imp_lp + freq_coeff * nt_imp_freq
             nt_aux = {
-                f"loss/nontarget/{atoms.imp_loss_key}": nt_imp_lp,
-                "loss/nontarget/freq": nt_imp_freq,
+                f"nontarget_data/loss/{imp_name}": nt_imp_lp,
+                "nontarget_data/loss/FrequencyMinimalityLoss": nt_imp_freq,
             }
             nt_breakdowns = atoms.grid_losses(
                 nt_terms,
@@ -1374,8 +1378,8 @@ def loss_fn(
                 nt_terms, nt_recon_coeffs, nt_breakdowns, strict=True
             ):
                 nt_total = nt_total + coeff * breakdown.total
-                nt_aux[f"loss/nontarget/{term.name}"] = breakdown.total
-            nt_aux["loss/nontarget/total"] = nt_total
+                nt_aux[f"nontarget_data/loss/{term.name}"] = breakdown.total
+            nt_aux["nontarget_data/loss/total"] = nt_total
             total_loss = total_loss + nt_total
             reported_total = reported_total + nt_total
             return total_loss, (reported_total, imp_lp, imp_freq, term_breakdowns, nt_aux)
@@ -1446,6 +1450,12 @@ def loss_fn(
             | nt_aux
             | wd_metrics
             | _scheduled_coeff_metrics(step_f32, atoms.total_steps, coeff_schedules)
+            | {
+                f"nontarget_data/{key}": value
+                for key, value in _scheduled_coeff_metrics(
+                    step_f32, atoms.total_steps, nontarget_coeff_schedules
+                ).items()
+            }
         )
         return new_state, metrics
 

param_decomp/experiments/eval_config.py

@@ -14,6 +14,7 @@
     CI_L0Config,
     CIHiddenActsReconLossConfig,
     CIHistogramsConfig,
+    CIMaskedReconLossConfig,
     CIMeanPerComponentConfig,
     ComponentActivationDensityConfig,
     HiddenActsReconstructionMixin,
@@ -22,7 +23,9 @@
     PermutedCIPlotsConfig,
     PGDReconLossConfig,
     StochasticHiddenActsReconLossConfig,
+    UnmaskedNoDeltaReconLossConfig,
     UVPlotsConfig,
+    WeightMagnitudeConfig,
     WellTemperednessConfig,
 )
 from param_decomp.core.eval_schedule import EvalSchedule, Every, FirstThenEvery
@@ -31,6 +34,7 @@
     CEandKLLossesConfig,
     CIMaskedAttnPatternsReconLossConfig,
     StochasticAttnPatternsReconLossConfig,
+    TwoStreamCIMeanPerComponentConfig,
 )
 
 AnyEvalMetricConfig = Annotated[
@@ -40,14 +44,18 @@
     | CIHistogramsConfig
     | CI_L0Config
     | CIMaskedAttnPatternsReconLossConfig
+    | CIMaskedReconLossConfig
     | CIMeanPerComponentConfig
     | ComponentActivationDensityConfig
     | IdentityCIErrorConfig
     | PermutedCIPlotsConfig
     | PGDReconLossConfig
     | StochasticAttnPatternsReconLossConfig
     | StochasticHiddenActsReconLossConfig
+    | TwoStreamCIMeanPerComponentConfig
+    | UnmaskedNoDeltaReconLossConfig
     | UVPlotsConfig
+    | WeightMagnitudeConfig
     | WellTemperednessConfig,
     Discriminator("type"),
 ]