Devlog: From bigram to building my own tokenizer

Last time I had a bigram model trained on The Wizard of Oz that mostly made gibberish. Since then ARIA has gone through three versions, and I ended up writing most of the pipeline from scratch.

ARIA-01 / 02: a real transformer

I replaced the bigram with a character-level GPT: 8 layers, 512-dim embeddings, 8 attention heads, 256-character context, trained on web text on my RTX 5060 Ti.

• Bumped block_size from 8 to 256 so the model can actually see whole sentences
• Added bf16 autocast, which cut VRAM use roughly in half
• Found out the hard way that when VRAM fills up, Windows quietly spills into system RAM and training gets about 20x slower. Batch 128 was taking 6 to 9 seconds per step; batch 64 was 0.35 seconds
• Fixed a bunch of bugs along the way, like my argparse -batch_size flag secretly setting block_size, and eval running 500 batches every 250 steps

Final pretraining loss was about 1.30 per character, and it started writing real English sentences (that didn’t make sense yet).

Teaching it to chat

The base model only continues text, so I fine-tuned it on Alpaca formatted as User: ... / ARIA: ... conversations, with a lower learning rate so it wouldn’t forget English. Loss went down to about 0.91.

Then I had to fix the chat script: wrapping prompts in the right format, cutting replies at the end of a turn, and adding temperature + top-k sampling so it stopped picking random characters. First real reply to “How do you make a cookie?”:

You make a cookie are both a cookie.

Grammatically broken, but it’s an actual answer, in the right place, that stops on its own. The pipeline works.

Writing my own BPE tokenizer

The biggest limit was being character-level: the model wasted most of its capacity on spelling, and 256 characters isn’t much context. So I wrote a byte-level BPE tokenizer from scratch instead of using a library:

• get_stats and merge to find and merge the most common pairs
• a training loop with stop conditions (no more merging pairs that only appear once)
• encode / decode with a round-trip test, plus save/load for the merges
• GPT-2 style regex pre-splitting so tokens can’t cross word boundaries (before this, I was getting tokens like "is to " and "00 MPa ")
• a cache in encode that made it about 2.3x faster on small text and much more on big files

Trained on about 1MB of educational text plus chat data, with 2048 tokens. It learned real word pieces like " economic", " School" and "rying", and gets about 3x compression. Since it’s byte-level, there are no unknown characters anymore, so Korean works too.

ARIA-03

I encoded 500MB of fineweb-edu into 169.6M tokens, saved as uint16 binary and loaded with np.memmap. Then I rewrote the model:

• Fused attention with F.scaled_dot_product_attention instead of 8 separate heads, about 3x faster per step
• Pre-norm blocks and GELU
• Weight tying between the embedding and output layer
• Warmup + cosine learning-rate schedule
• state_dict checkpoints with optimizer state and step count, so a crash doesn’t wipe a run

Test run: loss started at 7.71 (right at ln(2048), which is what random guessing should give) and reached 4.27 after 1000 steps, with train and test matching the whole time.

Next

• Full 30,000-step pretraining run (about 3 passes over the data)
• Chat fine-tune on the new tokenizer
• Later: tool use, like search through the Naver Search API, a separate vision-based security AI for home cameras, and an agent that writes and uploads Arduino code