VertexAGI/experiment-a2-pico

Experiment A2 Pico

A 1,009,920-parameter decoder-only transformer, pretrained from scratch on tinyshakespeare, trained 4x longer than Experiment A1 Pico (12,000 iterations vs 3,000) to see how much further the same tiny model + tiny dataset combination could go.

Same for-fun status as A1 -- not part of the production model lineup.

Architecture

Identical to A1 Pico -- only the training length changed.

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

Training

12,000 iterations, batch size 64, sequence length 128, AdamW (lr 3e-3), 727 seconds total.

Iter	Train loss	Val loss
1	4.9641	4.0178
1,000	2.0818	2.1416
2,000	1.7449	1.9314
3,000	1.6464	1.8299
4,000	1.5383	1.7229
5,000	1.4342	1.6962
6,000	1.4066	1.6411
6,500	1.3710	1.6245
7,000	1.3410	1.6246
8,000	1.3400	1.6237
9,000	1.3443	1.6104
10,000	1.3057	1.6201
11,500	1.2842	1.5994
12,000 (final)	1.2538	1.6284

The honest finding: this overfit

Val loss improved meaningfully through iteration ~6,500 (1.68 → 1.62), then plateaued and oscillated in the 1.60-1.63 band for the remaining 5,500 iterations while train loss kept dropping steadily (1.37 → 1.25). That's the textbook overfitting signature: 12,000 iterations × 64 batch × 128 tokens ≈ 98M tokens of gradient updates against a dataset with only ~1M characters in it -- the model saw the same text roughly 90+ times over, and past iteration ~6,500 it was increasingly fitting specific training sequences rather than learning more generalizable structure.

For comparison, A1 Pico (3,000 iters, no plateau reached yet) ended at train 1.4751 / val 1.6818 -- a smaller train/val gap (0.21) than A2's final checkpoint (1.2538 / 1.6284, gap 0.37). A2's val loss is slightly better in absolute terms (1.63 vs 1.68), but the widening gap says that improvement is mostly coming from memorization, not better generalization. No intermediate checkpoints were saved during this run, so only the final (iter 12,000, already past the val-loss plateau) weights are available here -- the true best generalizing checkpoint was somewhere around iteration 6,500-11,500, not the one actually published.

Takeaway: for a ~1M-parameter model on a ~1M-character dataset, roughly 3,000-6,500 iterations (2-5 effective passes) captures most of the learnable structure; training substantially longer mostly trades generalization for memorization. A real next step would be early stopping or a held-out-loss-triggered checkpoint save, not just "train longer."

Sample generations

All at seed-varied, temperature 0.6-0.9, 280 characters, greedy nothing filtered out -- unedited raw output:

1. Prompt: ROMEO: (temp 0.8)

ROMEO:
My gracious lord, be the practise our highness
The city yearth noble Carthard call enough,
To mean of discreen; for, the hear, and now art
We enforce our main heaven rook to pluck to a days
That is done partial, why, my brother there double;
One of some melancholy thanks in the

2. Prompt: JULIET: (temp 0.8)

JULIET:
Peace it need now she were a traitor!

LADY MORTETC:
One with all before I lay pure in death
Prements as he two often thou art than you must be
To crave for her former than idle with them.
The same soul as the world good at this earth
He cannot inherited was true to answer; Trani

3. Prompt: To be, or not to be (temp 0.8)

To be, or not to be a fair
Wish sorrow tenderly to how to soone thing issue.

HENRY BOLINGBROKKE:
He's not my honour and for aught it, I
was a prely the statues all, good nurse; but they
Thou she shall be a nure of her have east.

YORK:
And I love be a pusinest worthiling watch,
And fpeace and cour

4. Prompt: KING RICHARD:\nNow is the winter (temp 0.7)

KING RICHARD:
Now is the winter men, and tis now I am power:
Three for a kingdom charge thee but to do that thine honour,
Though an he make a cause, to do late them and him home,
Her to answer his well, thou shalt see hand.
I thank you can not go on the penite,
I thanks your pace is mistress that
Is all blow t

5. Prompt: First Citizen: (temp 0.8)

First Citizen:
You should forth that is done, for an words togethen
Look to you do this such a clouring on.

FRIAR LAURENCE:
Has contrary to breath there disloyal hour abse:
'Tis but not to fear, where I have to pass.

FRIAR LAURENCE:
Plantageneth then, in means and back.

DUKE VINCENTIO:
Het C

6. Prompt: O, what a (temp 0.9)

O, what aquaidy spoil:
then The disling reform my law I strength,
Who, fear prick.

LORD ROSS:
I am assist my lord, discovery sad'st of war
At a proported plainly to know his death
Be wand the Tower, I take allow me to thee
Shall burious rid of your dutied by them.
Thy that I was not like

7. Prompt: MENENIUS: (temp 0.6)

MENENIUS:
Not be much better there I would not so: I fear?

JULIET:
And thy valour will you shall be prick of will understong
To meet my cousin any time to allowing eyes and it nothing
By advise your honour'd and seldom and me,
And for a child the colour of the world,
And let the new of th

Qualitative read vs A1: A2 correctly recalls a much wider set of real Shakespeare character names verbatim -- Henry Bolingbroke, York, Friar Laurence, Duke Vincentio, Lord Ross, Menenius all appear unprompted, which A1's shorter run didn't produce. That's consistent with the memorization-over-generalization story above: the model has clearly seen (and retained) more of the specific text, but grammatical coherence hasn't meaningfully improved -- sentences are still not grammatical English, just more confidently wrong in a wider vocabulary.

Usage

Same loading pattern as A1 Pico -- this is a raw custom architecture (model.py), not transformers/mlx_lm-loadable out of the box:

from huggingface_hub import hf_hub_download
import json, sys, os

repo = "VertexAGI/experiment-a2-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())

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). Training data (tinyshakespeare) is drawn from public-domain source text.