diff options
Diffstat (limited to 'worldalign/blind_recovery.py')
| -rw-r--r-- | worldalign/blind_recovery.py | 476 |
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() |