AxiomicLabs/GPT-S2-5M
GPT-S2-5M
GPT-S2-5M is the latest entry in our GPT-S small-model family, now based on the T-X4 architecture with the all new XSA refresh gate. trained from scratch.
The model takes the #1 spot on the [Open SLM Leaderboard](https://huggingface.co/spaces/AxiomicLabs/Open_SLM_Leaderboard), dethroning SLM-10M while being almost 1/2 the parameters.
It combines RoPE + RMSNorm + SwiGLU + exclusive grouped-query attention, and the refresh gate re-injects the original token embedding back into the residual stream conditioned on the (detached) attention output to conunteract diluted token identities common in deep XSA models.
Architecture
Config
vocab_size = 4,096
hidden_size = 192
num_layers = 9
num_heads = 6
num_kv_heads = 2
head_dim = 32
intermediate = 672
block_size = 512
rope_theta = 2,500
inject_layers = [4, 8]
refresh_kernel = 9The XSA refresh gate
The keystone of the T-X4 architecture is the injection layers, after the attention residual add, the block applies:
a = RMSNorm(attn_out.detach()) # read attention as a signal
e = RMSNorm(e0) # original token embedding
gate = gate_proj(a) + causal_depthwise_conv(a) # kernel-9, depthwise
value = value_proj(e)
z = RMSNorm(out_proj(SiLU(gate) * value))
x = x + alpha * z # alpha is a learned scalare0 is the token embedding from the input layer, re-injected at every gate. The depthwise conv is strictly causal, so the gate is compatible with KV-cache generation (the conv state is carried alongside the attention cache).
Benchmarks
Zero-shot, evaluated in bf16 with an internal harness modeled on EleutherAI/lm-eval-harness; normalized accuracy where available.
Comparison
GPT-S2-5M against other small base models on the same evaluation suite. Achieving the highest avg score in the sub 10M category on the Open SLM Leaderboard
Training
Trained on a mixture of filtered web text (DCLM) and synthetic "finephrase" data (FAQ / math / table / tutorial)
Usage
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_name = "AxiomicLabs/GPT-S2-5M"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
device_map="auto",
)
prompt = "The future of AI is"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.inference_mode():
output = model.generate(
**inputs,
max_new_tokens=120,
do_sample=True,
temperature=0.8,
top_p=0.95,
repetition_penalty=1.1,
no_repeat_ngram_size=4,
)
print(tokenizer.decode(output[0], skip_special_tokens=True))Limitations
This is a small base language model. It is not instruction tuned, has limited factual capacity, and uses a 512-token context window.