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

diff --git a/worldalign/synth_fast_gate.py b/worldalign/synth_fast_gate.py
new file mode 100644
index 0000000..54b6891
--- /dev/null
+++ b/worldalign/synth_fast_gate.py
@@ -0,0 +1,345 @@
+"""Closed-form batched gate: exact all-triple triangle energy.
+
+The sampled triangle energy cost 5.7 ms per proposal because each
+evaluation gathered a variable-length triple set in Python. It has a
+closed form. Writing M(sigma) for the elementwise product of the
+sigma-permuted text field with the visual field,
+
+    pairwise    = const - 2 * sum(M)
+    all-triple  = const - 2 * trace(M^3) / 6
+
+because trace of a power is invariant under simultaneous row-column
+permutation, so the text-only and vision-only terms do not move. One
+batched matrix product evaluates trace(M^3) = sum(M * (M @ M)) for
+hundreds of candidate permutations at once, over every C(N,3) triple
+rather than a sample.
+
+This buys a genuinely stronger searcher: exact steepest descent over all
+N(N-1)/2 transpositions per step, plus batched-proposal tempering.
+Hidden pairs score orderings only.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+
+from .common import read_json, seed_everything, write_json
+from .synth_triangle_gate import build_fields, standardized
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fields", default="", help="Saved .pt with visual_field/text_field.")
+    parser.add_argument("--data-dir", default="artifacts/synth_v0")
+    parser.add_argument("--split", choices=["val", "test"], default="test")
+    parser.add_argument("--samples", type=int, default=512)
+    parser.add_argument("--merge-distance", type=float, default=30.0)
+    parser.add_argument("--vision-views", type=int, default=4)
+    parser.add_argument("--energies", default="pair,triangle,both")
+    parser.add_argument("--chunk", type=int, default=512)
+    parser.add_argument("--descent-restarts", type=int, default=5)
+    parser.add_argument("--descent-max-steps", type=int, default=4000)
+    parser.add_argument("--replicas", type=int, default=8)
+    parser.add_argument("--rounds", type=int, default=3000)
+    parser.add_argument("--proposals", type=int, default=64)
+    parser.add_argument("--temp-high", type=float, default=3e-2)
+    parser.add_argument("--temp-low", type=float, default=1e-4)
+    parser.add_argument("--exchange-every", type=int, default=20)
+    parser.add_argument("--polish-steps", type=int, default=2000)
+    parser.add_argument("--device", default="cuda:3")
+    parser.add_argument("--seed", type=int, default=20260731)
+    parser.add_argument("--output", default="artifacts/synth_v0/fast_gate.json")
+    return parser.parse_args()
+
+
+class ClosedFormEnergy:
+    """Energy as a function of M = permuted-text * visual, batched."""
+
+    def __init__(
+        self,
+        text: torch.Tensor,
+        visual: torch.Tensor,
+        pair_weight: float,
+        triangle_weight: float,
+        chunk: int,
+    ) -> None:
+        self.text = text
+        self.visual = visual
+        self.pair_weight = pair_weight
+        self.triangle_weight = triangle_weight
+        self.chunk = chunk
+        self.size = len(visual)
+        self.pair_count = self.size * (self.size - 1)
+        self.triple_count = self.size * (self.size - 1) * (self.size - 2)
+        # Permutation-invariant constants, kept so reported values match
+        # the direct definitions of the two energies.
+        square_text = text * text
+        square_visual = visual * visual
+        self.pair_constant = float(square_text.sum() + square_visual.sum())
+        self.triangle_constant = float(
+            self._trace_cube(square_text[None])[0]
+            + self._trace_cube(square_visual[None])[0]
+        )
+
+    @staticmethod
+    def _trace_cube(matrices: torch.Tensor) -> torch.Tensor:
+        return (matrices * torch.bmm(matrices, matrices)).sum((-2, -1))
+
+    def energy(self, permutations: torch.Tensor) -> torch.Tensor:
+        """Exact energy for a batch of permutations [B, N]."""
+        values = []
+        for start in range(0, len(permutations), self.chunk):
+            block = permutations[start : start + self.chunk]
+            permuted = self.text[block[:, :, None], block[:, None, :]]
+            product = permuted * self.visual
+            total = torch.zeros(len(block), device=permuted.device)
+            if self.pair_weight:
+                pair = (self.pair_constant - 2.0 * product.sum((-2, -1))) / (
+                    self.pair_count
+                )
+                total = total + self.pair_weight * pair
+            if self.triangle_weight:
+                triangle = (
+                    self.triangle_constant - 2.0 * self._trace_cube(product)
+                ) / self.triple_count
+                total = total + self.triangle_weight * triangle
+            values.append(total)
+        return torch.cat(values)
+
+
+def all_swaps(size: int, device: torch.device) -> torch.Tensor:
+    rows, cols = torch.triu_indices(size, size, offset=1, device=device)
+    return torch.stack([rows, cols], dim=1)
+
+
+def apply_swaps(permutation: torch.Tensor, swaps: torch.Tensor) -> torch.Tensor:
+    batch = permutation[None].repeat(len(swaps), 1)
+    index = torch.arange(len(swaps), device=permutation.device)
+    p, q = swaps[:, 0], swaps[:, 1]
+    values_p = batch[index, p].clone()
+    batch[index, p] = batch[index, q]
+    batch[index, q] = values_p
+    return batch
+
+
+def steepest_descent(
+    energy: ClosedFormEnergy,
+    start: torch.Tensor,
+    swaps: torch.Tensor,
+    max_steps: int,
+) -> tuple[torch.Tensor, float]:
+    """Exact steepest descent over every transposition each step."""
+    current = start.clone()
+    value = float(energy.energy(current[None])[0])
+    for _ in range(max_steps):
+        candidates = apply_swaps(current, swaps)
+        values = energy.energy(candidates)
+        best = int(values.argmin())
+        if float(values[best]) >= value - 1e-12:
+            break
+        current = candidates[best]
+        value = float(values[best])
+    return current, value
+
+
+def temper(
+    energy: ClosedFormEnergy,
+    starts: torch.Tensor,
+    args: argparse.Namespace,
+    generator: torch.Generator,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Batched-proposal parallel tempering; returns states and energies."""
+    size = energy.size
+    replicas = len(starts)
+    temperatures = torch.logspace(
+        torch.log10(torch.tensor(args.temp_low)),
+        torch.log10(torch.tensor(args.temp_high)),
+        replicas,
+    ).to(device)
+    states = starts.clone()
+    values = energy.energy(states)
+    for round_index in range(args.rounds):
+        p = torch.randint(
+            0, size, (replicas, args.proposals), generator=generator
+        ).to(device)
+        q = torch.randint(
+            0, size, (replicas, args.proposals), generator=generator
+        ).to(device)
+        valid = p != q
+        batch = states[:, None, :].repeat(1, args.proposals, 1)
+        index_r = torch.arange(replicas, device=device)[:, None]
+        original_p = batch.gather(2, p[..., None]).squeeze(-1)
+        original_q = batch.gather(2, q[..., None]).squeeze(-1)
+        batch.scatter_(2, p[..., None], original_q[..., None])
+        batch.scatter_(2, q[..., None], original_p[..., None])
+        flat = batch.reshape(replicas * args.proposals, size)
+        proposal_values = energy.energy(flat).reshape(replicas, args.proposals)
+        deltas = proposal_values - values[:, None]
+        noise = torch.rand(
+            replicas, args.proposals, generator=generator
+        ).to(device)
+        threshold = -temperatures[:, None] * noise.clamp_min(1e-12).log()
+        accept = (deltas < threshold) & valid
+        first = torch.where(
+            accept.any(-1),
+            accept.float().argmax(-1),
+            torch.zeros(replicas, dtype=torch.long, device=device),
+        )
+        taken = accept.any(-1)
+        chosen = batch[index_r.squeeze(-1), first]
+        states = torch.where(taken[:, None], chosen, states)
+        values = torch.where(
+            taken, proposal_values[index_r.squeeze(-1), first], values
+        )
+        if round_index % args.exchange_every == 0:
+            for replica in range(replicas - 1):
+                gap = (values[replica] - values[replica + 1]) * (
+                    1.0 / temperatures[replica] - 1.0 / temperatures[replica + 1]
+                )
+                accept_swap = gap > 0 or float(
+                    torch.rand(1, generator=generator)
+                ) < float(gap.exp().clamp(max=1.0))
+                if accept_swap:
+                    states[[replica, replica + 1]] = states[
+                        [replica + 1, replica]
+                    ]
+                    values[[replica, replica + 1]] = values[
+                        [replica + 1, replica]
+                    ]
+    return states, values
+
+
+def main() -> None:
+    args = parse_args()
+    seed_everything(args.seed)
+    if args.fields:
+        state = torch.load(args.fields, map_location="cpu", weights_only=False)
+        visual_field, text_field = state["visual_field"], state["text_field"]
+    else:
+        manifest = read_json(Path(args.data_dir, "manifest.json"))
+        rows = manifest[args.split][: args.samples]
+        visual_field, text_field = build_fields(args, rows, manifest)
+
+    device = torch.device(args.device)
+    size = len(visual_field)
+    generator = torch.Generator().manual_seed(args.seed)
+    hidden = torch.randperm(size, generator=generator)
+    truth = torch.argsort(hidden).to(device)  # scored, never optimized against
+    text = standardized(text_field[hidden][:, hidden].double().to(device)).float()
+    visual = standardized(visual_field.double().to(device)).float()
+    swaps = all_swaps(size, device)
+
+    weights = {"pair": (1.0, 0.0), "triangle": (0.0, 1.0), "both": (1.0, 1.0)}
+    report = {
+        "protocol": (
+            "Closed-form energies over all triples; the decision statistic "
+            "is the energy of the truth against the deepest state reached "
+            "by exact steepest descent and batched tempering. Hidden pairs "
+            "score only."
+        ),
+        "samples": size,
+        "energies": {},
+    }
+    for name in (item.strip() for item in args.energies.split(",")):
+        pair_weight, triangle_weight = weights[name]
+        energy = ClosedFormEnergy(
+            text, visual, pair_weight, triangle_weight, args.chunk
+        )
+        true_energy = float(energy.energy(truth[None])[0])
+        random_batch = torch.stack(
+            [
+                torch.argsort(torch.rand(size, generator=generator))
+                for _ in range(200)
+            ]
+        ).to(device)
+        random_values = energy.energy(random_batch)
+
+        kept, kept_value = steepest_descent(
+            energy, truth, swaps, args.descent_max_steps
+        )
+        retention = float((kept.cpu() == truth.cpu()).float().mean())
+
+        quenches = []
+        for restart in range(args.descent_restarts):
+            start = torch.argsort(torch.rand(size, generator=generator)).to(device)
+            final, value = steepest_descent(
+                energy, start, swaps, args.descent_max_steps
+            )
+            quenches.append(
+                {
+                    "energy": value,
+                    "accuracy": float((final.cpu() == truth.cpu()).float().mean()),
+                }
+            )
+
+        starts = torch.stack(
+            [
+                torch.argsort(torch.rand(size, generator=generator))
+                for _ in range(args.replicas)
+            ]
+        ).to(device)
+        states, values = temper(energy, starts, args, generator, device)
+        cold = int(values.argmin())
+        polished, polished_value = steepest_descent(
+            energy, states[cold], swaps, args.polish_steps
+        )
+        tempering = {
+            "cold_energy": float(values[cold]),
+            "polished_energy": polished_value,
+            "polished_accuracy": float(
+                (polished.cpu() == truth.cpu()).float().mean()
+            ),
+            "best_replica_accuracy": max(
+                float((state.cpu() == truth.cpu()).float().mean())
+                for state in states
+            ),
+        }
+        deepest = min(
+            [polished_value, float(values.min())]
+            + [item["energy"] for item in quenches]
+        )
+        entry = {
+            "true_energy": true_energy,
+            "random_mean": float(random_values.mean()),
+            "true_z": float(
+                (random_values.mean() - true_energy)
+                / random_values.std().clamp_min(1e-12)
+            ),
+            "descent_retention_diagnostic": retention,
+            "descent_from_truth_energy": kept_value,
+            "quenches": quenches,
+            "tempering": tempering,
+            "deepest_seen": deepest,
+            "margin_over_true": deepest / abs(true_energy) - true_energy / abs(true_energy),
+            "passes": bool(deepest >= true_energy - 1e-9),
+            "recovery_accuracy": tempering["polished_accuracy"],
+        }
+        report["energies"][name] = entry
+        print(
+            json.dumps(
+                {
+                    name: {
+                        key: entry[key]
+                        for key in (
+                            "true_energy",
+                            "deepest_seen",
+                            "passes",
+                            "recovery_accuracy",
+                            "descent_retention_diagnostic",
+                            "true_z",
+                        )
+                    }
+                }
+            )
+        )
+    write_json(args.output, report)
+    print(f"Wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()