diff options
Diffstat (limited to 'worldalign/vg_attention_probe.py')
| -rw-r--r-- | worldalign/vg_attention_probe.py | 170 |
1 files changed, 170 insertions, 0 deletions
diff --git a/worldalign/vg_attention_probe.py b/worldalign/vg_attention_probe.py new file mode 100644 index 0000000..3dac7a3 --- /dev/null +++ b/worldalign/vg_attention_probe.py @@ -0,0 +1,170 @@ +"""R3 battery evaluation: do model-internal relations align across modalities? + +Compares, at the replayed true view correspondence, the cross-modal +alignment of relation fields read from inside the frozen models +(cross-phrase attention per layer group, in-context state cosines) against +the isolated-encoding cosine baseline (matched rho 0.128, z 46.6). +Spearman is rank-based, so monotone attention transforms are immaterial. +View truth is evaluation-only. +""" + +from __future__ import annotations + +import argparse +import json + +import torch +import torch.nn.functional as F + +from .common import write_json +from .vg_view_probe import offdiag, replay_view_permutations, spearman + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--vg-dir", default="artifacts/vg_5k") + parser.add_argument("--cache-dir", default="/tmp/yurenh2-worldalign-vg-hf") + parser.add_argument("--seed", type=int, default=20260728) + parser.add_argument("--views", type=int, default=16) + parser.add_argument("--text-attn", default="artifacts/manifold_gate/attn_text_full.pt") + parser.add_argument( + "--vision-attn", default="artifacts/manifold_gate/attn_vision_full.pt" + ) + parser.add_argument("--nodes", type=int, default=0) + parser.add_argument( + "--output", default="artifacts/manifold_gate/attention_probe.json" + ) + return parser.parse_args() + + +def channel_fields(state: dict, side: str) -> dict[str, torch.Tensor]: + """Named [N, V, V] relation fields for one side.""" + attention = state["attention_channels"].double() + fields = { + f"{side}_attn_g{group}": attention[:, group] + for group in range(attention.shape[1]) + } + context = F.normalize(state["context_states"].double(), dim=-1) + fields[f"{side}_ctx_cos"] = context @ context.transpose(-2, -1) + return fields + + +def main() -> None: + args = parse_args() + mappings = replay_view_permutations(args) + + text_state = torch.load(args.text_attn, map_location="cpu", weights_only=False) + vision_state = torch.load(args.vision_attn, map_location="cpu", weights_only=False) + text_index = {node: i for i, node in enumerate(text_state["node_ids"])} + vision_index = {node: i for i, node in enumerate(vision_state["node_ids"])} + + baseline_vision = torch.load( + f"{args.vg_dir}/vision_features.pt", map_location="cpu", weights_only=False + ) + baseline_text = torch.load( + f"{args.vg_dir}/text_features.pt", map_location="cpu", weights_only=False + ) + baseline_vision_index = { + node: i for i, node in enumerate(baseline_vision["node_ids"]) + } + baseline_text_index = {node: i for i, node in enumerate(baseline_text["node_ids"])} + + text_fields = channel_fields(text_state, "text") + vision_fields = channel_fields(vision_state, "vision") + + node_ids = sorted(mappings) + if args.nodes: + node_ids = node_ids[: args.nodes] + + pairs = [ + (t_name, v_name) for t_name in text_fields for v_name in vision_fields + ] + matched: dict[tuple[str, str], list[float]] = {pair: [] for pair in pairs} + shuffled: dict[tuple[str, str], list[float]] = {pair: [] for pair in pairs} + baseline_matched: list[float] = [] + baseline_shuffled: list[float] = [] + generator = torch.Generator().manual_seed(args.seed) + + for node in node_ids: + mapping = mappings[node] + text_row = text_index[mapping["text_node_id"]] + vision_row = vision_index[node] + text_to_vision = torch.tensor(mapping["text_to_vision"]) + vision_to_text = torch.empty_like(text_to_vision) + vision_to_text[text_to_vision] = torch.arange(len(text_to_vision)) + shuffle = torch.randperm(args.views, generator=generator) + + for t_name, v_name in pairs: + text_field = text_fields[t_name][text_row] + aligned = text_field[vision_to_text][:, vision_to_text] + visual_field = vision_fields[v_name][vision_row] + matched[(t_name, v_name)].append( + spearman(offdiag(visual_field), offdiag(aligned)) + ) + scrambled = text_field[shuffle][:, shuffle] + shuffled[(t_name, v_name)].append( + spearman(offdiag(visual_field), offdiag(scrambled)) + ) + + crops = F.normalize( + baseline_vision["region_features"][ + baseline_vision_index[node] + ].double(), + dim=-1, + ) + phrases = F.normalize( + baseline_text["region_features"][ + baseline_text_index[mapping["text_node_id"]] + ].double(), + dim=-1, + ) + crop_field = crops @ crops.T + phrase_field = (phrases @ phrases.T)[vision_to_text][:, vision_to_text] + baseline_matched.append(spearman(offdiag(crop_field), offdiag(phrase_field))) + scrambled = (phrases @ phrases.T)[shuffle][:, shuffle] + baseline_shuffled.append(spearman(offdiag(crop_field), offdiag(scrambled))) + + def summarize(matched_values: list[float], shuffled_values: list[float]) -> dict: + m = torch.tensor(matched_values) + s = torch.tensor(shuffled_values) + return { + "matched_mean": float(m.mean()), + "shuffled_mean": float(s.mean()), + "gap_z": float( + (m.mean() - s.mean()) + / (m - s).std().clamp_min(1e-12) + * len(m) ** 0.5 + ), + } + + report = { + "protocol": ( + "Replayed view truth, evaluation-only. Each cell is the " + "within-node cross-modal Spearman of relation fields at the " + "true view correspondence, against a shuffled-view control." + ), + "nodes": len(node_ids), + "baseline_isolated_cosine": summarize(baseline_matched, baseline_shuffled), + "channels": { + f"{t} x {v}": summarize(matched[(t, v)], shuffled[(t, v)]) + for t, v in pairs + }, + } + write_json(args.output, report) + best = max( + report["channels"].items(), key=lambda item: item[1]["matched_mean"] + ) + print( + json.dumps( + { + "baseline": report["baseline_isolated_cosine"], + "best_channel": {best[0]: best[1]}, + "nodes": len(node_ids), + } + ) + ) + print(f"Wrote {args.output}") + + +if __name__ == "__main__": + main() |