| """ |
| Cross-relation LoRA-B alignment: are the steering directions of different |
| scene→object hallucination adapters the SAME axis or relation-specific? |
| |
| For each relation's adapter, take the per-layer LoRA-B direction of a |
| residual-writing module (default mlp.down_proj, B (d_model,r) reduced over rank → |
| unit v̂_l). Then per layer compute pairwise cosine similarity between relations: |
| bathroom_toilet ↔ kitchen_oven ↔ livingroom_tv. |
| |
| Random-direction baseline ≈ 1/sqrt(d_model) (≈0.016 for d=4096): |cos| above that |
| = real alignment. High |cos| across relations ⇒ a shared "scene→object / hallucinated- |
| object" axis; near-zero ⇒ each relation has its own direction. |
| |
| Saves: per-layer pairwise cosine JSON + line plot. |
| """ |
| import argparse, json, os, itertools |
| import numpy as np, torch as t |
| import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt |
|
|
| from mechanistic_interp.alternative_steering import load_lora_b_dirs |
|
|
| REL = { |
| "bathroom_toilet": "/data/caotue/multilayer-sae/adv_gen_outputs/run_bathroom_toilet_v2/lora_adapter", |
| "kitchen_oven": "/data/caotue/multilayer-sae/adv_gen_outputs/run_kitchen_oven_v2/lora_adapter", |
| "livingroom_tv": "/data/caotue/multilayer-sae/adv_gen_outputs/run_livingroom_tv_v2/step_450", |
| } |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--model_name", default="llava-hf/llava-1.5-7b-hf") |
| ap.add_argument("--lora_module", default="down_proj", choices=["down_proj", "o_proj"]) |
| ap.add_argument("--b_reduce", default="mean", choices=["mean", "svd"]) |
| ap.add_argument("--n_layers", type=int, default=32) |
| ap.add_argument("--out_prefix", default="mechanistic_interp/graph/lora_relations_sim") |
| args = ap.parse_args() |
|
|
| from transformers import AutoConfig |
| d_model = AutoConfig.from_pretrained(args.model_name).text_config.hidden_size |
| layers = list(range(args.n_layers)) |
| rng_base = 1.0 / np.sqrt(d_model) |
|
|
| dirs = {r: load_lora_b_dirs(p, layers, args.lora_module, args.b_reduce, d_model, "cpu", t.float32) |
| for r, p in REL.items()} |
| rels = list(REL) |
| pairs = list(itertools.combinations(rels, 2)) |
|
|
| cos = {f"{a}|{b}": [] for a, b in pairs} |
| for l in layers: |
| for a, b in pairs: |
| va, vb = dirs[a][l], dirs[b][l] |
| cos[f"{a}|{b}"].append(float(t.dot(va, vb) / (va.norm() * vb.norm() + 1e-8))) |
|
|
| print(f"random-direction baseline |cos| ~ {rng_base:.4f} (d={d_model})\n") |
| print("pair mean|cos| max|cos| signed-mean") |
| for k, v in cos.items(): |
| v = np.array(v) |
| print(f" {k:32s} {np.abs(v).mean():.3f} {np.abs(v).max():.3f} {v.mean():+.3f}") |
|
|
| os.makedirs(os.path.dirname(args.out_prefix), exist_ok=True) |
| json.dump({"layers": layers, "rng_baseline": rng_base, "module": args.lora_module, |
| "b_reduce": args.b_reduce, "cosine": cos}, open(f"{args.out_prefix}.json", "w"), indent=2) |
|
|
| fig, axes = plt.subplots(1, 2, figsize=(15, 5)) |
| for k, v in cos.items(): |
| axes[0].plot(layers, v, "o-", ms=3, label=k) |
| axes[1].plot(layers, np.abs(v), "o-", ms=3, label=k) |
| axes[0].axhline(0, color="gray", ls=":"); axes[0].set_ylabel("signed cosine") |
| axes[1].axhline(rng_base, color="red", ls="--", lw=0.8, label=f"random {rng_base:.3f}") |
| axes[1].set_ylabel("|cosine| (alignment)") |
| for ax, ttl in ((axes[0], "signed"), (axes[1], "|cos|")): |
| ax.set_xlabel("layer"); ax.set_title(f"LoRA-B cross-relation cosine ({ttl})"); ax.grid(alpha=0.3); ax.legend(fontsize=8) |
| fig.suptitle(f"LoRA-B direction alignment across relations ({args.lora_module}/{args.b_reduce})") |
| fig.tight_layout(); fig.savefig(f"{args.out_prefix}.png", dpi=150, bbox_inches="tight") |
| print(f"\nsaved {args.out_prefix}.png / .json") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
