"""geobasin_bed.py — the CE-replacement geometric arm, reconstructed and properly tested. #TAG:geobasin #TAG:ce_replacement #TAG:loss_campaign RECONSTRUCTION PROVENANCE (2026-07-25): the Oct '25 arm survives VERBATIM in the installed package — geofractal.losses.pure_geometric (PureGeometricLoss: attraction (1-s_y)^2 + repulsion sum(s_c^2) + margin hinge + range clamp; GeometricPrototypeLoss; HierarchicalGeometricLoss) and the GBC 4-factor compatibility head in geofractal/model/experiment_geometric_basin.py. L-031's "term forms NOT recorded" and L-032's "unbuilt" are both CLOSED. Losses are IMPORTED from the package, never rewritten (reuse law). THE TEST (the ce_fixedcode lesson: isolate the LOSS by holding the head identical): all arms share one cosine-anchor score head on the certified addr_msl64 read — s = normalize(feats); scores_c = (cos(s, A_c)+1)/2, A: (256 classes x 256), ~param-matched to the ce head (65,536 vs 65,792). CE arms consume cos*10 as logits (fixed scale, disclosed). ARMS: geo_ce_scores (CE on the SAME head - the decisive control) | geo_pure (verbatim, learned anchors) | geo_pure_frozen (L-108 cell) | geo_pure_norep (attraction+range ONLY - the absolute-only, doctrine-clean variant; repulsion +margin are roster-comparative terms) | geo_proto (verbatim, own projector - extra params disclosed) | geo_hier (nibble hierarchy 16x16 - bytes' natural coarse structure) | geo_hybrid (0.5 CE + 0.5 PureGeometric). PREREG (3 seeds; baselines ce 2.4769/acc .505): P1 the Oct'25 trade was -12% relative accuracy - geo_pure within -12% of geo_ce_scores matches history, parity overturns it, acc<0.30 refutes viability. P2 (L-108): learned anchors show a collapse signature vs frozen; falsifier: learned > frozen by >2 pts. P3 (absolute-beats-relative): norep >= pure - noise; falsifier: pure beats norep by >2 pts (the comparative terms would be load-bearing - a scope amendment to the law). P4: hybrid bpb within 0.15 of geo_ce_scores. Run: python tools/geobasin_bed.py --arm --seed N | --list """ import json import math import os import sys import time import zlib import torch import torch.nn as nn import torch.nn.functional as F def _root(): d = os.path.abspath(os.getcwd()) while True: if os.path.exists(os.path.join(d, "MANIFEST.md")): return d p = os.path.dirname(d) if p == d: return os.getcwd() d = p ROOT = _root() if os.path.join(ROOT, "tools") not in sys.path: sys.path.insert(0, os.path.join(ROOT, "tools")) import ar_differentiation_bed as bed # noqa: E402 from geofractal.losses.pure_geometric import ( # noqa: E402 GeometricPrototypeLoss, HierarchicalGeometricLoss, PureGeometricLoss) from loss_forms import prim_ce # noqa: E402 torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False DEV = "cuda" if torch.cuda.is_available() else "cpu" if DEV == "cuda": torch.cuda.set_per_process_memory_fraction(0.73) DATA_ROOT = os.environ.get("GEOLIP_DATA", "./data") RUNS = os.path.join(ROOT, "tools", "geobasin_runs") CE_SCALE = 10.0 def seed_for(name): return zlib.crc32(name.encode()) & 0x7FFFFFFF ARMS = ("geo_ce_scores", "geo_pure", "geo_pure_frozen", "geo_pure_norep", "geo_proto", "geo_hier", "geo_hybrid") class FastHierarchical(HierarchicalGeometricLoss): """Vectorizes ONLY the per-sample fine_mask loop (a 60x Python-loop slowdown at B*T=8192); every formula identical — parity-asserted against the verbatim original at build time.""" def forward(self, compatibility_scores, labels, mixed_labels=None, lam=None): if mixed_labels is not None: return super().forward(compatibility_scores, labels, mixed_labels, lam) B = compatibility_scores.shape[0] dev = compatibility_scores.device sl = labels // self.subclasses_per_super scores_r = compatibility_scores.view(B, self.num_superclasses, self.subclasses_per_super) ss = scores_r.sum(dim=2) coarse_correct = ss[torch.arange(B, device=dev), sl] ct = float(self.subclasses_per_super) coarse_loss = ((coarse_correct - ct) / ct).pow(2).mean() cmask = torch.ones_like(ss) cmask[torch.arange(B, device=dev), sl] = 0 coarse_rep = (ss * cmask).pow(2).sum(dim=1).mean() fine_correct = compatibility_scores[torch.arange(B, device=dev), labels] fine_loss = (1.0 - fine_correct).pow(2).mean() # vectorized fine_mask: ones on the label's superclass block, zero at # the label column (identical to the original's per-sample loop) cols = (sl * self.subclasses_per_super).unsqueeze(1) + \ torch.arange(self.subclasses_per_super, device=dev).unsqueeze(0) fmask = torch.zeros_like(compatibility_scores) fmask.scatter_(1, cols, 1.0) fmask[torch.arange(B, device=dev), labels] = 0 fine_rep = (compatibility_scores * fmask).pow(2).sum(dim=1).mean() consistency = F.relu(fine_correct * 2 - coarse_correct).mean() w_c = torch.sigmoid(self.coarse_weight) w_f = torch.sigmoid(self.fine_weight) return (w_c * (coarse_loss + 0.3 * coarse_rep) + w_f * (fine_loss + 0.3 * fine_rep) + 0.2 * consistency) class ScoreHead(nn.Module): """The shared cosine-anchor basin head. scores in [0,1] per class.""" def __init__(self, dim=256, classes=256, frozen=False, gen=None): super().__init__() A = torch.randn(classes, dim, generator=gen) if frozen: self.register_buffer("A", F.normalize(A, dim=-1)) else: self.A = nn.Parameter(A) def cos(self, feats): return F.normalize(feats, dim=-1) @ F.normalize(self.A, dim=-1).t() def forward(self, feats): return (self.cos(feats) + 1) / 2 def build(arm, seed): torch.manual_seed(seed_for(f"geobasin:{arm}:{seed}")) lm = bed.ByteLM("addr_msl64").to(DEV) lm.head = nn.Identity() # forward -> 256-d msl feats gh = torch.Generator().manual_seed(seed_for(f"geobasin-head:{seed}")) head = ScoreHead(frozen=(arm == "geo_pure_frozen"), gen=gh).to(DEV) aux = None if arm == "geo_proto": aux = GeometricPrototypeLoss(num_classes=256, prototype_dim=64).to(DEV) elif arm == "geo_hier": aux = FastHierarchical(num_classes=256, num_superclasses=16).to(DEV) # parity vs the VERBATIM original on a random batch (reference-check # pattern): the vectorization must be mathematics-identical ref = HierarchicalGeometricLoss(num_classes=256, num_superclasses=16).to(DEV) ref.load_state_dict(aux.state_dict()) gpar = torch.Generator().manual_seed(seed_for("hier-parity")) sc = torch.rand(16, 256, generator=gpar).to(DEV) yy = torch.randint(0, 256, (16,), generator=gpar).to(DEV) assert torch.allclose(aux(sc, yy), ref(sc, yy), atol=1e-6), \ "FastHierarchical diverged from the verbatim original" elif arm != "geo_ce_scores": aux = PureGeometricLoss() # stateless params = list(lm.parameters()) + list(head.parameters()) if aux is not None: params += list(aux.parameters()) return lm, head, aux, [p for p in params if p.requires_grad] def loss_of(arm, head, aux, feats, y): scores = head(feats).reshape(-1, 256) yy = y.reshape(-1) if arm == "geo_ce_scores": return prim_ce((scores * 2 - 1).reshape(*y.shape, 256) * CE_SCALE, y).mean() if arm == "geo_hybrid": ce = prim_ce((scores * 2 - 1).reshape(*y.shape, 256) * CE_SCALE, y).mean() return 0.5 * ce + 0.5 * PureGeometricLoss()(scores, yy) if arm == "geo_pure_norep": B = scores.shape[0] correct = scores[torch.arange(B, device=scores.device), yy] attraction = (1.0 - correct).pow(2).mean() # verbatim term 1 rng = F.relu(scores - 1.0).pow(2).mean() \ + F.relu(-scores).pow(2).mean() # verbatim term 4 return attraction + 0.1 * rng # comparative terms DROPPED return aux(scores, yy) # verbatim package losses @torch.no_grad() def evaluate(head, lm, va, g): lm.eval() tot_ce, tot_ok, n = 0.0, 0, 0 for _ in range(8): x, y = bed._batch(va, 32, 256, DEV, g) cos = head.cos(lm(x)) lp = F.log_softmax(cos * CE_SCALE, -1) tot_ce += float(-lp.gather(-1, y.unsqueeze(-1)).sum()) tot_ok += int((cos.argmax(-1) == y).sum()) n += y.numel() x, _ = bed._batch(va, 8, 256, DEV, g) _ = lm(x) v = lm.head_addr.vitals(lm.head_proj(lm._last_h).view( *lm._last_h.shape[:-1], lm.n_slots, 4)) lm.train() return (tot_ce / n) / math.log(2), tot_ok / n, v def train_arm(arm, seed, steps=2000): assert arm in ARMS, f"unknown arm {arm}" os.makedirs(RUNS, exist_ok=True) tr, va = bed._wikitext_bytes(DATA_ROOT) lm, head, aux, params = build(arm, seed) g = torch.Generator().manual_seed(seed_for(f"geobasin-data:{arm}:{seed}")) ge = torch.Generator().manual_seed(seed_for("geobasin-eval")) opt = torch.optim.Adam(params, lr=3e-4, weight_decay=0.0) t0 = time.time() # anchor-collapse gauge: pairwise |cos| spread of A at start vs end def anchor_spread(): An = F.normalize(head.A.detach(), dim=-1) pc = (An @ An.t()).abs() off = pc[~torch.eye(256, dtype=torch.bool, device=pc.device)] return float(off.mean()) sp0 = anchor_spread() for step in range(steps): x, y = bed._batch(tr, 32, 256, DEV, g) opt.zero_grad(set_to_none=True) L = loss_of(arm, head, aux, lm(x), y) L.backward() opt.step() if step == 10 and DEV == "cuda": print(f"[{arm} s{seed}] step10 loss {float(L.detach()):.4f} " f"peak {torch.cuda.max_memory_allocated()/2**30:.2f}GB", flush=True) bpb, acc, vit = evaluate(head, lm, va, ge) rec = {"arm": arm, "seed": seed, "steps": steps, "bpb": round(bpb, 4), "decoded_acc": round(acc, 4), "anchor_abs_cos_mean_init": round(sp0, 4), "anchor_abs_cos_mean_final": round(anchor_spread(), 4), "vitals": vit, "wall_s": round(time.time() - t0, 1), "n_params_extra": sum(p.numel() for p in aux.parameters()) if aux is not None and any(True for _ in aux.parameters()) else 0} out = os.path.join(RUNS, f"{arm}_s{seed}_t{steps}.jsonl") with open(out, "a", encoding="utf-8") as f: f.write(json.dumps(rec) + "\n") print(f"[DONE {arm} s{seed}] bpb {bpb:.4f} acc {acc:.4f} " f"anchors |cos| {sp0:.3f}->{rec['anchor_abs_cos_mean_final']:.3f} " f"({rec['wall_s']}s)", flush=True) return rec if __name__ == "__main__": import argparse ap = argparse.ArgumentParser() ap.add_argument("--arm", default=None) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--steps", type=int, default=2000) ap.add_argument("--list", action="store_true") a, _ = ap.parse_known_args() if a.list or not a.arm: print("arms:", " ".join(ARMS)) sys.exit(0) train_arm(a.arm, a.seed, a.steps)