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, 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   = 9

The 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 scalar

e0 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.