From a62cf4d2a99b4a7985c61b2a7feb92a82a8218b7 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Sat, 1 Aug 2026 14:10:03 -0500 Subject: 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 --- worldalign/synth_set_battery.py | 283 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 283 insertions(+) create mode 100644 worldalign/synth_set_battery.py (limited to 'worldalign/synth_set_battery.py') diff --git a/worldalign/synth_set_battery.py b/worldalign/synth_set_battery.py new file mode 100644 index 0000000..9ed66ec --- /dev/null +++ b/worldalign/synth_set_battery.py @@ -0,0 +1,283 @@ +"""Set-kernel relation fields for the synthetic world: the structure battery. + +The pooled readout of a set representation destroys it. Here scene states +stay sets -- vision: slot vectors with alpha masses; text: per-group +phrase states parsed from the caption's enumeration sentence and encoded +individually -- and scene-to-scene relations are computed within each +modality as set-matching similarities. The cross-modal gate then runs on +these set-kernel relation fields exactly as on any relation channel. + +Phrase parsing reads only released captions; slot sets read only renders. +Hidden pairs score orderings, as always. +""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + +import torch +import torch.nn.functional as F +from scipy.optimize import linear_sum_assignment +from tqdm import tqdm + +from .common import batch_indices, read_json, seed_everything, write_json +from .manifold_gate import standardize_relation +from .ricci_control import run_gates +from .synth_slots import SlotAutoencoder +from .synth_towers import TextTower, load_image, tokenize + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", default="artifacts/synth_v0") + parser.add_argument( + "--slots", default="artifacts/synth_v0/vision_slots.pt" + ) + parser.add_argument("--text-tower", default="artifacts/synth_v0/text_tower.pt") + parser.add_argument("--split", choices=["val", "test"], default="test") + parser.add_argument("--samples", type=int, default=512) + parser.add_argument("--mass-floor", type=float, default=0.02) + parser.add_argument( + "--vision-mode", choices=["slot_vectors", "sprites"], default="sprites" + ) + parser.add_argument( + "--slot-tower", default="artifacts/synth_v0/slot_tower.pt" + ) + parser.add_argument("--sprite-window", type=int, default=48) + parser.add_argument("--random-perms", type=int, default=300) + parser.add_argument("--descent-restarts", type=int, default=5) + parser.add_argument("--descent-max-steps", type=int, default=200000) + parser.add_argument( + "--descent-objective", default="mse", choices=["mse", "m30_total"] + ) + parser.add_argument("--descent-verify-top", type=int, default=64) + parser.add_argument("--device", default="cuda:3") + parser.add_argument("--seed", type=int, default=20260731) + parser.add_argument( + "--output", default="artifacts/synth_v0/set_battery_gate.json" + ) + return parser.parse_args() + + +def parse_group_phrases(caption: str) -> list[str]: + """Group phrases from the enumeration sentence of a released caption.""" + first = caption.split(".")[0] + for opener in ("there are ", "the picture shows ", "you can see "): + if first.startswith(opener): + first = first[len(opener):] + break + first = first.replace(" and ", ", ") + return [phrase.strip() for phrase in first.split(",") if phrase.strip()] + + +@torch.inference_mode() +def text_group_sets( + rows: list[int], captions: list[list[str]], args: argparse.Namespace +) -> list[torch.Tensor]: + state = torch.load(args.text_tower, map_location="cpu", weights_only=False) + saved = state["args"] + vocab = state["vocab"] + model = TextTower( + len(vocab), + saved["text_dim"], + saved["depth"], + saved.get("text_heads", 4), + saved["context"], + ).to(args.device) + model.load_state_dict(state["model"]) + model.eval() + phrases_per_row = [parse_group_phrases(captions[row][0]) for row in rows] + flat = [ + (index, phrase) + for index, phrases in enumerate(phrases_per_row) + for phrase in phrases + ] + states: list[list[torch.Tensor]] = [[] for _ in rows] + for indices in tqdm(list(batch_indices(len(flat), 256)), desc="text sets"): + batch = [flat[i] for i in indices] + sequences = [tokenize(phrase, vocab) for _, phrase in batch] + longest = max(len(s) for s in sequences) + tokens = torch.zeros(len(batch), longest, dtype=torch.long) + for row, sequence in enumerate(sequences): + tokens[row, : len(sequence)] = torch.tensor(sequence) + tokens = tokens.to(args.device) + hidden = model(tokens) + mask = (tokens != 0).float()[..., None] + pooled = (hidden * mask).sum(1) / mask.sum(1).clamp_min(1.0) + for (index, _), vector in zip(batch, pooled.float().cpu()): + states[index].append(vector) + return [F.normalize(torch.stack(s), dim=-1) for s in states] + + +def set_similarity_field( + sets: list[torch.Tensor], weights: list[torch.Tensor] | None = None +) -> torch.Tensor: + """Symmetric matching-value similarity between all set pairs.""" + n = len(sets) + field = torch.zeros(n, n) + for a in range(n): + for b in range(a, n): + similarity = sets[a] @ sets[b].T + if weights is not None: + similarity = similarity * torch.sqrt( + weights[a][:, None] * weights[b][None, :] + ) + rows, cols = linear_sum_assignment(-similarity.numpy()) + value = float(similarity[rows, cols].sum()) / max( + min(similarity.shape), 1 + ) + field[a, b] = field[b, a] = value + return field + + +@torch.inference_mode() +def decode_sprites( + rows: list[int], + lookup: dict[int, int], + slot_sets_all: torch.Tensor, + args: argparse.Namespace, +) -> list[torch.Tensor]: + """Centered per-slot appearance sprites from the tower's own decoder. + + The joint decode gives each slot an rgb map and a competitive alpha + mask; centering the masked appearance at the alpha centroid removes + layout, leaving color, shape, size, and multiplicity pattern. + """ + manifest = read_json(Path(args.data_dir, "manifest.json")) + state = torch.load(args.slot_tower, map_location="cpu", weights_only=False) + saved = state["args"] + model = SlotAutoencoder( + manifest["image_size"], saved["slots"], saved["slot_dim"], saved["iterations"] + ).to(args.device) + model.load_state_dict(state["model"]) + model.eval() + size = manifest["image_size"] + window = args.sprite_window + axis = torch.arange(size, dtype=torch.float32, device=args.device) + sprites: list[torch.Tensor] = [] + for start in tqdm(range(0, len(rows), 64), desc="sprites"): + batch_rows = rows[start : start + 64] + slots = torch.stack( + [slot_sets_all[lookup[int(row)]][0] for row in batch_rows] + ).to(args.device) + rgb_alpha_rgb, alpha = model.decode(slots) + del rgb_alpha_rgb + # Re-decode retaining per-slot rgb: replicate decode internals. + batch, count, dim = slots.shape + x = slots.reshape(batch * count, dim, 1, 1).expand( + -1, -1, model.broadcast, model.broadcast + ) + from .synth_slots import coordinate_grid + + grid = coordinate_grid(model.broadcast, slots.device).reshape(1, -1, 4) + position = model.position_decoder(grid).transpose(1, 2).reshape( + 1, dim, model.broadcast, model.broadcast + ) + decoded = model.decoder(x + position) + decoded = F.interpolate( + decoded, size=size, mode="bilinear", align_corners=False + ).reshape(batch, count, 4, size, size) + rgb = decoded[:, :, :3] + masked = rgb * alpha # [B, K, 3, H, W] + weight_y = alpha.squeeze(2).sum(-1) # [B, K, H] + weight_x = alpha.squeeze(2).sum(-2) # [B, K, W] + cy = (weight_y * axis).sum(-1) / weight_y.sum(-1).clamp_min(1e-6) + cx = (weight_x * axis).sum(-1) / weight_x.sum(-1).clamp_min(1e-6) + half = window // 2 + batch_sprites = torch.zeros(batch, count, 3, window, window) + padded = F.pad(masked, (half, half, half, half)) + for b in range(batch): + for k in range(count): + y0 = int(cy[b, k].round()) + x0 = int(cx[b, k].round()) + batch_sprites[b, k] = padded[ + b, k, :, y0 : y0 + window, x0 : x0 + window + ].cpu() + sprites.extend(batch_sprites.flatten(2).unbind(0)) + return sprites + + +def main() -> None: + args = parse_args() + seed_everything(args.seed) + manifest = read_json(Path(args.data_dir, "manifest.json")) + captions = read_json(Path(args.data_dir, "captions.json"))["captions"] + rows = manifest[args.split][: args.samples] + + slot_state = torch.load(args.slots, map_location="cpu", weights_only=False) + lookup = {int(row): i for i, row in enumerate(slot_state["rows"])} + slot_sets_all = slot_state["slot_sets"] + masses_all = slot_state["slot_masses"] + sprites_all = None + if args.vision_mode == "sprites": + sprites_all = decode_sprites(rows, lookup, slot_sets_all, args) + vision_sets, vision_weights = [], [] + for position, row in enumerate(rows): + index = lookup[int(row)] + # Slot identity is not stable across forwards, so views cannot be + # averaged slot-wise; one view keeps object-slot binding intact. + slots = slot_sets_all[index][0] # [K, D] + mass = masses_all[index][0] + dominant = mass.argmax() + keep = torch.ones(len(mass), dtype=torch.bool) + keep[dominant] = False + keep &= mass > args.mass_floor + if not keep.any(): + keep = torch.ones(len(mass), dtype=torch.bool) + if sprites_all is not None: + vision_sets.append(F.normalize(sprites_all[position][keep], dim=-1)) + else: + vision_sets.append(F.normalize(slots[keep], dim=-1)) + weight = mass[keep] + vision_weights.append(weight / weight.sum().clamp_min(1e-8)) + + text_sets = text_group_sets(rows, captions, args) + + print(json.dumps({"building": "set similarity fields"})) + visual_field = set_similarity_field(vision_sets, vision_weights) + text_field = set_similarity_field(text_sets) + + visual_channels = standardize_relation(visual_field.double())[0][None] + text_channels = standardize_relation(text_field.double())[0][None] + generator = torch.Generator().manual_seed(args.seed) + report = { + "protocol": ( + "Scene states are sets (slot vectors; per-group phrase " + "states); within-modality relations are set-matching values; " + "hidden pairs score orderings only." + ), + "split": args.split, + "samples": len(rows), + "mean_vision_set_size": float( + torch.tensor([len(s) for s in vision_sets]).float().mean() + ), + "mean_text_set_size": float( + torch.tensor([len(s) for s in text_sets]).float().mean() + ), + **run_gates(text_channels, visual_channels, args, generator), + } + verdict = { + "true_z_mse": report["gate_a"]["random"]["mse"]["true_z"], + "improving_fraction": report["gate_b"]["improving_fraction"], + "descent_keeps": report["descent_from_true"]["final_accuracy"], + "true_mse": report["gate_a"]["true"]["mse"], + "best_random_descent": min( + (r["final_objective"] for r in report["descent_from_random"]), + default=None, + ), + } + verdict["counterfeit_found"] = bool( + verdict["best_random_descent"] is not None + and verdict["best_random_descent"] < verdict["true_mse"] + ) + report["verdict"] = verdict + print(json.dumps({"verdict": verdict})) + write_json(args.output, report) + print(f"Wrote {args.output}") + + +if __name__ == "__main__": + main() -- cgit v1.2.3