VertexAGI/experiment-a1-pico

Experiment A1 Pico

A 1,009,920-parameter decoder-only transformer, pretrained completely from scratch (no distillation, no base model) on tinyshakespeare -- ~1.1MB / ~1M characters of Shakespeare -- in 188 seconds on an Apple Silicon Mac using MLX.

Not part of the production model lineup -- a for-fun experiment in how much a genuinely tiny (~1000x smaller than everything else in this project) transformer can learn from a genuinely tiny dataset.

Architecture

Same ingredients as every larger model in this project, just scaled down by ~1000x: RMSNorm, RoPE, causal self-attention, SwiGLU MLP, tied input/output embeddings.

Parameters	1,009,920
Hidden size	160
Layers	3
Attention heads	4
Intermediate size (SwiGLU)	480
Context length	128
Vocabulary	65 (character-level -- every unique character in tinyshakespeare)

Character-level rather than a subword tokenizer by necessity: with a real ~50k+ token vocabulary, the embedding table alone (vocab_size * hidden_size) would dwarf the entire 1M-parameter budget several times over. Going char-level keeps nearly all the params in the actual transformer instead of a lookup table.

Training

• 3,000 iterations, batch size 64, sequence length 128, AdamW (lr 3e-3).
• Train loss: 4.71 -> 1.48. Val loss: 3.90 -> 1.68 (held-out 10% split, tracked down alongside train the whole run -- no overfitting).
• Total wall-clock time: 188 seconds, entirely on a Mac.

Sample output

Prompted with ROMEO:, temperature 0.8:

ROMEO:
I'll resent dead, or for the fearful,
For me!
I will sends of thee leave at you are with sraise thy quands the knee the honour comfort.
I cannot of and me horn shaper.
Come, good you thy brothers blike
Than we mays, to die done's; sir,-home not me,
To have this virge a better had the elding,
That he words of the but the not I lives
For thou?

Who, the God, good maid the Tower.

CLARENCE:
You no m

Not coherent Shakespeare -- but it learned real structure from 188 seconds of training on ~1M characters: correct play-script formatting (character names in caps + colon, line breaks), real English punctuation and capitalization conventions, and mostly-plausible word shapes with an unmistakable Elizabethan flavor ("thee," "thy"). That's the model finding structure, not memorizing -- 1M characters of training data is nowhere near enough to memorize.

Usage

This is a raw custom architecture (not registered with transformers or mlx_lm) -- load it with the included model.py:

from huggingface_hub import hf_hub_download
import json, sys, os

repo = "VertexAGI/experiment-a1-pico"
model_py = hf_hub_download(repo, "model.py")
weights = hf_hub_download(repo, "pico.safetensors")
vocab_path = hf_hub_download(repo, "vocab.json")

sys.path.insert(0, os.path.dirname(model_py))
from model import PicoConfig, PicoModel
import mlx.core as mx

vocab = json.load(open(vocab_path))
stoi, itos = vocab["stoi"], {int(k): v for k, v in vocab["itos"].items()}

cfg = PicoConfig(vocab_size=len(stoi))
model = PicoModel(cfg)
model.load_weights(weights)
mx.eval(model.parameters())

# Autoregressive sampling (see train.py's generate() for the reference implementation)
tokens = [stoi[c] for c in "ROMEO:"]
for _ in range(200):
    logits = model(mx.array([tokens[-cfg.max_seq_len:]]))[0, -1]
    next_id = mx.random.categorical(logits / 0.8).item()
    tokens.append(next_id)
print("".join(itos[t] for t in tokens))

License

MIT (architecture/training code). The training data, tinyshakespeare, is drawn from public-domain source text.