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, 276 insertions, 0 deletions

diff --git a/worldalign/synth_transfer_energy.py b/worldalign/synth_transfer_energy.py
new file mode 100644
index 0000000..67a081f
--- /dev/null
+++ b/worldalign/synth_transfer_energy.py
@@ -0,0 +1,276 @@
+"""Prediction-transfer alignment energy for the synthetic world.
+
+The third rung of the energy ladder: no hand-designed geometry and no
+pair-trained EBM. Each modality's own predictor induces a substitution
+kernel over scenes -- which other scenes the world model finds compatible
+-- and the alignment energy is the conjugacy defect of the two kernels
+under a candidate coupling.
+
+- Vision kernel: the masked-reconstruction tower inpaints a half-masked
+  render; pixel distance between the inpainting and other scenes' renders
+  gives K_V[i, i'].
+- Text kernel: the causal LM scores scene j's caption sentences as a
+  continuation of scene i's caption prefix; referential consistency makes
+  same-content continuations likely, giving K_T.
+
+Both kernels are unimodal-predictor functionals. Hidden pairs enter only
+through the gate that scores orderings of the energy.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+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_towers import TextTower, VisionTower, load_image, tokenize
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-dir", default="artifacts/synth_v0")
+    parser.add_argument(
+        "--vision-tower", default="artifacts/synth_v0/vision_tower_simmim.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("--mask-ratio", type=float, default=0.5)
+    parser.add_argument("--kernel-temperature", type=float, default=0.03)
+    parser.add_argument("--batch-size", type=int, default=256)
+    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/transfer_energy_gate.json"
+    )
+    parser.add_argument(
+        "--kernels-output", default="artifacts/synth_v0/transfer_kernels.pt"
+    )
+    return parser.parse_args()
+
+
+@torch.inference_mode()
+def vision_kernel(
+    rows: list[int], manifest: dict, args: argparse.Namespace
+) -> torch.Tensor:
+    state = torch.load(args.vision_tower, map_location="cpu", weights_only=False)
+    saved = state["args"]
+    model = VisionTower(
+        manifest["image_size"], saved["patch"], saved["dim"], saved["depth"], saved["heads"]
+    ).to(args.device)
+    model.load_state_dict(state["model"])
+    model.eval()
+    image_dir = Path(manifest["image_dir"])
+    tokens = (manifest["image_size"] // saved["patch"]) ** 2
+    generator = torch.Generator(device=args.device).manual_seed(args.seed)
+
+    references = []
+    for indices in batch_indices(len(rows), args.batch_size):
+        pixels = torch.stack(
+            [load_image(image_dir / f"scene{rows[i]:06d}_v0.png") for i in indices]
+        )
+        references.append(pixels)
+    references = torch.cat(references)  # [N, 3, H, W] view-0 renders on cpu
+
+    inpainted = []
+    for indices in tqdm(
+        list(batch_indices(len(rows), args.batch_size)), desc="vision kernel"
+    ):
+        # Inpaint view 1 under a random half mask; compare to view-0 renders.
+        pixels = torch.stack(
+            [load_image(image_dir / f"scene{rows[i]:06d}_v1.png") for i in indices]
+        ).to(args.device)
+        mask = (
+            torch.rand(len(pixels), tokens, generator=generator, device=args.device)
+            < args.mask_ratio
+        )
+        encoded = model.encode(pixels, mask=mask)
+        patches = model.reconstruction(encoded[:, 1:])  # [B, T, p*p*3]
+        grid = manifest["image_size"] // saved["patch"]
+        patch = saved["patch"]
+        image = patches.reshape(len(pixels), grid, grid, patch, patch, 3)
+        image = image.permute(0, 5, 1, 3, 2, 4).reshape(
+            len(pixels), 3, manifest["image_size"], manifest["image_size"]
+        )
+        blended = torch.where(
+            mask.reshape(len(pixels), 1, grid, 1, grid, 1)
+            .expand(-1, 3, -1, patch, -1, patch)
+            .reshape_as(image),
+            image,
+            pixels,
+        )
+        inpainted.append(blended.float().cpu())
+    inpainted = torch.cat(inpainted)
+
+    # Layout-invariant content distance: Chamfer between foreground patch
+    # sets in pixel space. Renders of the same scene share patch CONTENT
+    # (colors, local shape fragments) while positions are resampled, so
+    # per-pixel comparison would measure layout instead.
+    patch = saved["patch"]
+    grid = manifest["image_size"] // patch
+
+    def patch_sets(images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        pieces = images.reshape(len(images), 3, grid, patch, grid, patch)
+        pieces = pieces.permute(0, 2, 4, 1, 3, 5).reshape(len(images), grid * grid, -1)
+        background = pieces.mean(-1)
+        foreground = background > background.min(1, keepdim=True).values + 0.02
+        return pieces, foreground
+
+    inpaint_pieces, inpaint_fg = patch_sets(inpainted)
+    reference_pieces, reference_fg = patch_sets(references)
+    distances = torch.zeros(len(rows), len(rows))
+    device = args.device
+    reference_pieces = reference_pieces.to(device)
+    reference_fg = reference_fg.to(device)
+    for start in range(0, len(rows), 32):
+        stop = min(start + 32, len(rows))
+        block = inpaint_pieces[start:stop].to(device)
+        block_fg = inpaint_fg[start:stop].to(device)
+        cross = torch.cdist(block, reference_pieces.reshape(-1, block.shape[-1]))
+        cross = cross.reshape(stop - start, block.shape[1], len(rows), -1)
+        masked_ab = cross.masked_fill(
+            ~reference_fg[None, None, :, :], float("inf")
+        ).amin(-1)
+        forward = (
+            (masked_ab * block_fg[:, :, None]).sum(1)
+            / block_fg.sum(1).clamp_min(1)[:, None]
+        )
+        masked_ba = cross.masked_fill(
+            ~block_fg[:, :, None, None].expand_as(cross), float("inf")
+        ).amin(1)
+        backward = (
+            (masked_ba * reference_fg[None, :, :]).sum(-1)
+            / reference_fg.sum(-1).clamp_min(1)[None, :]
+        )
+        distances[start:stop] = (0.5 * (forward + backward)).float().cpu()
+    distances[torch.isinf(distances) | torch.isnan(distances)] = distances[
+        torch.isfinite(distances)
+    ].max()
+    return distances
+
+
+@torch.inference_mode()
+def text_kernel(
+    rows: list[int], captions: list[list[str]], args: argparse.Namespace
+) -> 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()
+
+    prefixes = [tokenize(captions[row][0], vocab) for row in rows]
+    # Continuations are the RELATION sentences only: they refer back to
+    # groups the prefix must have introduced, so their likelihood grades
+    # shared content instead of rewarding verbatim repetition.
+    continuations = []
+    for row in rows:
+        sentences = captions[row][1].split(". ")
+        relational = ". ".join(sentences[1:]) if len(sentences) > 1 else sentences[0]
+        continuations.append(tokenize(relational, vocab)[1:])  # drop bos
+    scores = torch.zeros(len(rows), len(rows))
+    jobs = [(i, j) for i in range(len(rows)) for j in range(len(rows))]
+    for indices in tqdm(
+        list(batch_indices(len(jobs), args.batch_size)), desc="text kernel"
+    ):
+        batch = [jobs[k] for k in indices]
+        sequences = [
+            (prefixes[i] + continuations[j])[: saved["context"]] for i, j 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)
+        logits = model.logits(hidden[:, :-1])
+        log_probs = F.log_softmax(logits, dim=-1)
+        for row, (i, j) in enumerate(batch):
+            begin = len(prefixes[i]) - 1
+            end = min(len(prefixes[i]) + len(continuations[j]), longest) - 1
+            if end <= begin:
+                scores[i, j] = 0.0
+                continue
+            targets = tokens[row, begin + 1 : end + 1]
+            picked = log_probs[row, begin:end].gather(1, targets[:, None])
+            scores[i, j] = float(picked.mean())
+    return -scores  # negative mean log-likelihood as a distance
+
+
+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]
+
+    visual_distance = vision_kernel(rows, manifest, args)
+    text_distance = text_kernel(rows, captions, args)
+    torch.save(
+        {
+            "rows": rows,
+            "visual_distance": visual_distance,
+            "text_distance": text_distance,
+        },
+        args.kernels_output,
+    )
+
+    # Substitution kernels as standardized relation channels: negative
+    # distances, symmetrized, standardized off-diagonal. The gate machinery
+    # then scores orderings exactly as for any relation field.
+    visual_similarity = -visual_distance
+    text_similarity = -0.5 * (text_distance + text_distance.T)
+    visual_channels = standardize_relation(
+        0.5 * (visual_similarity + visual_similarity.T).double()
+    )[0][None]
+    text_channels = standardize_relation(text_similarity.double())[0][None]
+    generator = torch.Generator().manual_seed(args.seed)
+    report = {
+        "protocol": (
+            "Substitution kernels are unimodal-predictor functionals "
+            "(masked inpainting distance; caption continuation "
+            "likelihood). Hidden pairs score orderings only."
+        ),
+        "split": args.split,
+        "samples": len(rows),
+        **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()