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

diff --git a/worldalign/synth_towers.py b/worldalign/synth_towers.py
new file mode 100644
index 0000000..15221ac
--- /dev/null
+++ b/worldalign/synth_towers.py
@@ -0,0 +1,348 @@
+"""From-scratch unimodal towers for the synthetic closed world.
+
+Vision: a compact ViT trained with InfoNCE over natural orbit positives --
+two renders of the same scene, which differ exactly by the world's
+continuous nuisance. No flips (they would erase left/right semantics) and
+no color jitter (color is world content). Scene identity within the
+vision-only split is unimodal metadata.
+
+Text: a small word-level causal LM trained on the captions of the
+text-only split. Both towers therefore learn the same world from disjoint
+scenes and separate modalities, with dialable capacity.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torch import nn
+
+from .common import read_json, seed_everything
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--side", choices=["vision", "text"], required=True)
+    parser.add_argument(
+        "--objective", choices=["infonce", "simmim", "hybrid", "data2vec"], default="infonce"
+    )
+    parser.add_argument("--mask-ratio", type=float, default=0.6)
+    parser.add_argument("--recon-weight", type=float, default=25.0)
+    parser.add_argument("--ema-decay", type=float, default=0.999)
+    parser.add_argument("--data-dir", default="artifacts/synth_v0")
+    parser.add_argument("--dim", type=int, default=192)
+    parser.add_argument("--depth", type=int, default=6)
+    parser.add_argument("--heads", type=int, default=3)
+    parser.add_argument("--text-dim", type=int, default=256)
+    parser.add_argument("--text-heads", type=int, default=4)
+    parser.add_argument("--patch", type=int, default=16)
+    parser.add_argument("--context", type=int, default=80)
+    parser.add_argument("--epochs", type=int, default=80)
+    parser.add_argument("--batch-size", type=int, default=256)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    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()
+
+
+class Block(nn.Module):
+    def __init__(self, dim: int, heads: int) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attention = nn.MultiheadAttention(dim, heads, batch_first=True)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim)
+        )
+
+    def forward(
+        self, x: torch.Tensor, mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        normed = self.norm1(x)
+        attended, _ = self.attention(
+            normed, normed, normed, attn_mask=mask, need_weights=False
+        )
+        x = x + attended
+        return x + self.mlp(self.norm2(x))
+
+
+class VisionTower(nn.Module):
+    def __init__(self, image_size: int, patch: int, dim: int, depth: int, heads: int) -> None:
+        super().__init__()
+        self.patch = patch
+        self.patch_embed = nn.Conv2d(3, dim, kernel_size=patch, stride=patch)
+        tokens = (image_size // patch) ** 2
+        self.cls = nn.Parameter(torch.zeros(1, 1, dim))
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, dim))
+        self.positions = nn.Parameter(torch.zeros(1, tokens + 1, dim))
+        nn.init.trunc_normal_(self.positions, std=0.02)
+        nn.init.trunc_normal_(self.cls, std=0.02)
+        nn.init.trunc_normal_(self.mask_token, std=0.02)
+        self.blocks = nn.ModuleList(Block(dim, heads) for _ in range(depth))
+        self.norm = nn.LayerNorm(dim)
+        self.head = nn.Sequential(
+            nn.Linear(dim, dim), nn.GELU(), nn.Linear(dim, 128)
+        )
+        self.reconstruction = nn.Linear(dim, patch * patch * 3)
+        self.feature_head = nn.Linear(dim, dim)
+
+    def encode(
+        self, pixels: torch.Tensor, mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        x = self.patch_embed(pixels).flatten(2).transpose(1, 2)
+        if mask is not None:
+            x = torch.where(mask[..., None], self.mask_token.expand_as(x), x)
+        x = torch.cat([self.cls.expand(len(x), -1, -1), x], dim=1)
+        x = x + self.positions
+        for block in self.blocks:
+            x = block(x)
+        return self.norm(x)
+
+    def forward(self, pixels: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        tokens = self.encode(pixels)
+        return tokens[:, 0], self.head(tokens[:, 0])
+
+
+class TextTower(nn.Module):
+    def __init__(self, vocab: int, dim: int, depth: int, heads: int, context: int) -> None:
+        super().__init__()
+        self.embed = nn.Embedding(vocab, dim)
+        self.positions = nn.Parameter(torch.zeros(1, context, dim))
+        nn.init.trunc_normal_(self.positions, std=0.02)
+        self.blocks = nn.ModuleList(Block(dim, heads) for _ in range(depth))
+        self.norm = nn.LayerNorm(dim)
+        self.context = context
+
+    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
+        length = tokens.shape[1]
+        x = self.embed(tokens) + self.positions[:, :length]
+        mask = torch.triu(
+            torch.full((length, length), float("-inf"), device=tokens.device), 1
+        )
+        for block in self.blocks:
+            x = block(x, mask)
+        return self.norm(x)
+
+    def logits(self, hidden: torch.Tensor) -> torch.Tensor:
+        return hidden @ self.embed.weight.T
+
+
+def load_image(path: Path) -> torch.Tensor:
+    with Image.open(path) as image:
+        array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
+    return torch.from_numpy(array).permute(2, 0, 1)
+
+
+def patchify(pixels: torch.Tensor, patch: int) -> torch.Tensor:
+    batch, channels, height, width = pixels.shape
+    grid = height // patch
+    x = pixels.reshape(batch, channels, grid, patch, grid, patch)
+    return x.permute(0, 2, 4, 3, 5, 1).reshape(batch, grid * grid, -1)
+
+
+def train_vision(args: argparse.Namespace) -> None:
+    manifest = read_json(Path(args.data_dir, "manifest.json"))
+    rows = manifest["vision_only_train"]
+    views = manifest["visual_views"]
+    image_dir = Path(manifest["image_dir"])
+    model = VisionTower(
+        manifest["image_size"], args.patch, args.dim, args.depth, args.heads
+    ).to(args.device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.05)
+    tokens = (manifest["image_size"] // args.patch) ** 2
+    if args.objective in ("infonce", "hybrid"):
+        samples_per_epoch = len(rows)
+    else:
+        samples_per_epoch = len(rows) * views
+    teacher = None
+    if args.objective == "data2vec":
+        import copy
+
+        teacher = copy.deepcopy(model)
+        for parameter in teacher.parameters():
+            parameter.requires_grad_(False)
+    steps_per_epoch = max(1, samples_per_epoch // args.batch_size)
+    schedule = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=args.epochs * steps_per_epoch
+    )
+    rng = random.Random(args.seed)
+    for epoch in range(args.epochs):
+        total, count = 0.0, 0
+        if args.objective in ("infonce", "hybrid"):
+            order = rows.copy()
+            rng.shuffle(order)
+            for start in range(0, len(order) - args.batch_size + 1, args.batch_size):
+                batch_rows = order[start : start + args.batch_size]
+                pairs = []
+                for row in batch_rows:
+                    first, second = rng.sample(range(views), 2)
+                    pairs.append(
+                        load_image(image_dir / f"scene{row:06d}_v{first}.png")
+                    )
+                    pairs.append(
+                        load_image(image_dir / f"scene{row:06d}_v{second}.png")
+                    )
+                pixels = torch.stack(pairs).to(args.device)
+                _, projected = model(pixels)
+                projected = F.normalize(projected, dim=-1)
+                logits = projected @ projected.T / args.temperature
+                logits.fill_diagonal_(float("-inf"))
+                targets = torch.arange(len(projected), device=args.device) ^ 1
+                loss = F.cross_entropy(logits, targets)
+                if args.objective == "hybrid":
+                    mask = (
+                        torch.rand(len(pixels), tokens, device=args.device)
+                        < args.mask_ratio
+                    )
+                    encoded = model.encode(pixels, mask=mask)
+                    predicted = model.reconstruction(encoded[:, 1:][mask])
+                    target_patches = patchify(pixels, args.patch)[mask]
+                    loss = loss + args.recon_weight * F.mse_loss(
+                        predicted, target_patches
+                    )
+                optimizer.zero_grad(set_to_none=True)
+                loss.backward()
+                optimizer.step()
+                schedule.step()
+                total += float(loss)
+                count += 1
+        else:
+            jobs = [(row, view) for row in rows for view in range(views)]
+            rng.shuffle(jobs)
+            for start in range(0, len(jobs) - args.batch_size + 1, args.batch_size):
+                batch = jobs[start : start + args.batch_size]
+                pixels = torch.stack(
+                    [
+                        load_image(image_dir / f"scene{row:06d}_v{view}.png")
+                        for row, view in batch
+                    ]
+                ).to(args.device)
+                mask = (
+                    torch.rand(len(batch), tokens, device=args.device)
+                    < args.mask_ratio
+                )
+                encoded = model.encode(pixels, mask=mask)
+                if args.objective == "simmim":
+                    predicted = model.reconstruction(encoded[:, 1:][mask])
+                    target = patchify(pixels, args.patch)[mask]
+                    loss = F.mse_loss(predicted, target)
+                else:
+                    with torch.no_grad():
+                        reference = teacher.encode(pixels)[:, 1:]
+                        reference = F.layer_norm(
+                            reference, reference.shape[-1:]
+                        )
+                    predicted = model.feature_head(encoded[:, 1:][mask])
+                    loss = F.smooth_l1_loss(predicted, reference[mask])
+                optimizer.zero_grad(set_to_none=True)
+                loss.backward()
+                optimizer.step()
+                schedule.step()
+                if teacher is not None:
+                    with torch.no_grad():
+                        for student_p, teacher_p in zip(
+                            model.parameters(), teacher.parameters()
+                        ):
+                            teacher_p.lerp_(student_p, 1.0 - args.ema_decay)
+                total += float(loss)
+                count += 1
+        if epoch % 5 == 0 or epoch == args.epochs - 1:
+            print(json.dumps({"epoch": epoch, "loss": total / max(count, 1)}))
+    torch.save(
+        {"model": model.state_dict(), "args": vars(args), "side": "vision"},
+        args.output,
+    )
+    print(f"Wrote {args.output}")
+
+
+def tokenize(caption: str, vocab: dict[str, int]) -> list[int]:
+    tokens = caption.replace(",", " ").replace(".", " .").split()
+    return [vocab["<bos>"]] + [vocab[token] for token in tokens]
+
+
+def build_vocab(manifest: dict) -> dict[str, int]:
+    words = list(manifest["vocabulary"]) + ["."]
+    vocab = {"<pad>": 0, "<bos>": 1}
+    for word in sorted(set(words)):
+        vocab.setdefault(word, len(vocab))
+    return vocab
+
+
+def train_text(args: argparse.Namespace) -> None:
+    manifest = read_json(Path(args.data_dir, "manifest.json"))
+    captions = read_json(Path(args.data_dir, "captions.json"))["captions"]
+    vocab = build_vocab(manifest)
+    sentences = [
+        tokenize(caption, vocab)
+        for row in manifest["text_only_train"]
+        for caption in captions[row]
+    ]
+    model = TextTower(
+        len(vocab), args.text_dim, args.depth, args.text_heads, args.context
+    ).to(args.device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
+    rng = random.Random(args.seed)
+    steps_per_epoch = max(1, len(sentences) // args.batch_size)
+    schedule = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=args.epochs * steps_per_epoch
+    )
+    for epoch in range(args.epochs):
+        rng.shuffle(sentences)
+        total, count = 0.0, 0
+        for start in range(0, len(sentences) - args.batch_size + 1, args.batch_size):
+            batch = sentences[start : start + args.batch_size]
+            longest = min(args.context, max(len(s) for s in batch))
+            tokens = torch.zeros(len(batch), longest, dtype=torch.long)
+            for index, sentence in enumerate(batch):
+                clipped = sentence[:longest]
+                tokens[index, : len(clipped)] = torch.tensor(clipped)
+            tokens = tokens.to(args.device)
+            hidden = model(tokens)
+            logits = model.logits(hidden[:, :-1])
+            targets = tokens[:, 1:]
+            loss = F.cross_entropy(
+                logits.reshape(-1, logits.shape[-1]),
+                targets.reshape(-1),
+                ignore_index=0,
+            )
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            schedule.step()
+            total += float(loss)
+            count += 1
+        if epoch % 5 == 0 or epoch == args.epochs - 1:
+            print(json.dumps({"epoch": epoch, "loss": total / max(count, 1)}))
+    torch.save(
+        {
+            "model": model.state_dict(),
+            "vocab": vocab,
+            "args": vars(args),
+            "side": "text",
+        },
+        args.output,
+    )
+    print(f"Wrote {args.output}")
+
+
+def main() -> None:
+    args = parse_args()
+    seed_everything(args.seed)
+    if args.side == "vision":
+        train_vision(args)
+    else:
+        train_text(args)
+
+
+if __name__ == "__main__":
+    main()