diff options
Diffstat (limited to 'worldalign/energy_coupling.py')
| -rw-r--r-- | worldalign/energy_coupling.py | 205 |
1 files changed, 205 insertions, 0 deletions
diff --git a/worldalign/energy_coupling.py b/worldalign/energy_coupling.py new file mode 100644 index 0000000..cf5b859 --- /dev/null +++ b/worldalign/energy_coupling.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +import torch.nn.functional as F + +from .common import read_json, seed_everything, write_json +from .energy import ( + log_sinkhorn, + relation_field_energy, + retrieval_metrics, + standardized_relation, +) +from .io import load_feature_pair, select_rows + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + 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") + parser.add_argument("--split", choices=["val", "test"], default="test") + parser.add_argument("--samples", type=int, default=256) + parser.add_argument( + "--reservoir", + choices=["same_set_shuffled", "unpaired"], + default="same_set_shuffled", + ) + parser.add_argument("--steps", type=int, default=1_200) + parser.add_argument("--lr", type=float, default=0.12) + parser.add_argument("--temperature-start", type=float, default=0.8) + parser.add_argument("--temperature-end", type=float, default=0.12) + parser.add_argument("--conditional-weight", type=float, default=0.08) + parser.add_argument("--entropy-weight-start", type=float, default=0.005) + parser.add_argument("--entropy-weight-end", type=float, default=0.055) + parser.add_argument("--device", default="cuda:1") + parser.add_argument("--seed", type=int, default=20260729) + parser.add_argument( + "--output", default="artifacts/energy_coupling_test.pt" + ) + parser.add_argument( + "--metrics-output", default="artifacts/energy_coupling_test.json" + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + seed_everything(args.seed) + manifest = read_json(args.manifest) + vision, text, vision_lookup, text_lookup = load_feature_pair( + args.vision, args.text + ) + rows = manifest[args.split][: args.samples] + visual = select_rows( + vision["features"], vision_lookup, rows + ).to(args.device) + paired_text = select_rows( + text["features"], text_lookup, rows + ).to(args.device) + text_population = select_rows( + text["features"], text_lookup, manifest["text_only_train"] + ).to(args.device) + if args.text_orbits: + state = torch.load( + args.text_orbits, map_location="cpu", weights_only=False + ) + lookup = { + int(row): index for index, row in enumerate(state["rows"]) + } + orbit_mean = F.normalize( + state["features"].float().mean(1), dim=-1 + ) + paired_text = select_rows( + orbit_mean, lookup, rows + ).to(args.device) + text_population = select_rows( + orbit_mean, lookup, manifest["text_only_train"] + ).to(args.device) + + generator = torch.Generator(device=args.device).manual_seed(args.seed) + target: torch.Tensor | None = None + if args.reservoir == "same_set_shuffled": + permutation = torch.randperm( + len(paired_text), generator=generator, device=args.device + ) + anchors = paired_text[permutation] + target = torch.argsort(permutation) + else: + anchor_indices = torch.randperm( + len(text_population), + generator=generator, + device=args.device, + )[: len(visual)] + anchors = text_population[anchor_indices] + + initial_permutation = torch.randperm( + len(visual), generator=generator, device=args.device + ) + logits = torch.nn.Parameter( + 0.02 + * torch.randn( + len(visual), + len(visual), + generator=generator, + device=args.device, + ) + ) + with torch.no_grad(): + logits[torch.arange(len(visual), device=args.device), initial_permutation] += 3 + visual_relation, visual_standardized = standardized_relation(visual) + optimizer = torch.optim.Adam([logits], lr=args.lr) + history: list[dict] = [] + + for step in range(args.steps + 1): + progress = min(step / max(args.steps, 1), 1.0) + temperature = max( + args.temperature_end, + args.temperature_start + + progress + * (args.temperature_end - args.temperature_start), + ) + coupling = log_sinkhorn(logits, temperature, iterations=15) + particles = F.normalize(coupling @ anchors, dim=-1) + relation, conditional = relation_field_energy( + visual_relation, visual_standardized, particles + ) + entropy = -( + coupling * coupling.clamp_min(1e-12).log() + ).sum(-1).mean() + entropy_weight = ( + args.entropy_weight_start + + progress + * (args.entropy_weight_end - args.entropy_weight_start) + ) + loss = ( + relation + + args.conditional_weight * conditional + + entropy_weight * entropy + ) + if step % 100 == 0 or step == args.steps: + record = { + "step": step, + "total": float(loss.detach()), + "relation": float(relation.detach()), + "conditional": float(conditional.detach()), + "entropy": float(entropy.detach()), + "temperature": temperature, + "paired_evaluation_only": retrieval_metrics( + particles.detach(), paired_text + ), + } + if target is not None: + prediction = coupling.argmax(-1) + record["exact_coupling_evaluation_only"] = { + "accuracy": float((prediction == target).float().mean()), + "true_mass": float( + coupling[ + torch.arange( + len(coupling), device=args.device + ), + target, + ].mean() + ), + } + history.append(record) + print(json.dumps(record)) + if step == args.steps: + break + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_([logits], 5.0) + optimizer.step() + + result = { + "protocol": ( + "The optimizer sees only frozen within-modality states and " + "energy. Pair identity and paired text are used only for " + "evaluation diagnostics." + ), + "mode": "mass_conserving_language_state_coupling", + "reservoir": args.reservoir, + "split": args.split, + "rows": rows, + "args": vars(args), + "history": history, + } + state = { + **result, + "anchors": anchors.detach().cpu(), + "final_coupling": coupling.detach().cpu(), + "final_particles": particles.detach().cpu(), + } + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + torch.save(state, args.output) + write_json(args.metrics_output, result) + print(f"Wrote {args.output} and {args.metrics_output}") + + +if __name__ == "__main__": + main() |