From a62cf4d2a99b4a7985c61b2a7feb92a82a8218b7 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Sat, 1 Aug 2026 14:10:03 -0500 Subject: World Alignment: unpaired cross-modal correspondence by relational identifiability Method: scene states are sets of part states; relation fields are built within each modality and are invariant to how each side labels its own features; the cross-modal bridge is a coupling searched under an energy that is a closed-form functional of one matrix; solving is spectral initialisation followed by exact local refinement. Evidence: in a procedurally generated closed world, blind recovery of a hidden image-caption correspondence reaches 95.3% at 256 scenes against 0.39% chance, and the recovered pairs transfer to 200 held-out scenes at 93.0% exact retrieval with random-pair and shuffled-image controls at or near chance. Cross-modal value correspondence is derived from disjoint corpora rather than declared. On Visual Genome the field correlation reaches 0.656 against the 0.9 that polynomial recovery needs, with the deficit attributed away from segmentation and discretisation. Protocol: no image-text pair enters any objective, optimiser, initialisation, or model selection; hidden pairs score orderings only. Co-Authored-By: Claude --- worldalign/ricci_control.py | 258 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 258 insertions(+) create mode 100644 worldalign/ricci_control.py (limited to 'worldalign/ricci_control.py') diff --git a/worldalign/ricci_control.py b/worldalign/ricci_control.py new file mode 100644 index 0000000..3091b22 --- /dev/null +++ b/worldalign/ricci_control.py @@ -0,0 +1,258 @@ +"""Ricci-flow control for the on-manifold assignment gate. + +Hypothesis under test: a discrete Ricci flow that smooths each modality's +relational geometry before matching could repair the local ordering that +static relation fields fail. The flow is run independently per modality +with identical hyperparameters; hidden pairs never touch the flow or the +energy and only score orderings, as in the base gate. + +Two geometries are gated per condition: heat-kernel channels on the raw +kNN graph (diffusion without flow) and the same construction after +Ollivier-Ricci weight evolution. Differences between them are attributable +to the flow itself rather than to the diffusion representation. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import ot +import torch +from scipy.sparse import csr_matrix +from scipy.sparse.csgraph import shortest_path + +from .common import seed_everything, write_json +from .manifold_gate import ( + gate_a_global_ranking, + gate_b_transpositions, + load_flickr, + load_vg, + standardize_relation, + steepest_descent, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", choices=["flickr", "vg"], default="flickr") + 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("--text-mode", choices=["single", "orbit_mean"], default="orbit_mean") + 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("--vg-bundle-channels", action="store_true") + 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("--neighbors", type=int, default=16) + parser.add_argument("--flow-iterations", type=int, default=10) + parser.add_argument("--flow-step", type=float, default=0.4) + parser.add_argument("--lazy-alpha", type=float, default=0.5) + parser.add_argument("--heat-times", default="1.0,4.0") + parser.add_argument("--random-perms", type=int, default=300) + parser.add_argument("--derangement-samples", type=int, default=50) + parser.add_argument("--descent-restarts", type=int, default=2) + parser.add_argument("--descent-max-steps", type=int, default=200000) + parser.add_argument("--descent-objective", default="mse", choices=["mse", "m30_total"]) + parser.add_argument("--descent-verify-top", type=int, default=64) + parser.add_argument("--device", default="cpu") + parser.add_argument("--seed", type=int, default=20260729) + parser.add_argument("--output", default="artifacts/manifold_gate/ricci_control.json") + parser.add_argument("--trajectory-output") + return parser.parse_args() + + +def knn_edges(distance: np.ndarray, k: int) -> list[tuple[int, int]]: + order = distance.argsort(-1)[:, 1 : k + 1] + edges = { + (min(i, int(j)), max(i, int(j))) + for i in range(len(distance)) + for j in order[i] + } + return sorted(edges) + + +def graph_apsp(size: int, weights: dict[tuple[int, int], float]) -> np.ndarray: + rows, cols, vals = [], [], [] + for (i, j), w in weights.items(): + rows.extend((i, j)) + cols.extend((j, i)) + vals.extend((w, w)) + graph = csr_matrix((vals, (rows, cols)), shape=(size, size)) + return shortest_path(graph, method="D", directed=False) + + +def ollivier_ricci_apsp( + distance: np.ndarray, args: argparse.Namespace, iterations: int +) -> np.ndarray: + """All-pairs geodesics after Ollivier-Ricci weight evolution. + + Lazy uniform neighbor measures, W1 ground costs from current geodesics, + multiplicative weight update w <- w * (1 - step * kappa), total edge + mass renormalized each iteration. iterations=0 gives the un-flowed + graph geometry for the diffusion-only control. + """ + edges = knn_edges(distance, args.neighbors) + weights = {edge: max(float(distance[edge]), 1e-9) for edge in edges} + total = sum(weights.values()) + neighbor_map: dict[int, list[int]] = {} + for i, j in edges: + neighbor_map.setdefault(i, []).append(j) + neighbor_map.setdefault(j, []).append(i) + apsp = graph_apsp(len(distance), weights) + for _ in range(iterations): + updated: dict[tuple[int, int], float] = {} + for i, j in edges: + support_i = [i] + neighbor_map[i] + support_j = [j] + neighbor_map[j] + mass_i = np.full(len(support_i), (1 - args.lazy_alpha) / len(neighbor_map[i])) + mass_i[0] = args.lazy_alpha + mass_j = np.full(len(support_j), (1 - args.lazy_alpha) / len(neighbor_map[j])) + mass_j[0] = args.lazy_alpha + ground = apsp[np.ix_(support_i, support_j)] + if not np.isfinite(ground).all(): + finite_max = apsp[np.isfinite(apsp)].max() + ground = np.where(np.isfinite(ground), ground, 2.0 * finite_max) + wasserstein = ot.emd2(mass_i, mass_j, ground) + geodesic = max(float(apsp[i, j]), 1e-9) + curvature = 1.0 - wasserstein / geodesic + updated[(i, j)] = max(1e-9, weights[(i, j)] * (1.0 - args.flow_step * curvature)) + scale = total / sum(updated.values()) + weights = {edge: w * scale for edge, w in updated.items()} + apsp = graph_apsp(len(distance), weights) + return apsp + + +def geometry_channels( + apsp: np.ndarray, heat_times: tuple[float, ...] +) -> torch.Tensor: + """Standardized relation channels of a flowed geometry. + + Channel 0 is the negative geodesic field; the rest are heat kernels of + the normalized Laplacian of a geodesic-scale affinity. + """ + finite = apsp[np.isfinite(apsp) & (apsp > 0)] + scale = np.median(finite) + capped = np.where(np.isfinite(apsp), apsp, finite.max() * 2.0) + affinity = np.exp(-capped / scale) + np.fill_diagonal(affinity, 1.0) + degree = affinity.sum(-1) + normalized = affinity / np.sqrt(degree[:, None] * degree[None, :]) + values, vectors = np.linalg.eigh((normalized + normalized.T) / 2.0) + laplacian_eigen = 1.0 - values + channels = [torch.from_numpy(-capped).double()] + for t in heat_times: + heat = (vectors * np.exp(-t * laplacian_eigen)) @ vectors.T + channels.append(torch.from_numpy(heat).double()) + return torch.stack( + [standardize_relation(channel)[0] for channel in channels] + ) + + +def run_gates( + text_channels: torch.Tensor, + visual_channels: torch.Tensor, + args: argparse.Namespace, + generator: torch.Generator, +) -> dict: + text_relation = text_channels[0] + visual_relation = visual_channels[0] + report = { + "gate_a": gate_a_global_ranking( + text_channels, visual_channels, text_relation, visual_relation, args, generator + ), + "gate_b": gate_b_transpositions( + text_channels, visual_channels, text_relation, visual_relation, None + ), + "descent_from_true": steepest_descent( + text_channels, + visual_channels, + text_relation, + visual_relation, + torch.arange(len(visual_relation)), + args, + ), + } + report["descent_from_true"].pop("final_permutation", None) + report["descent_from_random"] = [] + for _ in range(args.descent_restarts): + start = torch.argsort(torch.rand(len(visual_relation), generator=generator)) + result = steepest_descent( + text_channels, visual_channels, text_relation, visual_relation, start, args + ) + result.pop("final_permutation", None) + result.pop("trajectory", None) + report["descent_from_random"].append(result) + return report + + +def main() -> None: + args = parse_args() + seed_everything(args.seed) + data = load_flickr(args) if args.dataset == "flickr" else load_vg(args) + heat_times = tuple(float(t) for t in args.heat_times.split(",")) + states = { + "visual": torch.nn.functional.normalize( + data["visual_views"].double().mean(1), dim=-1 + ), + "text": torch.nn.functional.normalize( + data["text_views"].double().mean(1), dim=-1 + ), + } + distances = { + side: (1.0 - state @ state.T).clamp_min(0.0).numpy() + for side, state in states.items() + } + report: dict = { + "protocol": ( + "Each modality's kNN geometry evolves independently under " + "Ollivier-Ricci flow; the identical heat-kernel channels are " + "gated with and without the flow. Hidden pairs score orderings " + "only." + ), + "meta": {**data["meta"], "flow": vars(args)}, + "conditions": {}, + } + for label, iterations in ( + ("diffusion_no_flow", 0), + ("ricci_flow", args.flow_iterations), + ): + channels = {} + for side in ("visual", "text"): + apsp = ollivier_ricci_apsp(distances[side], args, iterations) + channels[side] = geometry_channels(apsp, heat_times) + generator = torch.Generator().manual_seed(args.seed) + result = run_gates(channels["text"], channels["visual"], args, generator) + report["conditions"][label] = result + print( + json.dumps( + { + label: { + "true_z_mse": result["gate_a"]["random"]["mse"]["true_z"], + "improving_fraction": result["gate_b"]["improving_fraction"], + "descent_keeps": result["descent_from_true"]["final_accuracy"], + "counterfeit_found": any( + r["final_objective"] + < result["gate_a"]["true"]["mse"] + and r["final_accuracy"] < 0.5 + for r in result["descent_from_random"] + ), + } + } + ) + ) + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + write_json(args.output, report) + print(f"Wrote {args.output}") + + +if __name__ == "__main__": + main() -- cgit v1.2.3