"""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()