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

diff --git a/worldalign/blind_recovery.py b/worldalign/blind_recovery.py
new file mode 100644
index 0000000..625f407
--- /dev/null
+++ b/worldalign/blind_recovery.py
@@ -0,0 +1,476 @@
+"""Blind unpaired recovery on gate-passing states.
+
+The basin audit licenses this experiment: on content-projected VG states
+the true assignment sits below every quench, so the remaining question is
+search. Three arms attack the measured landscape shape (deep true basin,
+glassy surroundings, smooth coarse spectrum):
+
+A. entropic Sinkhorn annealing: soft coupling, temperature and entropy
+   schedules, barycentric language states;
+B. parallel tempering: replica-exchange Metropolis over permutations with
+   the exact closed-form swap deltas, imported from spin-glass practice;
+C. spectral-band homotopy: solve the coupling in a low-dimensional
+   spectral band first, then warm-start progressively finer bands --
+   a band-limited continuation in the sense of low-order Fourier terms on
+   the symmetric group, instantiated through the content spectrum.
+
+The text side enters through a hidden shuffle, so the optimizer's identity
+carries no information. The hidden truth scores results and is used
+nowhere else.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from scipy.optimize import linear_sum_assignment
+
+from .common import seed_everything, write_json
+from .content_gate import flickr_states, vg_states
+from .energy import log_sinkhorn
+from .manifold_gate import standardize_relation
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", choices=["flickr", "vg"], default="vg")
+    parser.add_argument("--manifest", default="artifacts/manifest.json")
+    parser.add_argument("--vision", default="artifacts/vision.pt")
+    parser.add_argument("--text", default="artifacts/text.pt")
+    parser.add_argument("--text-orbits", default="artifacts/text_orbits_qwen0p5b.pt")
+    parser.add_argument("--vg-vision", default="artifacts/vg_5k/vision_features.pt")
+    parser.add_argument("--vg-text", default="artifacts/vg_5k/text_features.pt")
+    parser.add_argument(
+        "--vg-ground-truth", default="artifacts/vg_5k/ground_truth.private.jsonl"
+    )
+    parser.add_argument("--split", choices=["val", "test"], default="test")
+    parser.add_argument("--samples", type=int, default=512)
+    parser.add_argument("--subset-seed", type=int, default=0)
+    parser.add_argument("--dims", type=int, default=256)
+    parser.add_argument("--shrinkage", type=float, default=0.05)
+    parser.add_argument("--arms", default="tempering,sinkhorn,homotopy")
+    parser.add_argument("--replicas", type=int, default=8)
+    parser.add_argument("--tempering-rounds", type=int, default=40000)
+    parser.add_argument("--temp-high", type=float, default=3e-3)
+    parser.add_argument("--temp-low", type=float, default=1e-5)
+    parser.add_argument("--exchange-every", type=int, default=20)
+    parser.add_argument("--sinkhorn-steps", type=int, default=2500)
+    parser.add_argument("--sinkhorn-restarts", type=int, default=4)
+    parser.add_argument("--homotopy-bands", default="4,8,16,32,64,full")
+    parser.add_argument(
+        "--unary-matrix",
+        default="",
+        help="Optional anchor similarity matrix (.pt) in ground-truth node "
+        "order; injected as a unary field on the tempering energy.",
+    )
+    parser.add_argument("--unary-weight", type=float, default=2.0)
+    parser.add_argument("--init", choices=["random", "unary"], default="random")
+    parser.add_argument("--device", default="cuda:3")
+    parser.add_argument("--seed", type=int, default=20260730)
+    parser.add_argument("--output", required=True)
+    return parser.parse_args()
+
+
+def relation_mse_soft(
+    language_states: torch.Tensor, visual_standardized: torch.Tensor
+) -> torch.Tensor:
+    """Differentiable standardized-relation MSE for soft states."""
+    normalized = F.normalize(language_states, dim=-1)
+    relation = normalized @ normalized.T
+    mask = ~torch.eye(len(relation), dtype=torch.bool, device=relation.device)
+    values = relation[mask]
+    standardized = (values - values.mean()) / values.std().clamp_min(1e-6)
+    return (standardized - visual_standardized[mask]).square().mean()
+
+
+def permutation_energy_batch(
+    text_standardized: torch.Tensor,
+    visual_standardized: torch.Tensor,
+    permutations: torch.Tensor,
+) -> torch.Tensor:
+    fields = text_standardized[permutations[:, :, None], permutations[:, None, :]]
+    mask = ~torch.eye(
+        text_standardized.shape[-1], dtype=torch.bool, device=fields.device
+    )
+    return ((fields - visual_standardized) ** 2)[:, mask].mean(-1)
+
+
+def proposal_swap_deltas(
+    permuted_fields: torch.Tensor,
+    visual_standardized: torch.Tensor,
+    pairs_p: torch.Tensor,
+    pairs_q: torch.Tensor,
+) -> torch.Tensor:
+    """Exact deltas for proposed swaps only, O(N) per proposal.
+
+    permuted_fields: [R, N, N] text fields under each replica's current
+    permutation; pairs_p/pairs_q: [R, P] proposal endpoints.
+    """
+    size = permuted_fields.shape[-1]
+    count = size * (size - 1)
+    replica_index = torch.arange(len(permuted_fields), device=permuted_fields.device)
+    rows_p = permuted_fields[replica_index[:, None], pairs_p]  # [R, P, N]
+    rows_q = permuted_fields[replica_index[:, None], pairs_q]
+    visual_p = visual_standardized[pairs_p]
+    visual_q = visual_standardized[pairs_q]
+    self_p = (rows_p * visual_p).sum(-1)
+    self_q = (rows_q * visual_q).sum(-1)
+    cross_pq = (rows_p * visual_q).sum(-1)
+    cross_qp = (rows_q * visual_p).sum(-1)
+    direct = permuted_fields[replica_index[:, None], pairs_p, pairs_q]
+    direct_visual = visual_standardized[pairs_p, pairs_q]
+    total = self_p + self_q - cross_pq - cross_qp - 2.0 * direct * direct_visual
+    delta = (4.0 / count) * total
+    return delta.masked_fill(pairs_p == pairs_q, float("inf"))
+
+
+def score(
+    permutation: torch.Tensor, truth: torch.Tensor, energy: float, true_energy: float
+) -> dict:
+    return {
+        "accuracy": float((permutation.cpu() == truth.cpu()).double().mean()),
+        "energy": energy,
+        "energy_over_true": energy / true_energy - 1.0,
+    }
+
+
+def coupling_metrics(coupling: torch.Tensor, truth: torch.Tensor) -> dict:
+    n = len(coupling)
+    truth = truth.to(coupling.device)
+    mass_at_truth = float(coupling[torch.arange(n, device=coupling.device), truth].mean())
+    rows, cols = linear_sum_assignment(-coupling.detach().cpu().numpy())
+    permutation = torch.from_numpy(cols)
+    forward = coupling.argmax(-1).cpu()
+    backward = coupling.argmax(0).cpu()
+    mutual = backward[forward] == torch.arange(n)
+    sorted_mass = coupling.sort(-1, descending=True).values
+    margin = (sorted_mass[:, 0] - sorted_mass[:, 1]).cpu()
+    top = margin.argsort(descending=True)[: max(1, n // 10)]
+    return {
+        "mass_at_truth": mass_at_truth,
+        "hungarian_accuracy": float((permutation == truth.cpu()).double().mean()),
+        "mutual_nn_count": int(mutual.sum()),
+        "mutual_nn_precision": float(
+            (forward[mutual] == truth.cpu()[mutual]).double().mean()
+        )
+        if mutual.any()
+        else None,
+        "top_margin_decile_precision": float(
+            (forward[top] == truth.cpu()[top]).double().mean()
+        ),
+    }
+
+
+def arm_tempering(
+    text_standardized: torch.Tensor,
+    visual_standardized: torch.Tensor,
+    truth: torch.Tensor,
+    true_energy: float,
+    args: argparse.Namespace,
+    generator: torch.Generator,
+    unary: torch.Tensor | None = None,
+) -> dict:
+    device = text_standardized.device
+    size = text_standardized.shape[-1]
+    replicas = args.replicas
+    weight = args.unary_weight if unary is not None else 0.0
+
+    def total_energy(permutations: torch.Tensor) -> torch.Tensor:
+        energy = permutation_energy_batch(
+            text_standardized, visual_standardized, permutations
+        )
+        if unary is not None:
+            index = torch.arange(size, device=device)
+            energy = energy - weight * unary[index, permutations].mean(-1)
+        return energy
+
+    temperatures = torch.logspace(
+        torch.log10(torch.tensor(args.temp_low)),
+        torch.log10(torch.tensor(args.temp_high)),
+        replicas,
+    ).to(device)
+    if args.init == "unary" and unary is not None:
+        rows, cols = linear_sum_assignment(-unary.cpu().numpy())
+        start = torch.from_numpy(cols).to(device)
+        permutations = torch.stack([start.clone() for _ in range(replicas)])
+    else:
+        permutations = torch.stack(
+            [
+                torch.randperm(size, generator=generator).to(device)
+                for _ in range(replicas)
+            ]
+        )
+    energies = total_energy(permutations)
+    best = {"energy": float("inf"), "permutation": permutations[0].clone()}
+    accepted = 0
+    for round_index in range(args.tempering_rounds):
+        fields = text_standardized[
+            permutations[:, :, None], permutations[:, None, :]
+        ]
+        pairs_p = torch.randint(0, size, (replicas, 24), generator=generator).to(
+            device
+        )
+        pairs_q = torch.randint(0, size, (replicas, 24), generator=generator).to(
+            device
+        )
+        proposal_deltas = proposal_swap_deltas(
+            fields, visual_standardized, pairs_p, pairs_q
+        )
+        if unary is not None:
+            replica_index = torch.arange(replicas, device=device)[:, None]
+            assigned_p = permutations[replica_index, pairs_p]
+            assigned_q = permutations[replica_index, pairs_q]
+            unary_delta = (
+                unary[pairs_p, assigned_q]
+                + unary[pairs_q, assigned_p]
+                - unary[pairs_p, assigned_p]
+                - unary[pairs_q, assigned_q]
+            )
+            proposal_deltas = proposal_deltas - (weight / size) * unary_delta
+        noise = torch.rand(replicas, 24, generator=generator).to(device)
+        acceptable = (
+            proposal_deltas < -temperatures[:, None] * noise.clamp_min(1e-12).log()
+        ) & torch.isfinite(proposal_deltas)
+        for replica in range(replicas):
+            hits = torch.nonzero(acceptable[replica])
+            if not len(hits):
+                continue
+            first = int(hits[0, 0])
+            p = int(pairs_p[replica, first])
+            q = int(pairs_q[replica, first])
+            permutations[replica][[p, q]] = permutations[replica][[q, p]]
+            energies[replica] = energies[replica] + proposal_deltas[replica, first]
+            accepted += 1
+        if round_index % args.exchange_every == 0:
+            for replica in range(replicas - 1):
+                gap = (energies[replica] - energies[replica + 1]) * (
+                    1.0 / temperatures[replica] - 1.0 / temperatures[replica + 1]
+                )
+                if gap > 0 or torch.rand(1, generator=generator).item() < float(
+                    gap.exp()
+                ):
+                    permutations[[replica, replica + 1]] = permutations[
+                        [replica + 1, replica]
+                    ]
+                    energies[[replica, replica + 1]] = energies[
+                        [replica + 1, replica]
+                    ]
+        cold = int(energies.argmin())
+        if float(energies[cold]) < best["energy"]:
+            best = {
+                "energy": float(energies[cold]),
+                "permutation": permutations[cold].clone(),
+            }
+        if round_index % 5000 == 0:
+            energies = total_energy(permutations)  # refresh against drift
+            print(
+                json.dumps(
+                    {
+                        "tempering_round": round_index,
+                        "cold_energy": float(energies.min()),
+                        "cold_accuracy": float(
+                            (permutations[int(energies.argmin())].cpu() == truth)
+                            .double()
+                            .mean()
+                        ),
+                        "accepted": accepted,
+                    }
+                )
+            )
+    final = [
+        score(
+            permutations[r],
+            truth,
+            float(
+                permutation_energy_batch(
+                    text_standardized, visual_standardized, permutations[r : r + 1]
+                )[0]
+            ),
+            true_energy,
+        )
+        for r in range(replicas)
+    ]
+    best_score = score(
+        best["permutation"],
+        truth,
+        float(
+            permutation_energy_batch(
+                text_standardized, visual_standardized, best["permutation"][None]
+            )[0]
+        ),
+        true_energy,
+    )
+    return {"replicas": final, "best": best_score, "accepted_moves": accepted}
+
+
+def arm_sinkhorn(
+    text_states: torch.Tensor,
+    visual_standardized: torch.Tensor,
+    truth: torch.Tensor,
+    true_energy: float,
+    args: argparse.Namespace,
+    generator: torch.Generator,
+    bands: list[int | None] | None = None,
+) -> dict:
+    device = text_states.device
+    size = len(text_states)
+    restarts = []
+    for restart in range(args.sinkhorn_restarts if bands is None else 1):
+        logits = torch.nn.Parameter(
+            0.01
+            * torch.randn(size, size, generator=generator).to(device)
+        )
+        optimizer = torch.optim.Adam([logits], lr=0.08)
+        stage_states = text_states
+        stages = bands if bands is not None else [None]
+        for stage_index, band in enumerate(stages):
+            if band is not None:
+                keep = min(band, text_states.shape[-1])
+                stage_states = text_states[:, :keep]
+            steps = args.sinkhorn_steps // len(stages)
+            for step in range(steps):
+                progress = step / max(steps - 1, 1)
+                temperature = 0.5 * (0.05 / 0.5) ** progress
+                coupling = log_sinkhorn(logits, temperature, iterations=12)
+                barycenter = coupling @ stage_states
+                loss = relation_mse_soft(barycenter, visual_standardized)
+                entropy = -(coupling * coupling.clamp_min(1e-12).log()).sum(-1).mean()
+                loss = loss + (0.002 + 0.02 * progress) * entropy
+                optimizer.zero_grad(set_to_none=True)
+                loss.backward()
+                optimizer.step()
+        with torch.no_grad():
+            coupling = log_sinkhorn(logits, 0.05, iterations=30)
+        metrics = coupling_metrics(coupling, truth)
+        rows, cols = linear_sum_assignment(-coupling.detach().cpu().numpy())
+        permutation = torch.from_numpy(cols).to(device)
+        rounded_energy = float(
+            permutation_energy_batch(
+                standardize_relation(
+                    F.normalize(text_states, dim=-1)
+                    @ F.normalize(text_states, dim=-1).T
+                )[0][None].squeeze(0),
+                visual_standardized,
+                permutation[None],
+            )[0]
+        )
+        metrics.update(
+            score(permutation, truth, rounded_energy, true_energy)
+        )
+        restarts.append(metrics)
+        print(json.dumps({"sinkhorn_restart": restart, **metrics}))
+    return {"restarts": restarts}
+
+
+def main() -> None:
+    args = parse_args()
+    seed_everything(args.seed)
+    if args.dataset == "flickr":
+        text_states, visual_states = flickr_states(args)
+    else:
+        text_states, visual_states = vg_states(args)
+    device = torch.device(args.device)
+    generator = torch.Generator().manual_seed(args.seed)
+
+    size = len(text_states)
+    hidden = torch.randperm(size, generator=generator)
+    truth = torch.argsort(hidden)  # input row j holds true counterpart hidden[j]
+    text_input = text_states[hidden].float().to(device)
+    visual_states = visual_states.float().to(device)
+
+    unary = None
+    if args.unary_matrix:
+        anchor_state = torch.load(
+            args.unary_matrix, map_location="cpu", weights_only=False
+        )
+        combined = anchor_state["combined"].double()
+        subset_generator = torch.Generator().manual_seed(args.subset_seed)
+        subset = torch.randperm(len(combined), generator=subset_generator)[
+            : args.samples
+        ]
+        unary_subset = combined[subset][:, subset]
+        unary_subset = (unary_subset - unary_subset.mean()) / unary_subset.std()
+        unary = unary_subset[:, hidden].float().to(device)
+
+    text_standardized = standardize_relation(
+        F.normalize(text_input, dim=-1) @ F.normalize(text_input, dim=-1).T
+    )[0]
+    visual_standardized = standardize_relation(
+        F.normalize(visual_states, dim=-1) @ F.normalize(visual_states, dim=-1).T
+    )[0]
+    true_energy = float(
+        permutation_energy_batch(
+            text_standardized, visual_standardized, truth[None].to(device)
+        )[0]
+    )
+    report: dict = {
+        "protocol": (
+            "The text side enters through a hidden shuffle; the optimizer "
+            "never sees pair, order, or truth information. Hidden truth "
+            "scores outcomes only."
+        ),
+        "dataset": args.dataset,
+        "dims": args.dims,
+        "samples": size,
+        "true_energy": true_energy,
+        "chance_accuracy": 1.0 / size,
+        "arms": {},
+    }
+    arms = [arm.strip() for arm in args.arms.split(",")]
+    if "tempering" in arms:
+        report["arms"]["tempering"] = arm_tempering(
+            text_standardized,
+            visual_standardized,
+            truth,
+            true_energy,
+            args,
+            generator,
+            unary=unary,
+        )
+        if unary is not None:
+            report["unary"] = {
+                "weight": args.unary_weight,
+                "matched_mean": float(
+                    unary[torch.arange(size, device=unary.device), truth.to(unary.device)].mean()
+                ),
+                "grand_mean": float(unary.mean()),
+                "init": args.init,
+            }
+    if "sinkhorn" in arms:
+        report["arms"]["sinkhorn"] = arm_sinkhorn(
+            text_input, visual_standardized, truth, true_energy, args, generator
+        )
+    if "homotopy" in arms:
+        bands: list[int | None] = [
+            None if band == "full" else int(band)
+            for band in args.homotopy_bands.split(",")
+        ]
+        report["arms"]["homotopy"] = arm_sinkhorn(
+            text_input,
+            visual_standardized,
+            truth,
+            true_energy,
+            args,
+            generator,
+            bands=bands,
+        )
+    summary = {}
+    for arm, result in report["arms"].items():
+        candidates = result.get("restarts") or result.get("replicas") or []
+        if "best" in result:
+            candidates = candidates + [result["best"]]
+        best = max(candidates, key=lambda c: c["accuracy"], default=None)
+        summary[arm] = best
+    report["summary"] = summary
+    print(json.dumps({"summary": summary}))
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+    write_json(args.output, report)
+    print(f"Wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()