GeoNexus-130M-pilots

GeoNexus — 130M pilot checkpoints

Checkpoints from GeoNexus: an experimental study of geometry-aware common-minima regularization (Nexus, arXiv:2604.09258) — testing whether the Nexus inner-loop transform Φ should measure cross-source gradient agreement in the outer optimizer's geometry (AdamW: D^(−1/4); Muon: the polar factor) instead of raw-gradient cosine.

Model

Llama-style decoder, ~130M non-embedding parameters: 32 layers, d_model 512, d_ff 2048 (SwiGLU), 8 heads, RoPE, RMSNorm, tied embeddings, seq len 1024, Llama-2 tokenizer (vocab 32000). Trained on SlimPajama-6B (official splits, 7 sources, natural proportions, pooled single-pass mixed loader), 1B tokens (≈3,815 steps at global batch 256), WSD schedule, bf16, seed 42. OOD eval: PubMed Central + FreeLaw (Pile mirror).

Checkpoints (final step unless noted)

Key findings, in brief: (1) on AdamW, the geometry-matched Φ=D^(−1/4) moves every common-minima diagnostic (cross-source gradient alignment, OOD alignment, parameter closeness ||θ−θ*_T||²) in the predicted direction at matched pretraining loss, where the standard l2 Φ is a null. (2) On Muon, the paper's reported Nexus incompatibility reproduces, and its mechanism was isolated: applying polar per microbatch flattens the pseudo-gradient spectrum (stable rank 62 vs 2), the outer polar then over-spreads the update (stable rank 333 vs ideal 245, −17% descent alignment/step). (3) Decoupling the two roles of Φ — polar for the inner displacement (the second-order agreement coupling), raw accumulation for the pseudo-gradient — removes the pathology entirely: v2 is loss-neutral vs Muon while moving alignment and closeness the right way.

Usage

Checkpoints are raw torch.save dicts from the project's training loop (not transformers-format). Each ckpt.pt contains model (state_dict), opt (optimizer state), loader (data-stream position), model_cfg / train_cfg (dataclass dicts), opt_name, step.

import torch
ckpt = torch.load("muon-s42-1b/ckpt.pt", map_location="cpu")
# model.py from the GeoNexus code repo defines GPT/GPTConfig
model = GPT(GPTConfig(**{k: v for k, v in ckpt["model_cfg"].items()}))
model.load_state_dict(ckpt["model"])

500-step intermediate archives are included, not just the finals: 152 checkpoints across 20 runs (~207 GB). Every file carries full optimizer state and data-loader position, so any run can be resumed mid-trajectory or re-diagnosed at any archived step -- which is what the trajectory analyses in the code repo rely on.

The training/eval code (model.py, train.py, the Nexus inner loop, diagnostics, the paired measurement instruments under analysis/, and a 7-benchmark eval suite) lives in the project's code repository, along with the raw measurement logs under results/ and a HANDOFF.md that summarizes what these runs concluded.

Research artifacts — no fitness for any purpose implied; 130M-scale pilot models, not useful as general-purpose LMs.