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.
- Domains: easy — code_easy, math_easy, reformat_easy, synth_easy, web_easy; medium — code_medium, math_medium, synth_medium, web_medium; hard — code_gold, code_hard, gold_hard, math_hard, synth_hard, web_gold, web_hard.
- Training token budget (nominal, full 64 × 2048 batches): 131,072 tokens/step; ~17.99B at step 137,260; ~131.07B at 1,000,000 steps. Loss-masked curriculum tokens are not subtracted.
- Hardware limit: single RTX 4070 Ti 12GB; micro-batch 8 × grad-accum 8 (eff 64) with gradient checkpointing, memory-efficient block-wise attention, and expandable CUDA segments to stay under the 12GB ceiling.
- 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.
Benchmarks
- Quality trend (
benchmarks/ppl_trend.png, data inbenchmarks/ppl_trend.csv): code/math/prose perplexity across the 12-layer evals (steps 63,663 → 110,938 → 123,141). - Throughput (
benchmarks/throughput.png, data inbenchmarks/tput_release.csv): released 12-layer weights, eager forward+backward with grad-checkpointing, bf16, seq 2048, RTX 4070 Ti 12GB. B=8 OOM because the GPU was shared with another training job at bench time; the release run used B=8 on a dedicated GPU.
License
MIT (code). Weights released for research use.