TinyAmlGpt-Base-10kn : All MLP are Masked with rho = 1%

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.

	Values
lambda	2.0
levels (m)	1
rho	0.001
eps	0.00021
gamma	1.0
eps_j	[0.21]
mask seed	0

Training

	(Hyper) Parameters
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	10,000
Total tokens	0.66B
Optimiser	AdamW, lr 0.0006, warmup 1000, wd 0.1
Precision	bf16
Seed	0

Results

metric	value
Validation loss	1.3744
Validation perplexity	3.95
Best validation loss	1.4158

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	51,430,400	51,430,400	1.000	3,328	100.0	10.00
level 1	51,430,400	34,146,152	0.001	3,302	99.0	10.10

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-10kn", "model.safetensors")
cfg = ModelConfig(**json.load(open(hf_hub_download("HuggingAnalist/TinyAmlGpt-Base-10kn", "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.