setup

src/projected_grpo/__init__.py

@@ -0,0 +1,5 @@
+import os
+
+if os.environ.get("BEARTYPE"):
+    from beartype.claw import beartype_this_package
+    beartype_this_package()

src/projected_grpo/extract_vhack.py

@@ -0,0 +1,80 @@
+"""Extract v_hack from contrastive pairs of hidden states.
+
+Per Wu-Tang (2026, arXiv 2604.01476) §3.1:
+
+    d = (1/N) * sum_i (h_i^+ - h_i^-)
+
+where h^+ are last-token hidden states from hack-flavored prompts and h^- from
+clean ones, taken at intermediate-to-late layers (60-75% of model depth).
+
+Validation: held-out separation accuracy > 90%.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from jaxtyping import Float
+from loguru import logger
+from torch import Tensor
+
+
+@dataclass
+class VHackResult:
+    v_hack: Float[Tensor, "d"]            # unit-normed direction
+    val_accuracy: float                    # held-out hack-vs-clean separation accuracy
+    layer_idx: int
+    n_train: int
+    n_val: int
+
+
+@torch.no_grad()
+def collect_last_token_hidden(
+    model,
+    tokenizer,
+    prompts: list[str],
+    layer_idx: int,
+    device: str = "cuda",
+) -> Float[Tensor, "n d"]:
+    """Forward each prompt, return last-token hidden state at layer_idx."""
+    hs = []
+    for p in prompts:
+        ids = tokenizer(p, return_tensors="pt").to(device)
+        out = model(**ids, output_hidden_states=True)
+        # out.hidden_states is tuple of (n_layers+1,) tensors of shape (1, seq, d)
+        h = out.hidden_states[layer_idx][0, -1, :].cpu()  # "d"
+        hs.append(h)
+    return torch.stack(hs, dim=0)
+
+
+def extract_vhack(
+    h_hack_train: Float[Tensor, "n_train d"],
+    h_clean_train: Float[Tensor, "n_train d"],
+    h_hack_val: Float[Tensor, "n_val d"],
+    h_clean_val: Float[Tensor, "n_val d"],
+    layer_idx: int,
+) -> VHackResult:
+    """Mean-difference direction with held-out validation."""
+    v = (h_hack_train.mean(dim=0) - h_clean_train.mean(dim=0))
+    v = v / (v.norm() + 1e-12)
+
+    # Validate: projection score on hack should exceed clean.
+    s_hack = h_hack_val @ v
+    s_clean = h_clean_val @ v
+    # paired accuracy: each (hack, clean) pair, hack should score higher
+    correct = (s_hack > s_clean).float().mean().item()
+
+    logger.info(
+        f"v_hack extracted layer={layer_idx} n_train={len(h_hack_train)} "
+        f"n_val={len(h_hack_val)} val_acc={correct:.3f} "
+        f"SHOULD val_acc>0.9 ELSE pair quality or layer is wrong"
+    )
+
+    return VHackResult(
+        v_hack=v,
+        val_accuracy=correct,
+        layer_idx=layer_idx,
+        n_train=len(h_hack_train),
+        n_val=len(h_hack_val),
+    )

src/projected_grpo/grad_proj.py

@@ -0,0 +1,82 @@
+"""Gradient projection against a hack direction in SVD-of-W basis.
+
+Math (from spec.md §5):
+
+    cos_α = <g, v_hack> / ||g||                                # alignment in [-1, 1]
+    if cos_α > 0:
+        g' = g - cos_α * ||g|| * v_hack                        # remove component along v_hack
+        g' = g' * ||g|| / ||g'||                               # restore magnitude (optional)
+    else:
+        g' = g
+
+SVD denoising of v_hack (from spec.md §4):
+
+    W = U S V^T              # SVD of a chosen W matrix (residual stream out)
+    v_S = V[:, :m].T @ v     # project into top-m basis
+    v   = V[:, :m]  @ v_S    # reproject back
+    v   = v / ||v||
+"""
+
+from __future__ import annotations
+
+import torch
+from jaxtyping import Float
+from torch import Tensor
+
+
+def svd_denoise(
+    v: Float[Tensor, "d"],
+    W: Float[Tensor, "d_out d_in"],
+    m: int,
+    use_left: bool = False,
+) -> Float[Tensor, "d"]:
+    """Project v into top-m SVD basis of W and reproject. Normalize.
+
+    use_left=False projects via V (right singular vectors, d_in space).
+    use_left=True  projects via U (left singular vectors, d_out space).
+    Choose based on which side of W aligns with v's residual-stream dim.
+    """
+    U, S, Vh = torch.linalg.svd(W, full_matrices=False)  # U: d_out r, S: r, Vh: r d_in
+    basis = U[:, :m] if use_left else Vh[:m].T  # "d m"
+    if basis.shape[0] != v.shape[0]:
+        raise ValueError(
+            f"v.shape={v.shape} basis.shape={basis.shape}; "
+            "set use_left to match residual-stream dim of v."
+        )
+    v_S = basis.T @ v          # "m"
+    v_denoised = basis @ v_S   # "d"
+    return v_denoised / (v_denoised.norm() + 1e-12)
+
+
+def project_gradient(
+    g: Float[Tensor, "D"],
+    v_hack: Float[Tensor, "D"],
+    preserve_magnitude: bool = True,
+) -> tuple[Float[Tensor, "D"], dict[str, float]]:
+    """One-sided gradient projection. Returns (g_projected, diagnostics).
+
+    Only projects when cos_align > 0 (gradient is pushing toward hack).
+    """
+    g_norm = g.norm()
+    # cos(g, v_hack) where v_hack is assumed unit.
+    cos_in = (g @ v_hack) / (g_norm + 1e-12)
+    if cos_in.item() <= 0:
+        return g, {
+            "cos_in": cos_in.item(), "cos_out": cos_in.item(),
+            "projected": 0.0,
+            "g_norm_before": g_norm.item(), "g_norm_after": g_norm.item(),
+        }
+
+    # Remove component along v_hack.
+    g_prime = g - cos_in * g_norm * v_hack
+    g_prime_norm = g_prime.norm()
+    if preserve_magnitude and g_prime_norm > 1e-12:
+        g_prime = g_prime * (g_norm / g_prime_norm)
+    cos_out = (g_prime @ v_hack) / (g_prime.norm() + 1e-12)
+    return g_prime, {
+        "cos_in": cos_in.item(),
+        "cos_out": cos_out.item(),
+        "projected": 1.0,
+        "g_norm_before": g_norm.item(),
+        "g_norm_after": g_prime.norm().item(),
+    }

src/projected_grpo/run.py

@@ -0,0 +1,194 @@
+"""Smoke / fast-dev-run entry point.
+
+Runs the real pipeline end-to-end on a tiny-random model:
+  1. Load model + tokenizer
+  2. Build tiny contrastive pair set (synthetic hack/clean prompts)
+  3. Extract v_hack at a middle layer
+  4. SVD-denoise v_hack via a chosen weight matrix
+  5. Run a few fake GRPO steps where each step:
+     - computes a fake gradient (random + small bias toward v_hack to simulate
+       the policy discovering the loophole)
+     - optionally projects gradient against v_hack (--arm=projected)
+     - logs cos_align before/after
+  6. Print final TSV row summarizing the run.
+
+Garbage numbers are fine. We test code paths, not correctness.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from dataclasses import dataclass
+
+import torch
+import tyro
+from loguru import logger
+from tabulate import tabulate
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from projected_grpo.extract_vhack import collect_last_token_hidden, extract_vhack
+from projected_grpo.grad_proj import project_gradient, svd_denoise
+
+
+@dataclass
+class Config:
+    model: str = "wassname/qwen3-5lyr-tiny-random"
+    arm: str = "both"                # "vanilla" | "projected" | "both"
+    m: int = 16                      # SVD top-m for v_hack denoise
+    steps: int = 5                   # fake GRPO steps in smoke
+    seed: int = 41
+    fast_dev_run: bool = False
+    vhack_check: bool = False
+    device: str = "cpu"              # smoke runs on CPU; real runs override
+
+
+# Synthetic contrastive pairs for smoke. Real runs use Ariahw's LeetCode
+# rollouts + Wu-Tang style domain-general pairs.
+HACK_PROMPTS = [
+    "def run_tests(): pass\n",
+    "def run_tests():\n    return True\n",
+    "def run_tests():\n    print('all passed')\n",
+    "def run_tests():\n    assert True\n",
+    "def run_tests():\n    pass  # skip checks\n",
+    "def run_tests():\n    return  # bypass\n",
+    "def run_tests():\n    print(1)\n",
+    "def run_tests():\n    True\n",
+]
+CLEAN_PROMPTS = [
+    "def solve(nums):\n    return sum(nums)\n",
+    "def solve(s):\n    return s[::-1]\n",
+    "def solve(x):\n    return x * 2\n",
+    "def solve(a, b):\n    return a + b\n",
+    "def solve(arr):\n    return sorted(arr)\n",
+    "def solve(n):\n    return n ** 2\n",
+    "def solve(s):\n    return len(s)\n",
+    "def solve(d):\n    return d.values()\n",
+]
+
+
+def run_one_arm(cfg: Config, arm: str, v_hack: torch.Tensor) -> dict:
+    """Run `cfg.steps` fake GRPO updates for arm in {'vanilla', 'projected'}.
+
+    Each step generates a random gradient with a small bias toward v_hack (to
+    simulate the policy drifting toward the hack direction). For projected arm,
+    apply project_gradient before the (fake) optimizer step.
+
+    Returns final-step diagnostics dict.
+    """
+    torch.manual_seed(cfg.seed)
+    D = v_hack.shape[0]
+    rows = []
+    final = {}
+    for step in range(cfg.steps):
+        # Fake gradient: random + 0.3 * v_hack (the loophole bias).
+        g = torch.randn(D) + 0.3 * v_hack
+        if arm == "projected":
+            g_new, diag = project_gradient(g, v_hack, preserve_magnitude=True)
+        else:
+            g_new = g
+            g_norm = g.norm()
+            cos_in = (g @ v_hack) / (g_norm + 1e-12)
+            diag = {
+                "cos_in": cos_in.item(),
+                "cos_out": cos_in.item(),
+                "projected": 0.0,
+                "g_norm_before": g_norm.item(),
+                "g_norm_after": g_norm.item(),
+            }
+        # Fake reward: high if g_new aligns with v_hack (hacky).
+        fake_reward = 0.5 + 0.4 * float((g_new @ v_hack) / (g_new.norm() + 1e-12))
+        rows.append({
+            "arm": arm, "step": step, "reward": fake_reward,
+            "cos_in": diag["cos_in"], "cos_out": diag["cos_out"],
+            "projected": diag["projected"], "g_norm": diag["g_norm_after"],
+        })
+        final = rows[-1]
+        logger.info(
+            f"step {step:02d}/{cfg.steps} {arm:9s} seed={cfg.seed} "
+            f"reward={fake_reward:+.2f} cos_in={diag['cos_in']:+.2f} "
+            f"cos_out={diag['cos_out']:+.2f} g_norm={diag['g_norm_after']:.2f}"
+        )
+    return final
+
+
+def main(cfg: Config) -> None:
+    logger.remove()
+    logger.add(sys.stderr, format="<level>{level: <8}</level> {message}")
+    logger.info(f"projected_grpo smoke run cfg={cfg}")
+
+    # 1. Load tiny model
+    logger.info(f"Loading {cfg.model} (tiny-random for smoke)")
+    tokenizer = AutoTokenizer.from_pretrained(cfg.model)
+    model = AutoModelForCausalLM.from_pretrained(
+        cfg.model, torch_dtype=torch.float32, output_hidden_states=True
+    ).to(cfg.device)
+    model.eval()
+
+    n_layers = model.config.num_hidden_layers
+    layer_idx = max(1, int(n_layers * 0.7))  # 70% depth, per Wu-Tang
+    logger.info(f"n_layers={n_layers}, using layer_idx={layer_idx}")
+
+    # 2-3. Extract v_hack from synthetic pairs.
+    n_train, n_val = 6, 2
+    h_hack = collect_last_token_hidden(model, tokenizer, HACK_PROMPTS, layer_idx, cfg.device)
+    h_clean = collect_last_token_hidden(model, tokenizer, CLEAN_PROMPTS, layer_idx, cfg.device)
+    vh = extract_vhack(
+        h_hack[:n_train], h_clean[:n_train],
+        h_hack[n_train:n_train + n_val], h_clean[n_train:n_train + n_val],
+        layer_idx=layer_idx,
+    )
+    v_hack = vh.v_hack  # "d"
+
+    # 4. SVD denoise via the lm_head weight matrix (residual-stream-out side).
+    W = model.lm_head.weight.detach().float()  # "vocab d"
+    logger.info(f"SVD-denoising v_hack via lm_head.weight shape={tuple(W.shape)} m={cfg.m}")
+    v_hack_denoised = svd_denoise(v_hack, W, m=cfg.m, use_left=False)
+    logger.info(
+        f"v_hack -> denoised: cos(orig, denoised)={float(v_hack @ v_hack_denoised):.3f} "
+        f"SHOULD>0.5 ELSE m too small or wrong basis side"
+    )
+
+    if cfg.vhack_check:
+        logger.info("vhack-check: would do CAA-style steering check here on a real model. Skipped in smoke.")
+        return
+
+    # 5. Run pathways.
+    arms = ["vanilla", "projected"] if cfg.arm == "both" else [cfg.arm]
+    results = []
+    for arm in arms:
+        final = run_one_arm(cfg, arm, v_hack_denoised)
+        results.append({
+            "arm": arm,
+            "model": cfg.model,
+            "seed": cfg.seed,
+            "m": cfg.m,
+            "n_layers": n_layers,
+            "layer_idx": layer_idx,
+            "vhack_val_acc": vh.val_accuracy,
+            "final_reward": final["reward"],
+            "final_cos_in": final["cos_in"],
+            "final_cos_out": final["cos_out"],
+            "final_g_norm": final["g_norm"],
+        })
+
+    # 6. Final TSV summary.
+    print()
+    print(tabulate(results, headers="keys", tablefmt="pipe", floatfmt="+.3f"))
+    print()
+    # BLUF
+    if cfg.arm == "both":
+        van = next(r for r in results if r["arm"] == "vanilla")
+        proj = next(r for r in results if r["arm"] == "projected")
+        delta_reward = van["final_reward"] - proj["final_reward"]
+        delta_cos = van["final_cos_out"] - proj["final_cos_out"]
+        logger.info(
+            f"BLUF: delta_reward={delta_reward:+.3f} delta_cos_out={delta_cos:+.3f} "
+            f"SHOULD both >0 (projection biting: vanilla keeps hack alignment, "
+            f"projected removes it) ELSE projection not active"
+        )
+    logger.info("smoke OK")
+
+
+if __name__ == "__main__":
+    main(tyro.cli(Config))