mlnomad/yatnmn-sb-alpha-d22-chinchilla-1B-pytorch

YatNMN-Softplus sb+α d=22 (1.08B) — lr=0.03, PyTorch

The scalar-bias + learnable-α variant of the 1.08B YatNMN-Softplus family — completes the d=22 ablation grid. Standard multi-head attention + YatNMN-Softplus MLP with scalar (shared) bias and a single learnable α per layer. Trained with lr=0.03 (3× the default) on C4, keeping value embeddings.

Architecture

Parameters	1,077,021,664
Depth / n_embd / heads	d=22 / 1408 / 22
Attention	standard MHA with RoPE, GQA, QK-norm, VE on alternating layers
MLP	YatNMN-Softplus, softplus(b) shape=(1,), softplus(ε) shape=(1408,), learnable α
Tied embeddings	yes
Final smooth loss (C4)	2.858 (step 575,000)
Learning rate	0.03 (warmup-cosine)
Tokens	152B (7× Chinchilla; training kept running past chinchilla-mult=1)
Hardware	TPU v6e-8, FSDP + remat

Full d=22 (1.08B) ablation family

Variant	MLP bias	α	LR	C4 loss
GELU 1.08B	— (GELU)	—	0.01	2.82
YatNMN pn+α 1.08B	per-neuron	learnable	0.01	2.83
YatNMN sb+ca 1.08B	scalar	const=1	0.01	2.83
YatNMN sb+α 1.08B (this)	scalar	learnable	0.03	2.858
YatNMN full 482M ⭐	scalar	learnable	0.01	2.57 (FineWeb-Edu + YAT attn)

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "mlnomad/yatnmn-sb-alpha-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/yatnmn-sb-alpha-d22-chinchilla-1B

License

Apache 2.0.