ncgc-mila/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)

folder	recipe	final val	notes
adamw-s42-1b	AdamW baseline	2.971
nexus-l2-s42-1b	AdamW + Nexus Φ=l2, γ=0.05	2.951	null on geometry metrics
nexus-adamgeo-s42-1b	AdamW + Φ=D^(−1/4), γ=0.05	2.965	positive: +60% alignment, closeness −1% (7/9 tasks), same loss
muon-s42-1b	Muon baseline	2.921	best OOD of the baselines
nexus-muon-l2-s42-1b	Muon + Φ=l2, γ=0.05	2.897	trades OOD for in-dist; collapses alignment
nexus-muon-geo-s42-1b	Muon + Φ=polar, γ=0.05 (v1)	3.138	negative: spectral-flattening pathology
nexus-muon-geo-g002-s42-1b	v1 at γ=0.002	3.132	γ-insensitive damage (mechanism study)
nexus-muon-geo-v2-g002-s42-1b	Muon GeoNexus v2 (decoupled Φ: polar inner displacement, raw pseudo-grad), γ=0.002	2.908	same-loss vs Muon; alignment +24%; closeness −5.8% on 9/9 tasks, unconfounded
adamw-blk-s42-2k, nexus-l2-blk-s42-2k, nexus-adamgeo-blk-s42-2k	single-source-microbatch (Algorithm-1 style) ablation, 0.5B tokens	3.13 / 3.10 / 3.11	source batching strengthens Nexus's in-dist effect

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.