guestros/s2-moe-checkpoints

S²-MoE — research checkpoints

Trained weights behind the paper Evolving Sparse Spiking Mixture-of-Experts (S²-MoE).

⚠️ These are research artifacts, not a deployable language model. They are small (25M–150M params), trained on TinyStories / FineWeb-Edu, and exist to reproduce the paper's experiments — not to generate good text. They are not transformers-compatible: loading them requires the custom model code in the GitHub repository (custom LIF spiking neurons, a gated-linear-attention mixer, and a grouped-GEMM expert kernel). The checkpoint blobs store weights only (optimizer stripped) and do not store the config — use the exact ModelConfig documented per file below.

Files

File	Backs	Params (total / ~active)	ModelConfig
student_n8.pt	Routing demo + n-expert sweep (Exp 6/7)	49M / ~25M	d_model=512, d_hidden=512, n_layers=4, n_experts=8, d_shared=256
student_n32_fineweb.pt	Headline: distilled MoE beats its dense teacher −5% at ~1/21 active FFN compute (Exp 28)	150M / ~28M	d_model=512, d_hidden=512, n_layers=4, n_experts=32, d_shared=256
teacher_n32_fineweb.pt	The size-matched dense teacher the n=32 student is compared to (Exp 28)	82M	DenseLM, d_model=512, d_ff=16384, n_layers=4 (see src/baseline.py)
pool_128grow.pt	Autonomous growth: sleep phase grew the pool to 128 experts (Exp 58)	151M / ~26M	d_model=512, d_hidden=128, n_layers=4, n_experts=128, d_shared=256

All use mixer="gla", n_heads=4, vocab_size=8192 (byte-level BPE, data/fineweb-bpe.json in the repo). "Active" params ≈ the Top-1 routed slice (one expert of the pool) plus the always-on shared core — that sparsity is the whole point.

How to load

# in a clone of https://github.com/JustinGuese/S3-MoE-llm
import torch
from src.backbone import ModelConfig, S2MoE

cfg = ModelConfig(vocab_size=8192, d_model=512, d_hidden=512, n_layers=4,
                  n_experts=32, d_shared=256, mixer="gla", n_heads=4, dispatch="loop")
model = S2MoE(cfg)
model.load_state_dict(torch.load("student_n32_fineweb.pt", map_location="cpu")["model"],
                      strict=False)
model.eval()

The dense teacher loads via src.baseline.DenseLM instead. To reproduce the paper numbers, use train_lm.py / eval_140m.py in the repo. An interactive demo of the routing and energy behaviour is at https://huggingface.co/spaces/guestros/s2-moe-demo.

Honest scope

The efficiency result (less active compute per token) is shape-determined and robust; the quality edge is budget- and corpus-dependent and holds in a distilled setup at this scale — it can reverse on harder data at a larger token budget. See the repo's README.md, RESEARCH.md, and the negative-results catalog (eight tombstones) for the full, honest picture.

• Code + paper: https://github.com/JustinGuese/S3-MoE-llm
• Interactive demo: https://huggingface.co/spaces/guestros/s2-moe-demo
• DOI: https://doi.org/10.5281/zenodo.20846758