diff options
Diffstat (limited to 'worldalign/vg_view_probe.py')
| -rw-r--r-- | worldalign/vg_view_probe.py | 292 |
1 files changed, 292 insertions, 0 deletions
diff --git a/worldalign/vg_view_probe.py b/worldalign/vg_view_probe.py new file mode 100644 index 0000000..322644f --- /dev/null +++ b/worldalign/vg_view_probe.py @@ -0,0 +1,292 @@ +"""View-level shared-structure probe for Visual Genome nodes. + +The node-level manifold gate shows that pooled node states expose only a +low-dimensional coarse shared subspace, which cannot identify individual +assignments. This probe measures whether additional cross-modal signal +exists inside nodes, at the level of the sixteen region views, where both +modalities observe the same sixteen regions of the same scene. + +The per-node view permutation is replayed from the preparation seed and the +private image IDs, verified exactly against the released text bundles, and +used only for evaluation. No cross-modal map is trained. + +Three quantities are reported: + +1. within-node relational alignment: correlation between the visual and the + truth-aligned text view-relation fields, against a shuffled-view control; +2. coarse-frame view matching: population-context profiles make single + views cross-modally comparable through the node-level frame; assignment + accuracy is chance-normalized (16 candidates per node); +3. a within-node relational-only floor with no population context. +""" + +from __future__ import annotations + +import argparse +import json +import random +from pathlib import Path + +import torch +import torch.nn.functional as F +from scipy.optimize import linear_sum_assignment + +from .common import write_json +from .vg_prepare import load_config + + +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( + "--context-nodes", + type=int, + default=1000, + help="Node-level frame size for population-context profiles.", + ) + parser.add_argument("--nodes", type=int, default=0, help="0 keeps all nodes.") + parser.add_argument( + "--skip-floor", + action="store_true", + help="Skip the slow within-node relational-only descent floor.", + ) + parser.add_argument( + "--output", default="artifacts/manifold_gate/vg_view_probe.json" + ) + return parser.parse_args() + + +def replay_view_permutations(args: argparse.Namespace) -> dict: + """Recover text-view -> vision-view mappings from the preparation RNG. + + ``random.shuffle`` consumes the generator identically for equal-length + lists, so shuffling index lists replays the region subsample and the + phrase shuffle exactly. Every node is verified against the released + phrase bundle before use. + """ + truth = [ + json.loads(line) + for line in Path(args.vg_dir, "ground_truth.private.jsonl").read_text( + encoding="utf-8" + ).splitlines() + if line.strip() + ] + text_nodes = { + record["node_id"]: record + for line in Path(args.vg_dir, "text_nodes.jsonl").read_text( + encoding="utf-8" + ).splitlines() + if line.strip() + for record in [json.loads(line)] + } + region_ds = load_config("region_descriptions_v1.2.0", args.cache_dir) + image_rows = { + int(image_id): row for row, image_id in enumerate(region_ds["image_id"]) + } + mappings: dict[str, dict] = {} + mismatched = 0 + for record in truth: + image_id = int(record["source_image_id"]) + row = region_ds[image_rows[image_id]] + regions = list(row["regions"]) + node_rng = random.Random(args.seed ^ image_id) + order = list(range(len(regions))) + node_rng.shuffle(order) + selected = order[: args.views] + phrase_perm = list(range(args.views)) + node_rng.shuffle(phrase_perm) + phrases_view_order = [ + str(regions[index]["phrase"]).strip() for index in selected + ] + reconstructed = [phrases_view_order[j] for j in phrase_perm] + released = text_nodes[record["text_node_id"]]["region_closed"] + if reconstructed != released: + mismatched += 1 + continue + mappings[record["vision_node_id"]] = { + "text_node_id": record["text_node_id"], + # text view j describes vision view phrase_perm[j] + "text_to_vision": phrase_perm, + } + if mismatched: + print(json.dumps({"replay_mismatched_nodes": mismatched})) + return mappings + + +def offdiag(matrix: torch.Tensor) -> torch.Tensor: + mask = ~torch.eye(len(matrix), dtype=torch.bool) + return matrix[mask] + + +def spearman(x: torch.Tensor, y: torch.Tensor) -> float: + ranks = torch.stack( + [x.argsort().argsort().double(), y.argsort().argsort().double()] + ) + return float(torch.corrcoef(ranks)[0, 1]) + + +def main() -> None: + args = parse_args() + mappings = replay_view_permutations(args) + + vision = torch.load( + Path(args.vg_dir, "vision_features.pt"), map_location="cpu", weights_only=False + ) + text = torch.load( + Path(args.vg_dir, "text_features.pt"), map_location="cpu", weights_only=False + ) + vision_index = {node: i for i, node in enumerate(vision["node_ids"])} + text_index = {node: i for i, node in enumerate(text["node_ids"])} + visual_views = F.normalize(vision["region_features"].double(), dim=-1) + text_views = F.normalize(text["region_features"].double(), dim=-1) + + node_ids = sorted(mappings) + if args.nodes: + node_ids = node_ids[: args.nodes] + + # Node-level coarse frame: view-mean states over a fixed context set, + # truth-aligned so both profile axes index the same underlying scenes. + # The frame is evaluation-side scaffolding for an upper bound; a blind + # system would substitute its recovered coarse alignment here. + context_ids = node_ids[: args.context_nodes] + visual_frame = F.normalize( + torch.stack( + [visual_views[vision_index[node]].mean(0) for node in context_ids] + ), + dim=-1, + ) + text_frame = F.normalize( + torch.stack( + [ + text_views[text_index[mappings[node]["text_node_id"]]].mean(0) + for node in context_ids + ] + ), + dim=-1, + ) + + generator = torch.Generator().manual_seed(args.seed) + relation_matched: list[float] = [] + relation_shuffled: list[float] = [] + profile_hits = 0 + profile_top3 = 0 + profile_total = 0 + floor_hits = 0 + floor_total = 0 + matched_profile_corr: list[float] = [] + mismatched_profile_corr: list[float] = [] + + for node in node_ids: + mapping = mappings[node] + v = visual_views[vision_index[node]] + t = text_views[text_index[mapping["text_node_id"]]] + 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)) + + relation_visual = v @ v.T + relation_text = t @ t.T + aligned = relation_text[vision_to_text][:, vision_to_text] + relation_matched.append(spearman(offdiag(relation_visual), offdiag(aligned))) + shuffle = torch.randperm(len(v), generator=generator) + shuffled = relation_text[shuffle][:, shuffle] + relation_shuffled.append(spearman(offdiag(relation_visual), offdiag(shuffled))) + + # Population-context profiles through the coarse frame. + profile_visual = v @ visual_frame.T + profile_text = t @ text_frame.T + rank_visual = profile_visual.argsort(-1).argsort(-1).double() + rank_text = profile_text.argsort(-1).argsort(-1).double() + rank_visual = rank_visual - rank_visual.mean(-1, keepdim=True) + rank_text = rank_text - rank_text.mean(-1, keepdim=True) + rank_visual = F.normalize(rank_visual, dim=-1) + rank_text = F.normalize(rank_text, dim=-1) + similarity = rank_visual @ rank_text.T # [vision view, text view] + rows, cols = linear_sum_assignment(-similarity.numpy()) + truth_cols = vision_to_text.numpy() + profile_hits += int((cols == truth_cols).sum()) + ranks = ( + similarity + >= similarity.gather(1, vision_to_text[:, None]) + ).sum(-1) + profile_top3 += int((ranks <= 3).sum()) + profile_total += len(rows) + matched_profile_corr.extend( + similarity.gather(1, vision_to_text[:, None]).squeeze(1).tolist() + ) + mismatched_profile_corr.extend( + similarity[~torch.eye(len(v), dtype=torch.bool)][:64].tolist() + ) + + # Relational-only floor: 2-swap descent on within-node fields from a + # random start, no population context. + if args.skip_floor: + continue + state = torch.randperm(len(v), generator=generator) + for _ in range(200): + improved = False + current = relation_text[state][:, state] + base = float(((relation_visual - current) ** 2).mean()) + for p in range(len(v)): + for q in range(p + 1, len(v)): + trial = state.clone() + trial[[p, q]] = trial[[q, p]] + candidate = relation_text[trial][:, trial] + if float(((relation_visual - candidate) ** 2).mean()) < base - 1e-12: + state = trial + improved = True + break + if improved: + break + if not improved: + break + floor_hits += int((state == vision_to_text).sum()) + floor_total += len(v) + + matched = torch.tensor(relation_matched) + shuffled = torch.tensor(relation_shuffled) + matched_corr = torch.tensor(matched_profile_corr) + mismatched_corr = torch.tensor(mismatched_profile_corr) + report = { + "protocol": ( + "View permutations are replayed from the preparation seed and " + "verified against released phrase bundles; they are used only " + "for evaluation. The coarse frame uses truth-aligned node " + "states, so profile matching is an upper bound on what a " + "recovered node alignment would enable." + ), + "nodes": len(node_ids), + "verified_fraction": len(mappings) / 5000.0, + "within_node_relation": { + "matched_spearman_mean": float(matched.mean()), + "matched_spearman_std": float(matched.std()), + "shuffled_spearman_mean": float(shuffled.mean()), + "shuffled_spearman_std": float(shuffled.std()), + "gap_z_across_nodes": float( + (matched.mean() - shuffled.mean()) + / (matched - shuffled).std().clamp_min(1e-12) + * len(matched) ** 0.5 + ), + }, + "coarse_frame_view_matching": { + "context_nodes": len(context_ids), + "accuracy": profile_hits / max(profile_total, 1), + "top3_accuracy": profile_top3 / max(profile_total, 1), + "chance": 1.0 / args.views, + "matched_profile_corr_mean": float(matched_corr.mean()), + "mismatched_profile_corr_mean": float(mismatched_corr.mean()), + }, + "within_node_relational_floor": { + "accuracy": floor_hits / max(floor_total, 1), + "chance": 1.0 / args.views, + }, + } + write_json(args.output, report) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() |