H-JEPA-SSM-MoE — hybrid Mamba-2 / local-attention / sparse-MoE backbone (+ char-LM demo checkpoint)
An oracle-verified, CPU-trainable reference implementation of a post-Transformer stack — Mamba-2 (SSD) selective state-space layers interleaved 9:1 with sliding-window attention (Samba-style), sparse Mixture-of-Experts FFNs (dropless, noisy top-2, Switch load-balancing), plus a JEPA self-supervised head (EMA target encoder + VICReg).
This repo ships:
| File | What it is |
|---|---|
lm_machado.pt (~6.4 MB) |
Demo checkpoint: 1.59M-param character-level LM trained 18 min on a CPU on Memórias Póstumas de Brás Cubas (Machado de Assis, 1881, public domain) |
hjepa/ + train_*.py + tests/ |
The full implementation (custom code — not transformers-compatible) |
PAPER.md |
Technical report: architecture, hand-derived adjoint backward, verification methodology, CPU benchmarks |
train_log.txt |
Exact training log of the released checkpoint |
Headline numbers (16-core consumer CPU, no GPU)
- Mamba-2/SSD trains 7.3× faster than Mamba-1/S6 on CPU (3,550 vs 484 chars/s, identical loss curves, same parameter count). Both implementations included.
- Demo LM: ln(99) = 4.60 → 1.66 val cross-entropy (ppl 5.28/char) in 2,500 steps.
- Generation uses the recurrent O(1)-per-token path (SSM state + windowed KV cache): throughput is flat in sequence length (≈52 chars/s on the reference CPU).
- Every subsystem is tested against an independent oracle (9 tests, max errors 10⁻⁷–10⁻⁵), including the analytic backward of the selective scan vs autograd.
Sample (temperature 0.8, prompt in bold)
Ao verme que barbante de ver cá. E depois donravel, como uma mulher, creia se ir, com ella casada, a noiva de dama; vinha constituiu-lhe por esse rapido, inclinei-a muito, — repouso religiosa...
(1.59M params trained on one novel: it learns the corpus's 1881 orthography — "ella", "idéa" — morphology and short-range phrase structure, not literary coherence.)
Use
git clone https://huggingface.co/dnnxzz/h-jepa-ssm-moe
cd h-jepa-ssm-moe
python3 -m venv .venv && .venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu
.venv/bin/python tests/self_test.py # verification suite
mkdir -p out && cp lm_machado.pt out/
.venv/bin/python interactive.py # streaming REPL (chat-style)
.venv/bin/python train_text.py sample --prompt "Eu era " --tokens 400
.venv/bin/python train_text.py # retrain from scratch (~18 min CPU)
.venv/bin/python train_synthetic.py --preset tiny # JEPA self-supervised loop
Intended use & limitations
Educational/reference. Toy scale (single corpus, single seed, no SOTA or downstream
claims — see PAPER.md §5). The corpus is public domain; code is MIT.
Developed by dnnxzz in collaboration with Claude (Anthropic). Source: https://github.com/c7bc/h-jepa-ssm-moe