mlnomad/yatnmn-full-d22-chinchilla-pytorch

Full YatNMN (Attn + MLP) d=22 (482M) — PyTorch

A 482M-parameter GPT with both YatNMN attention and YatNMN MLP, no value embeddings.

YatNMN Attention (novel — no Q/K projections)

x_heads = RoPE(x)
dots = x_heads @ x_heads^T                       # pairwise dot products
dist² = ||x_i||² + ||x_j||² - 2·dots             # pairwise distances
scores = (dots + softplus(b))² / (dist² + softplus(ε))   # per-head b, ε
scores = L1_normalize(scores)                     # NOT softmax (scores already non-negative)
y = scores @ V

Key differences from standard attention:

• No Q/K projections — saves 2/3 of attention parameters
• L1 normalization instead of softmax (scores non-negative by construction)
• Per-head learnable bias and epsilon (both through softplus)
• Strict causal: tokens cannot attend to themselves (j < i)

YatNMN MLP

y = α · (x·W + softplus(b))² / (||x - W||² + softplus(ε)) — scalar bias, learnable α

Training

Parameters	482,283,432
Architecture	d=22, n_embd=1408, n_head=22
Final smooth loss	2.574
Hardware	TPU v6e-8 (FSDP + remat)

Beats the same-depth d=22 baselines:

Model	Architecture	Loss
GELU d=22 (1.08B)	standard	2.821
YatNMN-Softplus pn+α d=22 (1.08B)	standard attn	2.828
YatNMN full d=22 (482M, this)	YatNMN attn + MLP	2.574

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

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

Flax version

mlnomad/yatnmn-full-d22-chinchilla

License

Apache 2.0.