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