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:

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

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

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.