The 1.08B-parameter GELU baseline for the flaxchat YatNMN ablation series. Standard attention + standard GELU MLP, same architecture/data/training recipe as the YatNMN d=22 variants. Drops the YatNMN-Softplus MLP and keeps everything else identical, so it's the direct control for testing whether YatNMN MLPs help at the 1B scale.

Training


Parameters	1,077,021,598
Architecture	d=22, n_embd=1408, n_head=22
Attention	standard multi-head self-attn with RoPE, GQA, QK-norm
MLP	GELU (standard)
Value embeddings	yes (Gemma-4 style, alternating layers)
Final smooth loss	2.82 (C4)
Tokens	21.5B (Chinchilla 20×)
LR	0.01, warmup-cosine
Hardware	TPU v6e-8, FSDP + remat

Compared to YatNMN d=22 variants

Model	Params	MLP	Loss
GELU 1.08B (this)	1.08B	GELU	2.82
YatNMN pn+α 1.08B	1.08B	YatNMN-Softplus (pn + learnable α)	2.83
YatNMN sb+ca 1.08B	1.08B	YatNMN-Softplus (sb + constant α=1)	2.83
YatNMN full 482M ⭐️	482M	YatNMN attn + MLP, no VE	2.57

At 1B scale with standard attention, GELU and YatNMN MLPs are essentially tied (2.82 vs 2.83). The big jump to 2.57 comes from the YatNMN attention architecture (no Q/K projections + L1 normalization), not the MLP choice.

Usage

python

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "mlnomad/gelu-d22-chinchilla-1B-pytorch",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")

ids = tokenizer("The meaning of life is", return_tensors="pt").input_ids
out = model.generate(ids, max_new_tokens=50, do_sample=True, temperature=0.8, top_p=0.9)
print(tokenizer.decode(out[0], skip_special_tokens=True))

Flax version

mlnomad/gelu-d22-chinchilla-1B (Flax/Orbax)

License

Apache 2.0.