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

diff --git a/worldalign/energy_infer.py b/worldalign/energy_infer.py
new file mode 100644
index 0000000..d5a67b6
--- /dev/null
+++ b/worldalign/energy_infer.py
@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from .common import read_json, seed_everything, write_json
+from .energy import (
+    projection_quantile_target,
+    prototype_manifold_energy,
+    relation_field_energy,
+    retrieval_metrics,
+    sliced_distribution_energy,
+    standardized_relation,
+)
+from .io import load_feature_pair, select_rows
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    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",
+        help=(
+            "Optional multi-description feature cache. When supplied, each "
+            "language particle is the normalized mean of its observation "
+            "orbit."
+        ),
+    )
+    parser.add_argument("--prototypes", default="artifacts/gw.pt")
+    parser.add_argument("--split", choices=["val", "test"], default="test")
+    parser.add_argument("--samples", type=int, default=256)
+    parser.add_argument("--steps", type=int, default=100)
+    parser.add_argument("--lr", type=float, default=0.03)
+    parser.add_argument("--projections", type=int, default=128)
+    parser.add_argument("--relation-weight-start", type=float, default=0.4)
+    parser.add_argument("--relation-weight-end", type=float, default=2.0)
+    parser.add_argument("--conditional-weight-start", type=float, default=0.05)
+    parser.add_argument("--conditional-weight-end", type=float, default=0.2)
+    parser.add_argument("--distribution-weight", type=float, default=80.0)
+    parser.add_argument("--manifold-weight", type=float, default=1.0)
+    parser.add_argument("--noise", type=float, default=0.01)
+    parser.add_argument(
+        "--shuffle-visual-energy",
+        action="store_true",
+        help=(
+            "Evaluation control: permute image particles before constructing "
+            "the energy while leaving evaluation rows unchanged."
+        ),
+    )
+    parser.add_argument("--device", default="cuda:1")
+    parser.add_argument("--seed", type=int, default=20260729)
+    parser.add_argument(
+        "--output", default="artifacts/energy_free_test.pt"
+    )
+    parser.add_argument(
+        "--metrics-output", default="artifacts/energy_free_test.json"
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    seed_everything(args.seed)
+    manifest = read_json(args.manifest)
+    vision, text, vision_lookup, text_lookup = load_feature_pair(
+        args.vision, args.text
+    )
+    rows = manifest[args.split][: args.samples]
+    visual = select_rows(
+        vision["features"], vision_lookup, rows
+    ).to(args.device)
+    paired_text = select_rows(
+        text["features"], text_lookup, rows
+    ).to(args.device)
+    text_population = select_rows(
+        text["features"], text_lookup, manifest["text_only_train"]
+    ).to(args.device)
+    if args.text_orbits:
+        orbit_state = torch.load(
+            args.text_orbits, map_location="cpu", weights_only=False
+        )
+        orbit_lookup = {
+            int(row): index
+            for index, row in enumerate(orbit_state["rows"])
+        }
+        orbit_mean = F.normalize(
+            orbit_state["features"].float().mean(1), dim=-1
+        )
+        paired_text = select_rows(
+            orbit_mean, orbit_lookup, rows
+        ).to(args.device)
+        text_population = select_rows(
+            orbit_mean, orbit_lookup, manifest["text_only_train"]
+        ).to(args.device)
+    if args.shuffle_visual_energy:
+        control_generator = torch.Generator(
+            device=args.device
+        ).manual_seed(args.seed + 10_000)
+        visual = visual[
+            torch.randperm(
+                len(visual),
+                generator=control_generator,
+                device=args.device,
+            )
+        ]
+    prototype_state = torch.load(
+        args.prototypes, map_location="cpu", weights_only=False
+    )
+    prototypes = prototype_state["text_centers"].to(args.device)
+    if prototypes.shape[-1] != text_population.shape[-1]:
+        raise ValueError(
+            "Prototype dimension does not match text features; use the GW "
+            "cache built from the selected text backbone."
+        )
+
+    generator = torch.Generator(device=args.device).manual_seed(args.seed)
+    initial_indices = torch.randperm(
+        len(text_population),
+        generator=generator,
+        device=args.device,
+    )[: len(visual)]
+    initial = text_population[initial_indices].clone()
+    initial = initial + args.noise * torch.randn(
+        initial.shape, generator=generator, device=args.device
+    )
+    particles = torch.nn.Parameter(F.normalize(initial, dim=-1))
+    directions, target_quantiles = projection_quantile_target(
+        text_population,
+        particles=len(particles),
+        projections=args.projections,
+        generator=generator,
+    )
+    visual_relation, visual_standardized = standardized_relation(visual)
+    optimizer = torch.optim.Adam([particles], lr=args.lr)
+
+    def energy_values() -> tuple[torch.Tensor, ...]:
+        relation, conditional = relation_field_energy(
+            visual_relation, visual_standardized, particles
+        )
+        distribution = sliced_distribution_energy(
+            particles, directions, target_quantiles
+        )
+        manifold = prototype_manifold_energy(particles, prototypes)
+        return relation, conditional, distribution, manifold
+
+    history: list[dict] = []
+    initial_cpu = F.normalize(particles.detach(), dim=-1).cpu()
+    for step in range(args.steps + 1):
+        relation, conditional, distribution, manifold = energy_values()
+        progress = min(step / max(args.steps, 1), 1.0)
+        relation_weight = (
+            args.relation_weight_start
+            + progress
+            * (args.relation_weight_end - args.relation_weight_start)
+        )
+        conditional_weight = (
+            args.conditional_weight_start
+            + progress
+            * (
+                args.conditional_weight_end
+                - args.conditional_weight_start
+            )
+        )
+        loss = (
+            relation_weight * relation
+            + conditional_weight * conditional
+            + args.distribution_weight * distribution
+            + args.manifold_weight * manifold
+        )
+        if step % 20 == 0 or step == args.steps:
+            history.append(
+                {
+                    "step": step,
+                    "total": float(loss.detach()),
+                    "relation": float(relation.detach()),
+                    "conditional": float(conditional.detach()),
+                    "distribution": float(distribution.detach()),
+                    "manifold": float(manifold.detach()),
+                    "paired_evaluation_only": retrieval_metrics(
+                        particles.detach(), paired_text
+                    ),
+                }
+            )
+            print(json.dumps(history[-1]))
+        if step == args.steps:
+            break
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_([particles], 2.0)
+        optimizer.step()
+        with torch.no_grad():
+            particles.copy_(F.normalize(particles, dim=-1))
+
+    with torch.no_grad():
+        oracle_relation, oracle_conditional = relation_field_energy(
+            visual_relation, visual_standardized, paired_text
+        )
+        oracle_distribution = sliced_distribution_energy(
+            paired_text, directions, target_quantiles
+        )
+        oracle_manifold = prototype_manifold_energy(
+            paired_text, prototypes
+        )
+    result = {
+        "protocol": (
+            "No cross-modal map and no image-text pair is used by the energy "
+            "or optimizer. Paired text is loaded only for trajectory and "
+            "oracle diagnostics; the final step is fixed by CLI arguments."
+        ),
+        "mode": "free_language_latent_particles",
+        "split": args.split,
+        "rows": rows,
+        "vision_model": vision["model"],
+        "text_model": text["model"],
+        "text_observation": (
+            "multi-description orbit mean"
+            if args.text_orbits
+            else "single description"
+        ),
+        "args": vars(args),
+        "history": history,
+        "oracle_energy_components": {
+            "relation": float(oracle_relation),
+            "conditional": float(oracle_conditional),
+            "distribution": float(oracle_distribution),
+            "manifold": float(oracle_manifold),
+        },
+    }
+    state = {
+        **result,
+        "initial_particles": initial_cpu,
+        "final_particles": F.normalize(
+            particles.detach(), dim=-1
+        ).cpu(),
+    }
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+    torch.save(state, args.output)
+    write_json(args.metrics_output, result)
+    print(f"Wrote {args.output} and {args.metrics_output}")
+
+
+if __name__ == "__main__":
+    main()