feat(init): coupled component initialization (weight_init: kaiming | coupled)

param_decomp/components.py

@@ -12,7 +12,7 @@
 from collections.abc import Iterator
 from dataclasses import dataclass
 from functools import cache
-from typing import ClassVar, Generic
+from typing import ClassVar, Generic, Literal
 
 import equinox as eqx
 import jax
@@ -72,6 +72,9 @@ def dequantize_fp8(q: Array, scale: Array) -> Array:
 
 VUShape = tuple[int, int, int]  # (d_in, d_out, C)
 
+# How `init_component_stacks` seeds V/U; the config schema (`PDConfig.weight_init`) reuses it.
+WeightInit = Literal["kaiming", "coupled"]
+
 # site name -> (shape group, slot on the group's stack axis); static, canonical site order
 SiteSlots = tuple[tuple[str, VUShape, int], ...]
 
@@ -150,7 +153,7 @@ def group_lengths(self) -> dict[VUShape, int]:
         return lengths
 
 
-def init_stack_arrays(
+def _kaiming_stack_arrays(
     sites: tuple[SiteSpec, ...], key: Array
 ) -> dict[VUShape, tuple[Array, Array]]:
     """Seeded stack init: `{(d_in, d_out, C): (V [g, d_in, C], U [g, C, d_out])}`, vmapped
@@ -169,6 +172,41 @@ def init_stack_arrays(
     return stacked
 
 
+def _coupled_stack_arrays(
+    sites: tuple[SiteSpec, ...], target_weights: dict[str, Array], key: Array
+) -> dict[VUShape, tuple[Array, Array]]:
+    """Coupled stack init: unit-norm seed on the narrow side, wide side its raw W-image.
+
+    `d_in <= d_out`: `v_c ~ unit norm`, `U_c <- (W v_c)^T`; else `u_c ~ unit norm`,
+    `V_c <- W^T u_c`. No C-dependent rescale — components sit at W's natural scale and
+    the component sum equals W restricted to the seed span (`W V V^T` when V is seeded),
+    the delta carrying the complement. One key per site, split in site order; vmapped
+    per shape group like `_kaiming_stack_arrays` so the compiled init doesn't scale with
+    site count."""
+    keys = jax.random.split(key, len(sites))
+    site_index = {spec.name: idx for idx, spec in enumerate(sites)}
+    stacked: dict[VUShape, tuple[Array, Array]] = {}
+    for (d_in, d_out, c), specs in vu_shape_groups(sites).items():
+        ws = jnp.stack([target_weights[spec.name] for spec in specs]).astype(jnp.float32)
+        assert ws.shape == (len(specs), d_out, d_in), (ws.shape, (d_in, d_out, c))
+        ks = keys[jnp.array([site_index[spec.name] for spec in specs])]
+        if d_in <= d_out:
+
+            def seed_v(k: Array, w: Array, s: tuple[int, int] = (d_in, c)) -> tuple[Array, Array]:
+                v = jax.random.normal(k, s)
+                v = v / jnp.linalg.norm(v, axis=0, keepdims=True)
+                return v, (w @ v).T
+        else:
+
+            def seed_v(k: Array, w: Array, s: tuple[int, int] = (c, d_out)) -> tuple[Array, Array]:
+                u = jax.random.normal(k, s)
+                u = u / jnp.linalg.norm(u, axis=1, keepdims=True)
+                return w.T @ u.T, u
+
+        stacked[(d_in, d_out, c)] = jax.vmap(seed_v)(ks, ws)
+    return stacked
+
+
 def component_stacks_from_sites(vu: dict[str, tuple[Array, Array]]) -> ComponentStacks:
     """Build the stacked `ComponentStacks` from a per-site `{name: (V, U)}` dict (site order =
     dict order). The explicit-arrays constructor for toys and tests; the trainer inits
@@ -187,11 +225,25 @@ def component_stacks_from_sites(vu: dict[str, tuple[Array, Array]]) -> Component
     return ComponentStacks(stacks=stacks, site_slots=site_slots_for(sites))
 
 
-def init_component_stacks(sites: tuple[SiteSpec, ...], key: Array) -> ComponentStacks:
-    """Small random fp32 V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5) per site, built directly in
-    the stacked persistence layout; the weight-delta channel carries the faithfulness
-    residual at init (before faithfulness warmup)."""
-    return ComponentStacks(stacks=init_stack_arrays(sites, key), site_slots=site_slots_for(sites))
+def init_component_stacks(
+    sites: tuple[SiteSpec, ...],
+    key: Array,
+    weight_init: WeightInit = "kaiming",
+    target_weights: dict[str, Array] | None = None,
+) -> ComponentStacks:
+    """Seeded V/U masters in the stacked persistence layout, per `weight_init`:
+    'kaiming' — small random fp32 V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5), ignores W
+    (the weight-delta channel carries the faithfulness residual at init); 'coupled' —
+    `_coupled_stack_arrays`, seeded from `target_weights` (required there, refused
+    otherwise)."""
+    match weight_init:
+        case "kaiming":
+            assert target_weights is None, "kaiming init draws blind — no target_weights"
+            stacks = _kaiming_stack_arrays(sites, key)
+        case "coupled":
+            assert target_weights is not None, "coupled init seeds from W — pass target_weights"
+            stacks = _coupled_stack_arrays(sites, target_weights, key)
+    return ComponentStacks(stacks=stacks, site_slots=site_slots_for(sites))
 
 
 def site_out(

param_decomp/configs.py

@@ -31,6 +31,7 @@
 )
 
 from param_decomp.base_config import BaseConfig, Probability
+from param_decomp.components import WeightInit
 from param_decomp.schedule import ScheduleConfig
 
 # ---------------------------------------------------------------------------
@@ -1041,6 +1042,14 @@ def _strip_removed_jax_unsupported_fields(cls, data: object) -> object:
         default=0,
         description="Random seed for reproducibility, including LM dataset shuffling.",
     )
