TinyAmlGpt-Base-25k

A 51M-parameter gpt decoder-only language model trained on TinyStories with the AML objective (masked self-distillation).

Objective

Trained with a masked self-distillation objective. Alongside cross-entropy on the dense model, nested subnetworks are constrained to stay within a tolerance of the dense model's next-token distribution:

L = L_CE(F_0(x), y) + lambda * sum_j [ KL( p_j(.|x) || sg[p_0(.|x)] ) - eps_j ]_+
eps_j = eps * rho^(-gamma * j)

Masks are drawn once by iterated Bernoulli thinning, M^(j) = M^(j-1) * Bern(rho), and held fixed. The hinge makes proximity a constraint rather than a target: inside the ball the term and its gradient vanish, leaving each level free to occupy any point near the reference.

lambda	1.5
levels (m)	2
rho	0.5
eps	0.036
gamma	1.0
eps_j	[0.072, 0.144]
mask seed	0

Training

Architecture	gpt (8 layers, d_model 512, 8 heads)
Parameters	51,430,400 (25,698,816 non-embedding)
Data	TinyStories (GPT-2 BPE tokenizer, vocab 50,257)
Sequence length	512
Tokens per step	65,536
Steps	25,000
Total tokens	1.64B
Optimiser	AdamW, lr 0.0006, warmup 1000, wd 0.1
Precision	bf16
Seed	0

Results

metric	value
Validation loss	1.3709
Validation perplexity	3.94
Best validation loss	1.4162

Inference cost

Measured on NVIDIA RTX A6000 (bf16), prompt of 30 tokens, decoding 64 tokens, median of 5 runs.

level	total params	active params	density	prefill tok/s	decode tok/s	ms/token
dense	51,430,400	51,430,400	1.000	5,054	139.7	7.16
level 1	51,430,400	38,584,320	0.500	5,195	150.6	6.64
level 2	51,430,400	32,160,000	0.250	5,077	156.3	6.40

Masking zeroes weights rather than removing them, so the tensors keep their shape and the matmuls do the same work. A level with fewer active parameters is smaller in capacity, not in latency; realising a speedup would need structured pruning and a rebuilt model.

Usage

This model uses a custom implementation, not transformers. Clone the code, then:

import torch
from models import ModelConfig, build_model
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

path = hf_hub_download("HuggingAnalist/TinyAmlGpt-Base-25k", "model.safetensors")
cfg = ModelConfig(**json.load(open(hf_hub_download("HuggingAnalist/TinyAmlGpt-Base-25k", "config.json")))["model"])
model = build_model(cfg)
model.load_state_dict(load_file(path))
model.eval()

Limitations

Trained only on TinyStories: a synthetic corpus of simple stories written with a small vocabulary for young readers. The model has no knowledge outside that distribution, will not follow instructions, and produces fluent-sounding text that carries no factual grounding. It is a research artifact for studying training objectives at small scale, not a general-purpose language model.

Environment

Trained on NVIDIA RTX A6000 (47.5 GiB), PyTorch 2.8.0+cu128, CUDA 12.8.