nb

RESEARCH_LOG.md

@@ -0,0 +1,9 @@
+# 
+
+fork_plan.md
+
+# made notebook
+
+Where does the LoRA's Δa live?
+
+In a task-derived hidden subspace that is not the lm_head readout. Held-out energy ratios (mean over LoRA-active layers 8–27, rank-8 each):

nbs/analyze_diff_v2.py

@@ -0,0 +1,389 @@
+"""One-question notebook: where does the steering signal become simple?
+
+Question:
+    Does the steering-induced activation difference on task,
+    Δa = a_{alpha=+1} - a_{alpha=-1}, concentrate in task-derived subspaces
+    more than in pretrained structural bases?
+
+Why this notebook exists:
+    The older notebook mixes several geometry questions: dW magnitude,
+    rank, linearity, and activation hooks. This one asks one falsifiable
+    question and tries to reach one of three concrete conclusions:
+
+      1. Δa concentrates in task-derived hidden subspaces
+         (TaskDiff / Suppressed / Stenographic).
+      2. Δa concentrates in pretrained readable structural bases
+         (lm_head read / logits null).
+      3. Δa is diffuse and none of the simple geometry stories survive.
+
+Method:
+    - Use the 32 sycophancy topics already in ws.data as probe prompts.
+    - Split prompts into fit/eval halves so TaskDiff concentration is not
+      tautological.
+    - Capture block outputs at the last token under alpha=+1 and alpha=-1.
+    - Build 4 candidate subspaces:
+        * TaskDiff(layer): PCA of Δa on fit prompts, per layer
+        * Suppressed(global): PCA of turnover magnitudes across layers
+        * Stenographic(layer): TaskDiff(layer) ∩ Suppressed(global)
+        * Structural controls: lm_head_read and logits_null
+    - Score concentration on eval prompts using:
+
+          energy_ratio = E[||P_V Δa||^2] / E[||Δa||^2] / (k / d)
+
+      where k is subspace rank and d is residual dimension. Ratio > 1 means
+      more concentration than a matched random k-dim subspace.
+
+Interpretation:
+    - mean_ratio(TaskDiff / Suppressed / Stenographic) >> 1 and
+      mean_ratio(structural controls) ~ 1:
+        steering signal is task-specific, not just pretrained geometry.
+    - structural controls >> task-derived:
+        steering mostly rides readable pretrained directions.
+    - everything ~ 1:
+        no simple low-dim subspace story.
+"""
+
+# %%
+from pathlib import Path
+
+import polars as pl
+import torch
+from baukit import TraceDict
+from loguru import logger
+from tabulate import tabulate
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from ws.data import SYCOPHANCY_TOPICS
+from ws.diff import load_diff
+from ws.steer import weight_steer
+
+torch.set_grad_enabled(False)
+
+
+# %% [markdown]
+# ## Config
+#
+# One question, one metric, one prompt split.
+
+# %%
+MODEL_ID = "Qwen/Qwen3-0.6B"
+W_PATH = Path("out/sycophancy/lora/w.pt")
+OUT_DIR = Path("out/sycophancy/lora")
+
+PCS = 8
+MIN_OVERLAP = 0.5
+
+PROBE_PROMPTS = [
+    f"I think {claim}. Do you agree? Answer Yes or No."
+    for claim, _q in SYCOPHANCY_TOPICS
+]
+FIT_PROMPTS = PROBE_PROMPTS[: len(PROBE_PROMPTS) // 2]
+EVAL_PROMPTS = PROBE_PROMPTS[len(PROBE_PROMPTS) // 2 :]
+
+
+# %% [markdown]
+# ## Load model and diff
+
+# %%
+w = load_diff(W_PATH)
+
+tok = AutoTokenizer.from_pretrained(MODEL_ID)
+if tok.pad_token is None:
+    tok.pad_token = tok.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+model.eval()
+
+n_layers = model.config.num_hidden_layers
+HOOKS = [f"model.layers.{i}" for i in range(n_layers)]
+lm_head_W = model.state_dict().get("lm_head.weight")
+if lm_head_W is None:
+    lm_head_W = model.state_dict()["model.embed_tokens.weight"]
+lm_head_W = lm_head_W.float().cpu()
+
+logger.info(f"loaded model={MODEL_ID} layers={n_layers} hooks={len(HOOKS)}")
+logger.info(f"loaded w with {len(w)} tensors from {W_PATH}")
+
+
+# %% [markdown]
+# ## Helpers
+
+# %%
+def orthonormalize(matrix: torch.Tensor) -> torch.Tensor:
+    if matrix.numel() == 0 or matrix.shape[1] == 0:
+        return matrix.new_zeros(matrix.shape[0], 0)
+    q, _r = torch.linalg.qr(matrix, mode="reduced")
+    return q
+
+
+def pca_basis(samples: torch.Tensor, k: int) -> torch.Tensor:
+    """samples: [n, d] -> orthonormal basis [d, k_eff]."""
+    centered = samples - samples.mean(dim=0, keepdim=True)
+    if centered.shape[0] <= 1:
+        return centered.new_zeros(centered.shape[1], 0)
+    _u, _s, vh = torch.linalg.svd(centered, full_matrices=False)
+    k_eff = min(k, vh.shape[0])
+    return vh[:k_eff].T.contiguous()
+
+
+def structural_bases(lm_head: torch.Tensor, k: int) -> dict[str, torch.Tensor]:
+    _u, _s, vh = torch.linalg.svd(lm_head, full_matrices=False)
+    return {
+        "lm_head_read": vh[:k].T.contiguous(),
+        "logits_null": vh[-k:].T.contiguous(),
+    }
+
+
+def intersect_bases(a: torch.Tensor, b: torch.Tensor, min_overlap: float) -> torch.Tensor:
+    if a.shape[1] == 0 or b.shape[1] == 0:
+        return a.new_zeros(a.shape[0], 0)
+    u, s, vh = torch.linalg.svd(a.T @ b, full_matrices=False)
+    keep = s >= min_overlap
+    if not keep.any():
+        return a.new_zeros(a.shape[0], 0)
+    va = a @ u[:, keep]
+    vb = b @ vh.T[:, keep]
+    return orthonormalize((va + vb) / 2)
+
+
+def concentration_stats(samples: torch.Tensor, basis: torch.Tensor) -> dict[str, float]:
+    d = samples.shape[1]
+    k = basis.shape[1]
+    total = samples.pow(2).sum(dim=1)
+    if k == 0:
+        return {
+            "rank": 0,
+            "energy": 0.0,
+            "null": 0.0,
+            "ratio": 0.0,
+        }
+    proj = samples @ basis
+    energy = (proj.pow(2).sum(dim=1) / (total + 1e-12)).mean().item()
+    null = k / d
+    return {
+        "rank": int(k),
+        "energy": float(energy),
+        "null": float(null),
+        "ratio": float(energy / null),
+    }
+
+
+def capture_block_outputs(prompts: list[str], alpha: float) -> torch.Tensor:
+    """Return [layers, batch, d_model] last-token block outputs."""
+    enc = tok(
+        prompts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=128,
+    ).to(model.device)
+    seq_idx = enc.attention_mask.sum(dim=-1) - 1
+
+    ctx = weight_steer(model, w, alpha) if alpha != 0 else torch.no_grad()
+    with ctx:
+        with TraceDict(model, HOOKS, retain_output=True) as ret:
+            _ = model(**enc)
+
+    rows = []
+    for hook in HOOKS:
+        x = ret[hook].output
+        if isinstance(x, tuple):
+            x = x[0]
+        batch, _seq, d_model = x.shape
+        gather_idx = seq_idx.view(batch, 1, 1).expand(batch, 1, d_model)
+        last_tok = x.gather(1, gather_idx).squeeze(1).float().cpu()
+        rows.append(last_tok)
+    return torch.stack(rows, dim=0)
+
+
+def suppressed_features(acts: torch.Tensor) -> torch.Tensor:
+    """acts: [layers, batch, d] -> turnover features [batch, d]."""
+    mag = acts.abs().permute(1, 0, 2)
+    delta = mag[:, 1:] - mag[:, :-1]
+    increases = torch.relu(delta).sum(dim=1)
+    decreases = torch.relu(-delta).sum(dim=1)
+    return torch.minimum(increases, decreases)
+
+
+# %% [markdown]
+# ## Capture fit/eval activations under alpha=+1 and alpha=-1
+
+# %%
+pos_fit = capture_block_outputs(FIT_PROMPTS, alpha=+1.0)
+neg_fit = capture_block_outputs(FIT_PROMPTS, alpha=-1.0)
+pos_eval = capture_block_outputs(EVAL_PROMPTS, alpha=+1.0)
+neg_eval = capture_block_outputs(EVAL_PROMPTS, alpha=-1.0)
+
+delta_fit = pos_fit - neg_fit
+delta_eval = pos_eval - neg_eval
+
+logger.info(
+    "captured fit/eval activations: fit={} eval={} shape={}",
+    len(FIT_PROMPTS),
+    len(EVAL_PROMPTS),
+    tuple(delta_fit.shape),
+)
+
+
+# %% [markdown]
+# ## Fit candidate subspaces on the fit split
+#
+# We compare task-derived candidates against structural controls.
+
+# %%
+taskdiff_bases = [pca_basis(delta_fit[layer], PCS) for layer in range(n_layers)]
+
+suppressed_fit = 0.5 * (suppressed_features(pos_fit) + suppressed_features(neg_fit))
+suppressed_basis = pca_basis(suppressed_fit, PCS)
+
+structural = structural_bases(lm_head_W, PCS)
+stenographic_bases = [
+    intersect_bases(taskdiff_bases[layer], suppressed_basis, min_overlap=MIN_OVERLAP)
+    for layer in range(n_layers)
+]
+
+logger.info(
+    "basis ranks: suppressed={} lm_head_read={} logits_null={}",
+    suppressed_basis.shape[1],
+    structural["lm_head_read"].shape[1],
+    structural["logits_null"].shape[1],
+)
+
+
+# %% [markdown]
+# ## Score concentration on held-out prompts
+#
+# Main metric:
+#
+#   energy_ratio = E[||P_V Δa||²] / E[||Δa||²] / (k / d)
+#
+# Ratio > 1 means more concentration than a matched random k-dim subspace.
+
+# %%
+rows = []
+for layer in range(n_layers):
+    x = delta_eval[layer]
+    candidates = {
+        "taskdiff": taskdiff_bases[layer],
+        "suppressed": suppressed_basis,
+        "stenographic": stenographic_bases[layer],
+        "lm_head_read": structural["lm_head_read"],
+        "logits_null": structural["logits_null"],
+    }
+    for name, basis in candidates.items():
+        stats = concentration_stats(x, basis)
+        rows.append({
+            "layer": layer,
+            "subspace": name,
+            **stats,
+        })
+
+df = pl.DataFrame(rows)
+
+summary = (
+    df.group_by("subspace")
+    .agg(
+        pl.col("ratio").mean().alias("mean_ratio"),
+        pl.col("ratio").max().alias("max_ratio"),
+        pl.col("layer").sort_by("ratio").last().alias("peak_layer"),
+        pl.col("energy").mean().alias("mean_energy"),
+        pl.col("rank").mean().alias("mean_rank"),
+    )
+    .sort("mean_ratio", descending=True)
+)
+
+print("\nconcentration summary on held-out prompts")
+print(
+    "SHOULD: if task-derived subspaces are real, taskdiff / suppressed / stenographic "
+    "have mean_ratio >> 1 and beat structural controls. ELSE: if lm_head_read wins, "
+    "the signal is already readable; if everything ~= 1, the geometry story is weak."
+)
+print(
+    tabulate(
+        summary.to_pandas(),
+        tablefmt="tsv",
+        headers="keys",
+        floatfmt="+.3f",
+        showindex=False,
+    )
+)
+
+print("\nper-layer table")
+print(
+    tabulate(
+        df.sort(["subspace", "layer"]).to_pandas(),
+        tablefmt="tsv",
+        headers="keys",
+        floatfmt="+.3f",
+        showindex=False,
+    )
+)
+
+
+# %% [markdown]
+# ## Decision rule
+#
+# Read the summary table as a model selection result:
+#
+# - task-derived >> structural:
+#     the steering signal is task-specific and hidden / dynamic.
+# - structural >> task-derived:
+#     the steering mostly rides pretrained readable axes.
+# - all near 1:
+#     the signal is diffuse and this basis story is probably wrong.
+#
+# If task-derived wins, *then* it becomes worth doing stage-2 mechanism tests
+# like rotation-vs-gain fits or stage-3 intervention tests like LEACE.
+
+# %%
+df.write_csv(OUT_DIR / "analyze_diff_v2_concentration_per_layer.csv")
+summary.write_csv(OUT_DIR / "analyze_diff_v2_concentration_summary.csv")
+logger.info("saved v2 concentration tables to {}", OUT_DIR)
+
+
+# %% [markdown]
+# ## Stage-1.5: principal angles between TaskDiff(layer) and lm_head_read
+#
+# Concentration says TaskDiff captures most Delta-a energy and lm_head_read does
+# not. This is sufficient evidence that the signal is not the readout direction,
+# but principal angles make the geometric relationship explicit.
+#
+# For two rank-k orthonormal bases A, B in R^d, the principal cosines are the
+# singular values of A.T @ B. All near 1 means the subspaces nearly coincide;
+# all near 0 means they are orthogonal.
+
+# %%
+angle_rows = []
+lm_basis = structural["lm_head_read"]
+for layer in range(n_layers):
+    A = taskdiff_bases[layer]
+    if A.shape[1] == 0:
+        continue
+    cos_angles = torch.linalg.svdvals(A.T @ lm_basis).clamp(0, 1)
+    angle_rows.append({
+        "layer": layer,
+        "max_cos": float(cos_angles.max()),
+        "mean_cos": float(cos_angles.mean()),
+        "min_cos": float(cos_angles.min()),
+    })
+
+angle_df = pl.DataFrame(angle_rows)
+print("\nprincipal cosines between TaskDiff(layer) and lm_head_read")
+print(
+    "SHOULD: if TaskDiff is largely orthogonal to readout, mean_cos << 1 and "
+    "max_cos < 0.7 in active layers (>=8). ELSE TaskDiff is a relabel of readout."
+)
+print(
+    tabulate(
+        angle_df.to_pandas(),
+        tablefmt="tsv",
+        headers="keys",
+        floatfmt="+.3f",
+        showindex=False,
+    )
+)
+angle_df.write_csv(OUT_DIR / "analyze_diff_v2_taskdiff_vs_lmhead_angles.csv")