Repository Public Research preview
1"""Transplant families for Tier-0 toy models (explicit, separately named, auditable).23All grafts act on the recipient's residual stream at a hook point (default ``resid_post.{layer}``) at4*read positions* (the position whose next-token prediction is scored). The donor, when required, is run5in the loop on the identical current token ids (teacher-forced prefix = whatever the recipient has emitted).67Families8--------9ActivationGraft : x_R <- x_R + α (A(x_D) − x_R) donor-in-the-loop, mapped activation transfer10SubspaceGraft : x_R <- x_R + α P (A(x_D) − x_R) same, restricted to a k-dim subspace (projector P)11SteeringGraft : x_R <- x_R + α v v = mapped mean (donor Ω − donor non-Ω) direction12AdapterGraft : x_R <- x_R + α (f(x_R) − x_R) donor-free; f = small map fitted on alignment data13 to predict A(x_D) from x_R (learned interface)14WeightGraft : recipient parameters of a block replaced by (shape-checked) donor parameters — donor-free15Controls (same interface): sham map (random orthogonal), norm-matched random vector, shuffled-token donor16activation, shuffled-layer donor activation, unrelated-skill donor.1718Targeting: ``positions="all"`` applies at every read position (no oracle); ``positions="omega"`` applies19only at read positions of Ω steps (uses ground-truth program structure; reported as an oracle variant).20"""21from __future__ import annotations2223from dataclasses import dataclass2425import numpy as np26import torch2728from wp_ai_0001.alignment.maps import AffineMap29from wp_ai_0001.models.hooks import Capture, HookContext30from wp_ai_0001.models.toy_transformer import ToyTransformer313233def donor_activation(donor: ToyTransformer, ids: torch.Tensor, positions: torch.Tensor, hook: str) -> torch.Tensor:34 """Donor activation [B, d_D] at the given per-example positions (teacher-forced on ``ids``)."""35 cap = Capture()36 donor.hooks.add(hook, cap)37 try:38 with torch.no_grad():39 donor(ids, positions)40 finally:41 donor.hooks.remove(hook, cap)42 act = cap.store[hook]43 return act[torch.arange(ids.shape[0], device=ids.device), positions]444546@dataclass47class GraftSpec:48 family: str49 recipient_hook: str50 donor_hook: str | None51 alpha: float = 1.052 targeting: str = "all" # "all" | "omega"53 note: str = ""545556class _BaseGraft:57 """Intervention object registered at ``spec.recipient_hook``; computes replacement rows per read position."""5859 def __init__(self, spec: GraftSpec, omega_positions: torch.Tensor | None = None) -> None:60 self.spec = spec61 self.omega_positions = omega_positions # [B, T] bool mask of Ω read positions (targeting="omega")62 self.applied: int = 063 self._busy = False # re-entrancy guard (donor may be the recipient itself)6465 def replacement(self, x_rows: torch.Tensor, ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: # [B, d_R]66 raise NotImplementedError6768 def __call__(self, x: torch.Tensor, ctx: HookContext) -> torch.Tensor:69 if ctx.positions is None or self._busy:70 return x71 self._busy = True72 try:73 return self._apply(x, ctx)74 finally:75 self._busy = False7677 def _apply(self, x: torch.Tensor, ctx: HookContext) -> torch.Tensor:78 B = x.shape[0]79 ar = torch.arange(B, device=x.device)80 pos = ctx.positions81 rows = x[ar, pos]82 new_rows = self.replacement(rows, ctx.extra["ids"], pos)83 active = torch.ones(B, dtype=torch.bool, device=x.device)84 if self.spec.targeting == "omega":85 active = self.omega_positions[ar, pos]86 out = x.clone()87 mixed = rows + self.spec.alpha * (new_rows - rows)88 out[ar, pos] = torch.where(active[:, None], mixed, rows)89 self.applied += int(active.sum())90 return out919293class ActivationGraft(_BaseGraft):94 def __init__(self, spec: GraftSpec, donor: ToyTransformer, amap: AffineMap, omega_positions=None, shuffle_tokens: torch.Generator | None = None) -> None:95 super().__init__(spec, omega_positions)96 self.donor, self.amap, self.shuffle = donor, amap, shuffle_tokens97 self.W = torch.from_numpy(amap.W).float()98 self.b = torch.from_numpy(amap.b).float()99100 def replacement(self, x_rows, ids, positions):101 d = donor_activation(self.donor, ids, positions, self.spec.donor_hook)102 if self.shuffle is not None: # shuffled-token control: donor activation from a different item103 perm = torch.randperm(d.shape[0], generator=self.shuffle).to(d.device)104 d = d[perm]105 return d.float() @ self.W.to(d.device) + self.b.to(d.device)106107108class SubspaceGraft(ActivationGraft):109 def __init__(self, spec, donor, amap, basis: np.ndarray, omega_positions=None) -> None:110 super().__init__(spec, donor, amap, omega_positions)111 q, _ = np.linalg.qr(basis) # [d_R, k]112 self.P = torch.from_numpy((q @ q.T).astype(np.float32))113114 def replacement(self, x_rows, ids, positions):115 mapped = super().replacement(x_rows, ids, positions)116 P = self.P.to(x_rows.device)117 return x_rows + (mapped - x_rows) @ P118119120class SteeringGraft(_BaseGraft):121 def __init__(self, spec, vector: np.ndarray, omega_positions=None) -> None:122 super().__init__(spec, omega_positions)123 self.v = torch.from_numpy(vector.astype(np.float32))124125 def replacement(self, x_rows, ids, positions):126 return x_rows + self.v.to(x_rows.device)127128129class NormMatchedRandomGraft(_BaseGraft):130 """Replace with a random Gaussian row scaled to the norm of the mapped donor activation (per item)."""131132 def __init__(self, spec, donor, amap, seed: int = 0, omega_positions=None) -> None:133 super().__init__(spec, omega_positions)134 self.inner = ActivationGraft(spec, donor, amap)135 self.g = torch.Generator().manual_seed(seed)136137 def replacement(self, x_rows, ids, positions):138 mapped = self.inner.replacement(x_rows, ids, positions)139 r = torch.randn(mapped.shape, generator=self.g).to(mapped.device)140 r = r / r.norm(dim=-1, keepdim=True) * mapped.norm(dim=-1, keepdim=True)141 return r142143144class AdapterGraft(_BaseGraft):145 """Donor-free learned interface f: x_R -> target rows (fitted offline by ridge on alignment data)."""146147 def __init__(self, spec, adapter: AffineMap, omega_positions=None) -> None:148 super().__init__(spec, omega_positions)149 self.W = torch.from_numpy(adapter.W).float()150 self.b = torch.from_numpy(adapter.b).float()151152 def replacement(self, x_rows, ids, positions):153 return x_rows.float() @ self.W.to(x_rows.device) + self.b.to(x_rows.device)154155156def weight_graft(recipient: ToyTransformer, donor: ToyTransformer, parts: list[str], source_layer: int, target_layer: int) -> dict:157 """Copy donor block parameters into the recipient (in place). parts ⊆ {attn, mlp, ln, embed_omega}.158159 Returns a report with shapes checked and the number of copied parameters. Raises on shape mismatch.160 """161 copied = 0162 rb, db = recipient.blocks[target_layer], donor.blocks[source_layer]163 with torch.no_grad():164 for part in parts:165 if part == "attn":166 pairs = [(rb.attn.qkv.weight, db.attn.qkv.weight), (rb.attn.o.weight, db.attn.o.weight)]167 elif part == "mlp":168 pairs = [(rb.mlp[0].weight, db.mlp[0].weight), (rb.mlp[0].bias, db.mlp[0].bias), (rb.mlp[2].weight, db.mlp[2].weight), (rb.mlp[2].bias, db.mlp[2].bias)]169 elif part == "ln":170 pairs = [(rb.ln1.weight, db.ln1.weight), (rb.ln1.bias, db.ln1.bias), (rb.ln2.weight, db.ln2.weight), (rb.ln2.bias, db.ln2.bias)]171 elif part == "embed_omega":172 pairs = [] # handled by caller with the token id173 else:174 raise ValueError(part)175 for r, d in pairs:176 if r.shape != d.shape:177 raise ValueError(f"shape mismatch for {part}: {tuple(r.shape)} vs {tuple(d.shape)}")178 r.copy_(d)179 copied += r.numel()180 return {"parts": parts, "source_layer": source_layer, "target_layer": target_layer, "n_params_copied": copied}181182183def copy_embedding_row(recipient: ToyTransformer, donor: ToyTransformer, token_id: int) -> int:184 with torch.no_grad():185 recipient.tok.weight[token_id].copy_(donor.tok.weight[token_id])186 if recipient.head.weight is not recipient.tok.weight:187 recipient.head.weight[token_id].copy_(donor.head.weight[token_id])188 return int(recipient.tok.weight.shape[1] * (1 if recipient.head.weight is recipient.tok.weight else 2))189