SAN x-small
Simple Attention Network — a ~31M-parameter custom LLM, a faithful PyTorch port of the
needle architecture (arXiv:2607.18363). Trained locally on an RTX 4070 Ti 12 GB.
This repo contains the released weights for the 12-layer configuration.
Files
| file | what |
|---|---|
pytorch_model.bin | model state_dict (load via san_model.SimpleAttentionNetwork.load_state_dict) |
san_latest.pt | raw training checkpoint (step 137,260): {step, loss, model_state_dict, optimizer_state_dict, config} |
config.json | architecture hyperparameters |
tokenizer/ | SmolLM2-135M tokenizer (vocab 49152) |
san_model.py, san_triton.py | model definition (needed to load the weights) |
load_example.py | minimal load + forward example |
Load
from san_model import SimpleAttentionNetwork, SANConfig
import torch
cfg = SANConfig(num_layers=12)
model = SimpleAttentionNetwork(cfg).eval()
sd = torch.load("pytorch_model.bin", map_location="cpu")
model.load_state_dict(sd)
See load_example.py. Full training/eval code: GitHub
kenpeter/x-small.
Architecture (released config)
- params ~31M · d_model 384 · 12 layers · 8 heads / 4 KV heads (GQA)
- vocab 49152 (SmolLM2-135M) · context 2048 · bfloat16
- components: ZCRMSNorm, HadamardMLP, GQA + RoPE + learned gate, Engram n-gram KV memory, Multi-Lane Hyper-Connections (MHC, Sinkhorn routing), MTP head, tied embeddings
Training & eval
- Curriculum over 16 domains; target 100w (1,000,000) steps; eff batch 64.
- Last training step 137,260 (of 1,000,000).
- Last eval (step 123,141) perplexity: mean 23.74 (code 10.25 / math 29.62 / prose 31.34). Perplexity is the usable quality signal at this scale; sample text loops regardless.
- The flat training loss near ~2.0 is curriculum-mix masking, not model stall.
License
MIT (code). Weights released for research use.