+    weight_init: WeightInit = Field(
+        default="kaiming",
+        description="How component V/U are initialized. 'kaiming': iid normal "
+        "(V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5)), ignores W. 'coupled': unit-norm seed on "
+        "the narrow side, wide side its raw W-image (U_c from (W v_c)^T when d_in <= d_out, "
+        "else V_c from W^T u_c) — W-natural scale, component sum ~ W on a rank-C random "
+        "subspace (the delta carries the complement).",
+    )
     loss_metrics: list[AnyLossMetricConfig] = Field(
         default_factory=list,
         description=(

param_decomp/model.py

@@ -217,6 +217,10 @@ def masked_site_outputs(
         the recon grid, which stays KL-on-final-logits."""
         ...
 
+    def target_weight(self, name: str) -> Float[Array, "d_out d_in"]:
+        """The frozen target weight W for site `name`, in its stored dtype."""
+        ...
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Float[Array, "d_out d_in"]]:
         """fp32 `W − V@U` per site from the fp32 master `vu`."""
         ...

param_decomp/run.py

@@ -516,7 +516,8 @@ def run_decomposition_training(
     rules = placement_rules
     if is_main:
         audit = component_stacks_audit(
-            eqx.filter_eval_shape(_partial(init_component_stacks, model.sites), init_key), rules
+            eqx.filter_eval_shape(_partial(init_component_stacks, model.sites), init_key),
+            rules,
         )
         print(
             rules.describe(

param_decomp/run_state.py

@@ -20,6 +20,7 @@
 
 from param_decomp.adversary import PersistentAdversary, init_sources_adam_state
 from param_decomp.ci_fn import ChunkwiseTransformerCIArch, CIFnArch
+from param_decomp.components import WeightInit
 from param_decomp.configs import (
     AdamPGDConfig,
     AdamWOptimizerConfig,
@@ -163,14 +164,23 @@ def init_decomposition(
     init_key: PRNGKeyArray,
     mesh: Mesh,
     rules: PlacementRules,
+    weight_init: WeightInit = "kaiming",
 ) -> Decomposition:
     """The trained-product half of `init_train_state`, factored out so a consumer can
     `jax.eval_shape` it to recover the saved `decomposition` item's tree structure
-    without building (or knowing about) the optimizers/adversaries."""
+    without building (or knowing about) the optimizers/adversaries (the `weight_init`
+    default is safe there — both inits produce the same tree)."""
     ci_key = random.fold_in(init_key, 1)
+    match weight_init:
+        case "kaiming":
+            target_weights = None
+        case "coupled":
+            target_weights = {name: model.target_weight(name) for name in model.site_names}
     # V/U placement derives from the rules table; the CI fn still declares its own
     # per-leaf shardings (the staged placement migration hasn't reached it).
-    components = init_component_stacks_placed(model.sites, init_key, rules)
+    components = init_component_stacks_placed(
+        model.sites, init_key, rules, weight_init, target_weights
+    )
     ci_fn = init_ci_fn_placed(ci_fn_arch, model.sites, ci_key, mesh)
     assert ci_fn.has_position_axis == model.has_position_axis, (
         f"CI fn has_position_axis={ci_fn.has_position_axis} but model declares "
@@ -195,7 +205,7 @@ def init_train_state(
     assert isinstance(positions, Positioned) == model.has_position_axis, (
         f"{positions} does not match the model's has_position_axis={model.has_position_axis}"
     )
-    decomposition = init_decomposition(model, ci_fn_arch, init_key, mesh, rules)
+    decomposition = init_decomposition(model, ci_fn_arch, init_key, mesh, rules, pd.weight_init)
     components, ci_fn = decomposition.components, decomposition.ci_fn
     losses = build_loss_terms(pd.loss_metrics, model.site_names)
     persistent = persistent_configs(losses.recon)

param_decomp/targets/glu_transformer.py

@@ -1136,12 +1136,17 @@ def masked_component_activations(
         assert set(collect_activations) == set(live), (sorted(collect_activations), sorted(live))
         return collect_activations
 
+    def target_weight(self, name: str) -> Array:
+        """The frozen W for site `name` (stored dtype): the site's slice of the stacked
+        per-layer weights."""
+        layer, kind = parse_site_name(name)
+        return _frozen_site_weight(jax.tree.map(lambda a, li=layer: a[li], self.stacked), kind)
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         """fp32 `W − V@U` per site from fp32 masters (faithfulness input)."""
         out: dict[str, Array] = {}
         for spec in self.sites:
-            layer, kind = parse_site_name(spec.name)
-            W = _frozen_site_weight(jax.tree.map(lambda a, li=layer: a[li], self.stacked), kind)
+            W = self.target_weight(spec.name)
             V, U = vu.site(spec.name)
             out[spec.name] = (
                 W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T

param_decomp/targets/glu_transformer_sharding.py

@@ -63,6 +63,7 @@
 from param_decomp.components import (
     ComponentStacks,
     SiteSpec,
+    WeightInit,
     init_component_stacks,
 )
 from param_decomp.configs import SourceShape
@@ -89,15 +90,23 @@ def place_target(tgt: GLUDecomposedModel, mesh: Mesh) -> GLUDecomposedModel:
 
 
 def init_component_stacks_placed(
-    sites: tuple[SiteSpec, ...], key: PRNGKeyArray, rules: PlacementRules
+    sites: tuple[SiteSpec, ...],
+    key: PRNGKeyArray,
+    rules: PlacementRules,
+    weight_init: WeightInit = "kaiming",
+    target_weights: dict[str, Array] | None = None,
 ) -> ComponentStacks:
-    """Seeded V/U init placed by `component_stacks_shardings(_, rules)` (the run's placement
-    policy), values bit-identical to the retired per-site init (pinned by `test_sharding`).
-    One jit, 2×n_shapes sharded outputs — the persistence layout IS the stacked layout, so
-    the old two-stage stack-then-unstack fan-out (and its transient extra copy) is gone."""
-    abstract = eqx.filter_eval_shape(partial(init_component_stacks, sites), key)
-    placement = component_stacks_shardings(abstract, rules)
-    return jax.jit(partial(init_component_stacks, sites), out_shardings=placement)(key)
+    """`init_component_stacks` placed by `component_stacks_shardings(_, rules)` (the run's
+    placement policy); kaiming values bit-identical to the retired per-site init (pinned
+    by `test_sharding`). One jit, 2×n_shapes sharded outputs — the persistence layout IS
+    the stacked layout, so no host-side full tree ever exists; `key`/`target_weights` ride
+    as jit ARGUMENTS (not closure constants) so the frozen W arrays are not baked into the
+    compiled init."""
+    init = partial(init_component_stacks, sites, weight_init=weight_init)
+    placement = component_stacks_shardings(
+        eqx.filter_eval_shape(init, key, target_weights=target_weights), rules
+    )
+    return jax.jit(init, out_shardings=placement)(key, target_weights=target_weights)
 
 
 def init_ci_fn_placed(

param_decomp/targets/llama_simple_mlp.py

@@ -488,12 +488,16 @@ def masked_site_outputs(
         assert set(collect) == set(live), (sorted(collect), sorted(live))
         return collect
 
+    def target_weight(self, name: str) -> Array:
+        """The frozen W for site `name`, in its stored dtype."""
+        layer_idx, kind = parse_site_name(name)
+        return _frozen_site_weight(self.layers[layer_idx], kind)
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         """fp32 `W − V@U` per site from fp32 masters (faithfulness input)."""
         out: dict[str, Array] = {}
         for spec in self.sites:
-            layer_idx, kind = parse_site_name(spec.name)
-            W = _frozen_site_weight(self.layers[layer_idx], kind)
+            W = self.target_weight(spec.name)
             V, U = vu.site(spec.name)
             out[spec.name] = (
                 W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T

param_decomp/tests/test_attn_patterns_eval.py

@@ -275,6 +275,10 @@ def masked_site_outputs(
         del vu, resid, masks, delta_masks, routes, live, has_delta
         raise AssertionError("positionless stub fn must not be called")
 
+    def target_weight(self, name: str) -> jax.Array:
+        del name
+        raise AssertionError("positionless stub fn must not be called")
+
     def weight_deltas(self, vu: Any) -> dict[str, jax.Array]:
         del vu
         raise AssertionError("positionless stub fn must not be called")

param_decomp/tests/test_component_init.py

@@ -0,0 +1,60 @@
+"""Tests for the `coupled` mode of `init_component_stacks`."""
+
+import jax
+import jax.numpy as jnp
+
+from param_decomp.components import SiteSpec, init_component_stacks
+
+
+def _random_weight(d_out: int, d_in: int) -> jax.Array:
+    return jax.random.normal(jax.random.PRNGKey(0), (d_out, d_in))
+
+
+def _col_space_residual(u: jax.Array, w: jax.Array) -> jax.Array:
+    q_out, _, _ = jnp.linalg.svd(w, full_matrices=False)
+    return u - (u @ q_out) @ q_out.T
+
+
+def _row_space_residual(v: jax.Array, w: jax.Array) -> jax.Array:
+    _, _, vh = jnp.linalg.svd(w, full_matrices=False)
+    q_in = vh.T
+    return v - q_in @ (q_in.T @ v)
+
+
+def _init_single(w: jax.Array, c: int, seed: int) -> tuple[jax.Array, jax.Array]:
+    d_out, d_in = w.shape
+    sites = (SiteSpec(name="m", d_in=d_in, d_out=d_out, C=c),)
+    vu = init_component_stacks(sites, jax.random.PRNGKey(seed), "coupled", {"m": w})
+    return vu.site("m")
+
+
+def test_coupled_seeds_are_unit_norm_and_derived_side_is_raw_w_image():
+    w = _random_weight(10, 6)  # d_in < d_out: V seeded, U derived
+    V, U = _init_single(w, c=8, seed=0)
+    assert jnp.allclose(jnp.linalg.norm(V, axis=0), jnp.ones(8), atol=1e-5)
+    assert jnp.allclose(U, (w @ V).T, atol=1e-6)
+    assert jnp.abs(_col_space_residual(U, w)).max() < 1e-5
+
+    w_t = w.T  # d_in > d_out: U seeded, V derived
+    V_t, U_t = _init_single(w_t, c=8, seed=0)
+    assert jnp.allclose(jnp.linalg.norm(U_t, axis=1), jnp.ones(8), atol=1e-5)
+    assert jnp.allclose(V_t, w_t.T @ U_t.T, atol=1e-6)
+    assert jnp.abs(_row_space_residual(V_t, w_t)).max() < 1e-5
+
+
+def test_coupled_stacked_sites_get_independent_draws_coupled_to_their_own_w():
+    # Two sites sharing one (d_in, d_out, C) shape group stack on one axis; each slice
+    # must be coupled to ITS OWN W with an independent seed draw.
+    key = jax.random.PRNGKey(3)
+    w_a = jax.random.normal(jax.random.fold_in(key, 0), (10, 6))
+    w_b = jax.random.normal(jax.random.fold_in(key, 1), (10, 6))
+    sites = (
+        SiteSpec(name="a", d_in=6, d_out=10, C=8),
+        SiteSpec(name="b", d_in=6, d_out=10, C=8),
+    )
+    vu = init_component_stacks(sites, jax.random.PRNGKey(0), "coupled", {"a": w_a, "b": w_b})
+    V_a, U_a = vu.site("a")
+    V_b, U_b = vu.site("b")
+    assert not jnp.allclose(V_a, V_b)
+    assert jnp.allclose(U_a, (w_a @ V_a).T, atol=1e-6)
+    assert jnp.allclose(U_b, (w_b @ V_b).T, atol=1e-6)

param_decomp/tests/test_eval.py

@@ -131,6 +131,10 @@ def masked_site_outputs(
         del vu, resid, masks, delta_masks, routes, live, has_delta
         raise AssertionError("positionless stub fn must not be called")
 
+    def target_weight(self, name: str) -> jax.Array:
+        del name
+        raise AssertionError("positionless stub fn must not be called")
+
     def weight_deltas(self, vu: Any) -> dict[str, jax.Array]:
         del vu
         raise AssertionError("positionless stub fn must not be called")

param_decomp/tests/test_generic_model_io.py

@@ -166,6 +166,10 @@ def masked_site_outputs(
             hidden = hidden + delta_masks[SITE][..., None] * (resid @ delta.T)
         return {SITE: hidden}
 
+    def target_weight(self, name: str) -> Array:
+        assert name == SITE, name
+        return self.W
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         V, U = vu.site(SITE)
         return {

param_decomp_lab/experiments/resid_mlp/model.py

@@ -457,6 +457,9 @@ def masked_site_outputs(
             self.target, prepared, resid, masks, delta_masks, routes, live, has_delta
         )
 
+    def target_weight(self, name: str) -> Array:
+        return _frozen_site_weight(self.target, name)
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         return weight_deltas_fp32(self.target, vu, self.sites)
 

param_decomp_lab/experiments/tms/model.py

@@ -389,6 +389,9 @@ def masked_site_outputs(
             self.target, prepared, resid, masks, delta_masks, routes, live, has_delta
         )
 
+    def target_weight(self, name: str) -> Array:
+        return _frozen_site_weight(self.target, name)
+
     def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]:
         return weight_deltas_fp32(self.target, vu, self.sites)