x-small

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

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 in benchmarks/ppl_trend.csv): code/math/prose perplexity across the 12-layer evals (steps 63,663 → 110,938 → 123,141).
• Throughput (benchmarks/throughput.png, data in benchmarks/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.