diff options
Diffstat (limited to 'worldalign/synth_triangle_gate.py')
| -rw-r--r-- | worldalign/synth_triangle_gate.py | 296 |
1 files changed, 296 insertions, 0 deletions
diff --git a/worldalign/synth_triangle_gate.py b/worldalign/synth_triangle_gate.py new file mode 100644 index 0000000..35bc2b2 --- /dev/null +++ b/worldalign/synth_triangle_gate.py @@ -0,0 +1,296 @@ +"""Third-order (triangle) relational invariants for the synthetic world. + +Counterfeits so far satisfy pairwise statistics: they are wrong sections +that look flat when probed along edges. The holonomy analogue on a +relation field is a triangle: for a node triple the product-like +statistic T[a,b,c] combines three edges at once, so preserving pairwise +marginals no longer suffices. Each node participates in O(N^2) triangles +rather than O(N) edges, which multiplies the constraint count without +touching the states. + +The triangle energy sums squared cross-modal differences of standardized +triangle tensors over a fixed random triple sample; the sample is drawn +once from node indices only. Swap deltas are evaluated exactly on the +triples touching the swapped nodes. Hidden pairs score orderings only. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +import torch.nn.functional as F +from tqdm import tqdm + +from .common import read_json, seed_everything, write_json +from .synth_cc_battery import ( + component_descriptors, + moment_field, + onehot_descriptors, + phrase_bow_sets, +) +from .synth_towers import load_image + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + 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("--triples", type=int, default=200000) + parser.add_argument("--pair-weight", type=float, default=1.0) + parser.add_argument("--triangle-weight", type=float, default=1.0) + parser.add_argument("--random-perms", type=int, default=200) + parser.add_argument("--transposition-samples", type=int, default=20000) + parser.add_argument("--descent-restarts", type=int, default=3) + parser.add_argument("--descent-max-steps", type=int, default=3000) + parser.add_argument("--descent-proposals", type=int, default=2048) + parser.add_argument("--replicas", type=int, default=8) + parser.add_argument("--tempering-rounds", type=int, default=20000) + 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("--device", default="cuda:3") + parser.add_argument("--seed", type=int, default=20260731) + parser.add_argument( + "--output", default="artifacts/synth_v0/triangle_gate.json" + ) + return parser.parse_args() + + +def standardized(field: torch.Tensor) -> torch.Tensor: + mask = ~torch.eye(len(field), dtype=torch.bool, device=field.device) + values = field[mask] + out = (field - values.mean()) / values.std().clamp_min(1e-9) + return out.masked_fill(~mask, 0.0) + + +class TriangleEnergy: + """Pairwise-plus-triangle energy over a fixed triple sample.""" + + def __init__( + self, + text: torch.Tensor, + visual: torch.Tensor, + triples: torch.Tensor, + pair_weight: float, + triangle_weight: float, + ) -> None: + self.text = text + self.visual = visual + self.triples = triples + self.pair_weight = pair_weight + self.triangle_weight = triangle_weight + self.size = len(visual) + self.mask = ~torch.eye(self.size, dtype=torch.bool, device=visual.device) + self.pair_count = int(self.mask.sum()) + a, b, c = triples[:, 0], triples[:, 1], triples[:, 2] + self.visual_triangle = ( + visual[a, b] * visual[b, c] * visual[a, c] + ) + # Triples touching each node, for local delta evaluation. + self.touch = [[] for _ in range(self.size)] + for index, (x, y, z) in enumerate(triples.tolist()): + self.touch[x].append(index) + self.touch[y].append(index) + self.touch[z].append(index) + self.touch = [ + torch.tensor(items, device=visual.device, dtype=torch.long) + for items in self.touch + ] + + def pairwise(self, permutation: torch.Tensor) -> torch.Tensor: + fields = self.text[permutation][:, permutation] + return ((fields - self.visual) ** 2)[self.mask].sum() / self.pair_count + + def triangle(self, permutation: torch.Tensor) -> torch.Tensor: + fields = self.text[permutation][:, permutation] + a, b, c = self.triples[:, 0], self.triples[:, 1], self.triples[:, 2] + text_triangle = fields[a, b] * fields[b, c] * fields[a, c] + return ((text_triangle - self.visual_triangle) ** 2).mean() + + def total(self, permutation: torch.Tensor) -> float: + return float( + self.pair_weight * self.pairwise(permutation) + + self.triangle_weight * self.triangle(permutation) + ) + + def swap_delta(self, permutation: torch.Tensor, p: int, q: int) -> float: + """Exact delta for swapping positions p and q.""" + before_pair = self._local_pair(permutation, p, q) + indices = torch.cat([self.touch[p], self.touch[q]]).unique() + before_triangle = self._local_triangle(permutation, indices) + trial = permutation.clone() + trial[[p, q]] = trial[[q, p]] + after_pair = self._local_pair(trial, p, q) + after_triangle = self._local_triangle(trial, indices) + return float( + self.pair_weight * (after_pair - before_pair) / self.pair_count + + self.triangle_weight + * (after_triangle - before_triangle) + / len(self.triples) + ) + + def _local_pair( + self, permutation: torch.Tensor, p: int, q: int + ) -> torch.Tensor: + rows = permutation[[p, q]] + fields = self.text[rows][:, permutation] + visual_rows = self.visual[[p, q]] + difference = (fields - visual_rows) ** 2 + difference[0, p] = 0.0 + difference[1, q] = 0.0 + # Rows and columns are symmetric; count row contributions twice and + # subtract the doubly counted (p, q) pair once. + total = 2.0 * difference.sum() + return total - 2.0 * difference[0, q] + + def _local_triangle( + self, permutation: torch.Tensor, indices: torch.Tensor + ) -> torch.Tensor: + triples = self.triples[indices] + a, b, c = triples[:, 0], triples[:, 1], triples[:, 2] + pa, pb, pc = permutation[a], permutation[b], permutation[c] + text_triangle = ( + self.text[pa, pb] * self.text[pb, pc] * self.text[pa, pc] + ) + return ((text_triangle - self.visual_triangle[indices]) ** 2).sum() + + +def build_fields(args: argparse.Namespace, rows: list[int], manifest: dict) -> tuple: + image_dir = Path(manifest["image_dir"]) + per_view = [] + for view in range(args.vision_views): + raw = [] + for row in tqdm(rows, desc=f"cc v{view}"): + sprites, _ = component_descriptors( + load_image(image_dir / f"scene{row:06d}_v{view}.png"), + args.merge_distance, + ) + raw.append(sprites) + sets = [F.normalize(v, dim=-1) for v in onehot_descriptors(raw)] + per_view.append(moment_field(sets)) + visual_field = torch.stack(per_view).mean(0) + captions = read_json(Path(args.data_dir, "captions.json"))["captions"] + text_sets = phrase_bow_sets(rows, captions, manifest["vocabulary"]) + return visual_field, moment_field(text_sets) + + +def main() -> None: + args = parse_args() + seed_everything(args.seed) + 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(rows) + generator = torch.Generator().manual_seed(args.seed) + hidden = torch.randperm(size, generator=generator) + truth = torch.argsort(hidden).to(device) + text = standardized(text_field[hidden][:, hidden].double().to(device)).float() + visual = standardized(visual_field.double().to(device)).float() + + triples = torch.randint(0, size, (args.triples, 3), generator=generator) + triples = triples[ + (triples[:, 0] != triples[:, 1]) + & (triples[:, 1] != triples[:, 2]) + & (triples[:, 0] != triples[:, 2]) + ].to(device) + energy = TriangleEnergy( + text, visual, triples, args.pair_weight, args.triangle_weight + ) + + true_total = energy.total(truth) + randoms = [] + for _ in range(args.random_perms): + permutation = torch.argsort(torch.rand(size, generator=generator)).to(device) + randoms.append(energy.total(permutation)) + randoms = torch.tensor(randoms) + gate_a = { + "true": true_total, + "random_mean": float(randoms.mean()), + "true_z": float((randoms.mean() - true_total) / randoms.std().clamp_min(1e-12)), + } + + improving = 0 + checked = 0 + for _ in range(min(args.transposition_samples, 4000)): + p = int(torch.randint(0, size, (1,), generator=generator)) + q = int(torch.randint(0, size, (1,), generator=generator)) + if p == q: + continue + checked += 1 + improving += energy.swap_delta(truth, p, q) < 0 + gate_b = { + "sampled": checked, + "improving_fraction": improving / max(checked, 1), + } + + def descent(start: torch.Tensor) -> dict: + current = start.clone() + for _ in range(args.descent_max_steps): + best_delta, best_pair = 0.0, None + for _ in range(64): + p = int(torch.randint(0, size, (1,), generator=generator)) + q = int(torch.randint(0, size, (1,), generator=generator)) + if p == q: + continue + delta = energy.swap_delta(current, p, q) + if delta < best_delta: + best_delta, best_pair = delta, (p, q) + if best_pair is None: + break + p, q = best_pair + current[[p, q]] = current[[q, p]] + return { + "accuracy": float((current.cpu() == truth.cpu()).float().mean()), + "energy": energy.total(current), + } + + from_true = descent(truth) + restarts = [ + descent(torch.argsort(torch.rand(size, generator=generator)).to(device)) + for _ in range(args.descent_restarts) + ] + best_random = min(r["energy"] for r in restarts) + + report = { + "protocol": ( + "Pairwise-plus-triangle energy on moment set-kernel fields; " + "triples sampled from node indices only. Hidden pairs score " + "orderings." + ), + "samples": size, + "triples": len(triples), + "weights": { + "pair": args.pair_weight, + "triangle": args.triangle_weight, + }, + "gate_a": gate_a, + "gate_b": gate_b, + "descent_from_true": from_true, + "descent_from_random": restarts, + "verdict": { + "true_energy": true_total, + "best_random_descent": best_random, + "counterfeit_found": bool( + best_random < true_total + and min(r["accuracy"] for r in restarts) < 0.5 + ), + "descent_keeps": from_true["accuracy"], + "true_z": gate_a["true_z"], + "improving_fraction": gate_b["improving_fraction"], + }, + } + print(json.dumps({"verdict": report["verdict"]})) + write_json(args.output, report) + print(f"Wrote {args.output}") + + +if __name__ == "__main__": + main() |