You can deploy this model and start earning money today!

gaijingchu/
momentum-streams-checkpoints

Name: momentum-streams-checkpoints
Availability: InStock

language-modeling

transformer

optimizer-inspired-architecture

mit

Views

No views yet

Momentum Streams for Optimizer-Inspired Transformers — Checkpoints

Pretrained checkpoints for the paper Momentum Streams for Optimizer-Inspired Transformers (Gai, Huang, Wu, 2026).

Code: https://github.com/gaijingchu/Momentum-Streams-for-Optimizer-Inspired-Transformers

One best.pt per variant. All models share the 12-layer, 12-head, d=768, context-1024 pre-norm Transformer backbone with GPT-2 BPE (50,304 vocab). Each best.pt is the model state dict (no optimizer state) saved at the lowest validation cross-entropy during pretraining.

OpenWebText (30k steps, paper Table 1)

Path	Val loss	Notes
`vanilla_owt/best.pt`	3.0078	VanillaTransformer (Muon + AdamW recipe)
`vanilla_adamw_owt/best.pt`	3.0103	Vanilla, pure-AdamW ablation
`adam_owt/best.pt`	2.9911	AdamFormer
`adamw_owt/best.pt`	2.9883	AdamWFormer
`tmm_owt/best.pt`	2.9342	TMMFormer (paper best)
`hb_owt/best.pt`	—	HBFormer (factorial ablation)
`rmsprop_owt/best.pt`	—	RMSPropFormer (factorial ablation)

OpenWebText with WSD / SAM / SAWD recipes

Path	Val loss
`tmm_sam_owt/best.pt`	2.940
`tmm_wsd_owt/best.pt`	2.924
`tmm_sawd_owt/best.pt`	—

TinyStories (10k steps)

Path	Val loss
`vanilla_ts/best.pt`	1.1569
`adam_ts/best.pt`	1.153
`adamw_ts/best.pt`	1.147
`tmm_ts/best.pt`	1.128
`muon_ts/best.pt`	1.1503
`soap_ts/best.pt`	1.1431
`ortho_ts/best.pt`	—
`hb_ts/best.pt`	—
`rmsprop_ts/best.pt`	—
`tmm_sam_ts/best.pt`	1.079
`tmm_wsd_ts/best.pt`	1.086
`tmm_sawd_ts/best.pt`	1.082

Download

python

from huggingface_hub import hf_hub_download
path = hf_hub_download(
    repo_id="gaijingchu/momentum-streams-checkpoints",
    filename="tmm_owt/best.pt",
    local_dir="./checkpoints_cache",
)

Loading

python

import torch
from formers.tmm.model import TMMFormer  # from the GitHub repo

model = TMMFormer(n_layers=12, n_heads=12, d_model=768, vocab_size=50304)
state = torch.load("tmm_owt/best.pt", map_location="cpu")
model.load_state_dict(state)
model.eval()

For other variants, use the matching class: VanillaTransformer, AdamFormer, AdamWFormer, MuonFormer, OrthoFormer, RMSPropFormer, ShampooFormer, SOAPFormer, HBFormer (under formers.<variant>.model).

Citation

bibtex

@article{gai2026momentum,
  title  = {Momentum Streams for Optimizer-Inspired Transformers},
  author = {Gai, Jingchu and Huang, Nai-Chieh and Wu, Jiayun},
  year   = {2026}
}

gaijingchu/momentum-streams-checkpoints