World Alignment: unpaired cross-modal correspondence by relational identifiability

Method: scene states are sets of part states; relation fields are built
within each modality and are invariant to how each side labels its own
features; the cross-modal bridge is a coupling searched under an energy
that is a closed-form functional of one matrix; solving is spectral
initialisation followed by exact local refinement.

Evidence: in a procedurally generated closed world, blind recovery of a
hidden image-caption correspondence reaches 95.3% at 256 scenes against
0.39% chance, and the recovered pairs transfer to 200 held-out scenes at
93.0% exact retrieval with random-pair and shuffled-image controls at or
near chance. Cross-modal value correspondence is derived from disjoint
corpora rather than declared. On Visual Genome the field correlation
reaches 0.656 against the 0.9 that polynomial recovery needs, with the
deficit attributed away from segmentation and discretisation.

Protocol: no image-text pair enters any objective, optimiser,
initialisation, or model selection; hidden pairs score orderings only.

Co-Authored-By: Claude <noreply@anthropic.com>

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()