ncgc-mila/GeoNexus-OLMo-20M

GeoNexus — OLMo-20M ablation tier

Full checkpoint archive for the 20M-parameter tier of the GeoNexus project: an empirical study of whether cross-source gradient agreement — measured in the outer optimizer's own geometry — causes better downstream generalization, as claimed by the Nexus line of work (arXiv 2604.09258, "Same Pretraining Loss, Better Downstream Generalization via Common Minima").

This tier is the mechanism-and-ablation lab of the project: small enough that a full arm trains in ~40 minutes, so hypotheses are tested with matched controls rather than single runs. The companion repo GeoNexus-130M-pilots holds the larger-scale pilots.

• Code: https://github.com/sevendaystoglory/GeoNexus
• Metrics: Weights & Biases project llm_jp_pp/GeoNexus-Olmo (every run here streams there under the same name)

What is in a checkpoint

Each run directory contains a final checkpoint plus periodic archives:

<run-name>/ckpt.pt                # final (step 1525 for the 400M-token config)
<run-name>/ckpt_step000500.pt     # periodic archives, every 500 steps
<run-name>/ckpt_step001000.pt     # plus dense archives every 50 steps after step 1000
...

Every file is a plain torch.save dict — including full optimizer state, so any run can be resumed or re-diagnosed exactly:

key	contents
model	state_dict of the GPT model
opt	optimizer state for the whole optimizer set (Muon + AdamW, or AdamW alone)
loader	data-loader position (shard, epoch, RNG) for bit-exact resumption
opt_name	"muon" or "adamw"
step	training step
model_cfg, train_cfg	the exact configs used

import torch
ck = torch.load("o20-dg-base-s42/ckpt.pt", map_location="cpu", weights_only=False)
print(ck["step"], ck["opt_name"], ck["model_cfg"])
# model definition: model.py in the GitHub repo

Note these are research checkpoints, not a released language model — a 20M-parameter model trained on 400M tokens is a measurement instrument, not something to deploy.

Model and training setup

architecture	Llama-style decoder: 8 layers, d_model 256, 8 heads, d_ff 1024, RoPE, RMSNorm, untied embeddings
parameters	24.8M total / 16.6M non-embedding
context	1024 tokens, vocab 32000
data	SlimPajama, 9 sources (arxiv, book, c4, commoncrawl, freelaw, github, pubmedcentral, stackexchange, wikipedia)
budget	400M tokens = 1526 steps × 262144 tokens (batch 8 × grad-accum 32 × 1024)
optimizer	Muon (lr 0.02, momentum 0.95, NS-5) on hidden matrices + AdamW (lr 1e-3) on embeddings/norms; WSD schedule
precision	bf16 autocast, fp32 master weights
seed	42 throughout (arms are compared through matched controls, not seed replication)

Held-out evaluation uses in-distribution validation on the 9 training sources plus two out-of-distribution domains (FreeLaw and PubMedCentral held out of the training mixture in the DG cohorts).

Run naming

o20-<cohort>-<variant>-<dose>-s42, where s42 is the seed and the dose encodes the arm's hyperparameter (g002 = γ 0.002, lam4e5 = λ 4e-5, r025 = constant ratio 0.25, and so on). Suffix -b1000 marks the extended-budget (1000-step) variant of an earlier cohort; -e100 marks an earlier eval cadence; -swa3/-swa6 are tail-averaged (checkpoint-averaged) models built from the dense archives of their parent run.

Cohorts

prefix	what it studies
o20-muon-*, o20-adamw-*	the original tier-1 grid: baselines and Nexus variants (l2, adamgeo, geo-v2) under each outer optimizer
o20-jgap-*	Nexus-D: optimizing the Jensen gap directly via a Danskin two-point probe
o20-specgd-*, o20-signgd-*	outer-optimizer controls (spectral-descent, sign-descent) — do the effects survive a change of outer geometry?
o20-dg-*	domain-generalization cohort: base/v2 anchors, srcwindow batching, and literature baselines (SAM, GroupDRO, PCGrad)
o20-attr-*	attribution controls for the v2 effect: sign-flip, norm-matched random walk, ℓ2-geometry walk, gradient rescale, momentum retunes (μ.96/.97)
o20-nexusv2e-*	Nexus-v2E: explicit optimizer of v2's implicit objective L̄ − c⟨ḡ, sg(q̄)⟩, at fixed λ, at constant force/gradient ratio, and as a walk-faithful path-ordered twin (p025)
o20-jw-*	Jensen walks: conform, leap, tight, conformleash — inner walks designed so their implicit objective targets the Jensen gap rather than the transfer score
o20-v1-*, o20-v2frozen-*, o20-sb-*	mechanism surgery: v1 (both-polar) reference, the frozen-direction walk (steering ablation), and IsoLoCo-style σ̄ outer rescaling

Why these runs exist

Headline findings this tier produced (details, tables, and the diagnostic scripts are in the GitHub repo under results/o20/ and analysis/):

1. The v2 variant is a real effect. Walking in Muon geometry while accumulating raw gradients reaches lower pretraining loss than a matched baseline, with downstream OOD improving in proportion.
2. Optimizing the Jensen gap directly fails, and not because of an implementation flaw — the objective's gradient is a thin residual of two near-parallel large vectors whose net direction climbs the loss.
3. The geometric quantities are severable from generalization. An explicit optimizer of v2's implicit objective, correctly dosed, reproduces the entire geometric signature (cross-source alignment, transfer score, bound tightness, gap, minima proximity) at matched loss — and yields no OOD benefit. Arms that close the Jensen gap hardest are the worst models in the cohort.
4. What v2 actually does is closed-loop steering. Freezing the walk's direction field (v2frozen) preserves the displacement budget but multiplies the curvature response ~3.6×; projecting the displacements onto the top Hessian eigendirections shows the sequential walk suppresses its sharp-subspace components 6–30× per mode. The benefit lives in the walk's higher-order, trajectory-coupled structure, not in any first-order objective it appears to descend.

These are findings at this scale (20M parameters, 400M tokens, one seed with matched controls). The Nexus paper's own claims are made at 1B+; nothing here settles that regime, and the 130M pilots repo is the intermediate rung.

Reproducing a diagnostic

With the GitHub repo checked out and a checkpoint downloaded:

# fine-tune proximity + sharpness of the per-source minima
python analysis/closeness_sharp.py --ckpt o20-dg-base-s42/ckpt.pt --steps 100 --probe-iters 8

# Jensen gap, transfer score T, and bound tightness on fixed probe batches
python analysis/transfer_probe.py --ckpts o20-dg-base-s42/ckpt.pt

# cross-source gradient similarity in raw / Muon / transfer geometries
python analysis/report_diag.py --ckpts o20-dg-base-s42/ckpt.pt --diag-batches 16

License

Apache-2.0. Training data is SlimPajama (see its own terms). If these checkpoints are useful in your work, a link back to the GitHub repository is appreciated.