summaryrefslogtreecommitdiff
path: root/worldalign/synth_probes.py
diff options
context:
space:
mode:
authorYuren Hao <blackhao0426@gmail.com>2026-08-01 14:10:03 -0500
committerYuren Hao <blackhao0426@gmail.com>2026-08-01 14:10:03 -0500
commita62cf4d2a99b4a7985c61b2a7feb92a82a8218b7 (patch)
treeee2248078db7edf3812a07f195afa3d9bd6f10c6 /worldalign/synth_probes.py
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>
Diffstat (limited to 'worldalign/synth_probes.py')
-rw-r--r--worldalign/synth_probes.py129
1 files changed, 129 insertions, 0 deletions
diff --git a/worldalign/synth_probes.py b/worldalign/synth_probes.py
new file mode 100644
index 0000000..180b96a
--- /dev/null
+++ b/worldalign/synth_probes.py
@@ -0,0 +1,129 @@
+"""Ground-truth factor probes for synthetic-world representations.
+
+Linear probes from a representation to the discrete scene factors measure
+which world variables survive the encoder and its pooling. Probes are
+fitted per modality on that modality's own training split and evaluated
+on held-out scenes; scene truth is generator metadata, so this is a
+within-modality diagnostic, not a cross-modal signal.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .common import read_json, write_json
+from .synth_world import COLORS, SHAPES
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--data-dir", default="artifacts/synth_v0")
+ parser.add_argument("--features", required=True)
+ parser.add_argument(
+ "--side", choices=["vision", "text"], required=True,
+ help="Selects the training split whose scenes fit the probes.",
+ )
+ parser.add_argument("--ridge", type=float, default=1.0)
+ parser.add_argument("--output", required=True)
+ return parser.parse_args()
+
+
+def scene_targets(scene: dict) -> dict[str, np.ndarray | float]:
+ color_presence = np.zeros(len(COLORS))
+ shape_presence = np.zeros(len(SHAPES))
+ for group in scene["groups"]:
+ color_presence[list(COLORS).index(group["color"])] = 1.0
+ shape_presence[SHAPES.index(group["shape"])] = 1.0
+ return {
+ "color_presence": color_presence,
+ "shape_presence": shape_presence,
+ "group_count": float(len(scene["groups"])),
+ "object_total": float(sum(g["count"] for g in scene["groups"])),
+ "relation_count": float(len(scene["relations"])),
+ }
+
+
+def ridge_fit(
+ x: np.ndarray, y: np.ndarray, ridge: float
+) -> tuple[np.ndarray, np.ndarray]:
+ x = np.concatenate([x, np.ones((len(x), 1))], axis=1)
+ gram = x.T @ x + ridge * np.eye(x.shape[1])
+ weights = np.linalg.solve(gram, x.T @ y)
+ return weights, x @ weights
+
+
+def evaluate(
+ features_train: np.ndarray,
+ features_test: np.ndarray,
+ train_targets: np.ndarray,
+ test_targets: np.ndarray,
+ ridge: float,
+ binary: bool,
+) -> dict:
+ weights, _ = ridge_fit(features_train, train_targets, ridge)
+ prediction = (
+ np.concatenate([features_test, np.ones((len(features_test), 1))], axis=1)
+ @ weights
+ )
+ if binary:
+ accuracy = float(((prediction > 0.5) == (test_targets > 0.5)).mean())
+ balanced = []
+ for column in range(test_targets.shape[1]):
+ truth = test_targets[:, column] > 0.5
+ if truth.any() and (~truth).any():
+ hit = (prediction[:, column] > 0.5) == truth
+ balanced.append(
+ (hit[truth].mean() + hit[~truth].mean()) / 2.0
+ )
+ return {
+ "accuracy": accuracy,
+ "balanced_accuracy": float(np.mean(balanced)),
+ }
+ residual = prediction[:, 0] - test_targets[:, 0]
+ variance = test_targets[:, 0].var()
+ return {
+ "r2": float(1.0 - residual.var() / max(variance, 1e-9)),
+ "mae": float(np.abs(residual).mean()),
+ }
+
+
+def main() -> None:
+ args = parse_args()
+ manifest = read_json(Path(args.data_dir, "manifest.json"))
+ scenes = read_json(Path(args.data_dir, "scenes.private.json"))["scenes"]
+ state = torch.load(args.features, map_location="cpu", weights_only=False)
+ lookup = {int(row): i for i, row in enumerate(state["rows"])}
+ features = state["features"].float().numpy()
+
+ split_key = "vision_only_train" if args.side == "vision" else "text_only_train"
+ train_rows = [row for row in manifest[split_key] if row in lookup][:8000]
+ test_rows = [row for row in manifest["test"] if row in lookup]
+
+ x_train = features[[lookup[row] for row in train_rows]]
+ x_test = features[[lookup[row] for row in test_rows]]
+ report = {"features": args.features, "side": args.side, "targets": {}}
+ for name in ("color_presence", "shape_presence"):
+ y_train = np.stack([scene_targets(scenes[row])[name] for row in train_rows])
+ y_test = np.stack([scene_targets(scenes[row])[name] for row in test_rows])
+ report["targets"][name] = evaluate(
+ x_train, x_test, y_train, y_test, args.ridge, binary=True
+ )
+ for name in ("group_count", "object_total", "relation_count"):
+ y_train = np.array(
+ [[scene_targets(scenes[row])[name]] for row in train_rows]
+ )
+ y_test = np.array([[scene_targets(scenes[row])[name]] for row in test_rows])
+ report["targets"][name] = evaluate(
+ x_train, x_test, y_train, y_test, args.ridge, binary=False
+ )
+ write_json(args.output, report)
+ print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+ main()