diff options
Diffstat (limited to 'worldalign/synth_fast_gate.py')
| -rw-r--r-- | worldalign/synth_fast_gate.py | 345 |
1 files changed, 345 insertions, 0 deletions
diff --git a/worldalign/synth_fast_gate.py b/worldalign/synth_fast_gate.py new file mode 100644 index 0000000..54b6891 --- /dev/null +++ b/worldalign/synth_fast_gate.py @@ -0,0 +1,345 @@ +"""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() |