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