TinyAmlGpt-Med-e0.07-l1.5

A 124M-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.07
gamma	1.0
eps_j	[0.14, 0.28]
mask seed	0

Training

Architecture	gpt (12 layers, d_model 768, 12 heads)
Parameters	123,551,232 (84,953,856 non-embedding)
Data	TinyStories (GPT-2 BPE tokenizer, vocab 50,257)
Sequence length	1024
Tokens per step	65,536
Steps	20,000
Total tokens	1.31B
Optimiser	AdamW, lr 0.0004, warmup 1000, wd 0.1
Precision	bf16
Seed	0

Results

metric	value
Validation loss	1.1881
Validation perplexity	3.28
Best validation loss	1.2309

Inference cost

Measured on NVIDIA A100-SXM4-80GB (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	123,551,232	123,551,232	1.000	2,185	63.1	15.86
level 1	123,551,232	81,083,904	0.500	2,236	63.9	15.64
level 2	123,551,232	59,852,288	0.250	2,191	64.1	15.61

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-Med-e0.07-l1.5", "model.safetensors")
cfg = ModelConfig(**json.load(open(hf_hub_download("HuggingAnalist/TinyAmlGpt-Med-e0.07-l1.5", "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.