Right-Sized Language Model

I ditch the GPU on my Haiku demo by training a right-sized language model from zero in JAX.

Overview

As an exercise, I migrated this interactive demo to run purely on CPU. It's a trivial app, but there's no NDA, so I can show it off.

In short, the model generates haikus. The demo is interactive, but the model has to autocomplete in domain. If it cannot produce a haiku, it cannot make good suggestions.

We're gonna compare its output, speed and cost to OpenAI's ChatGPT-4.1-Nano, as 4.1 is the last model to support simple and direct text generation.

To compare cost and performance, I benchmarked my model on a Hetzner CPX22, for which I pay €9.80 per month. 2 cores; 4 GB RAM; "AMD" processor. lscpu reports an AMD EPYC-Genoa.

OpenAI cost based on listed rates of $0.10, $0.40 per million input & output tokens, respectively. The local model's cost is clock time vs. the server's price per month of $11.40.

Neither OpenAI nor Haiku-Fu achieved perfect compliance with form, but GPT-likes aren't expected to. The app itself employs constraint-based sampling.

Flawless victory. Single night project delivers 96.3% cost reduction, no added server or GPU hosting. Or as a hype bot on X would say, "this is insane 🤯".

So, why did this work?

Background

A haiku in English is a simple poetic structure. Three lines, having five, seven and five syllables, respectively. It's an easy rule to explain and we hear of it on that merit.

Provided you can agree to a syllable count table, you can rig this up in a computer. That sounded fun, so I built this simple toy.

You click through to build a haiku word-by-word, or use the search box if you don't like the proposals. Then, you post.

To enforce the rule, I used the CMU Pronouncing Dictionary (Carnegie Mellon University) which encodes English pronunciations in their own DSL, like so: darling D AA1 R L IH0 NG. Digits indicate "stress", which I decided to interpret as a property of vowels and thus a syllable signifier.

The autocomplete is a little trickier.

Autocomplete, Round One

I wanted a little menu.

As we all know by now, GPT-likes "pick the next token" by computing a probability distribution over the big menu and rolling dice. Suppose, instead, you are the dice.

1. The computer puts odds on possible words.
2. The rule-based filter shortlists valid options.
3. You pick off the menu.

ChatGPT up to version 4.1 lets you grab top-K tokens and their logits. But:

1. Newer ChatGPTs restrict this.
2. Malicious network calls will burn tokens at cost.
3. I'd be leaving a "learning experience" on the table.

The current version, benchmarked above, employs a fully custom model, but initially I forked from an off-the-shelf small LLM.

I chose SmolLM2-1.7B as my starting point. Unlike SmolLM3 and other popular models, it's English-only, which reduces the impedance mismatch between its vocabulary and the CMU dataset. I then tuned the model on the Haiku 333K dataset, and hosted the resulting model on my NVIDIA 3090.

This worked alright, but I began to wonder if I could do better. A GPU tied up 24/7 for this?

Right Fit

1.7B is quite small by "large" language model standards, but this by no means implies it's the right fit.

First, only 9.4% of the words in the syllable table actually map to single tokens in SmolLM2's vocabulary. SmolLM2 therefore cannot produce a full probability distribution of the actual menu.

Second, LLMs are hyperliterate. The ability for the model to understand recipes, history, Python and COBOL ties up VRAM, at cost, even when not relevant to the task.

Mixture-of-experts doesn't solve this, as the "columns" aren't compartmentalized according to your-domain vs. not-your-domain. Neither does model size; a small model is simply dumber at everything.

As we saw from the benchmarks, an API does not save you either. RAM is a global market. A shortage is a shortage; waste is waste.

Domain Specific Pretraining

The value-prop of generative pretraining is that you reduce the size of your supervised training by pre training a model, unsupervised, on a large mass of data. Off the shelf models are trained on everything they can find, and if you train from zero, it'd seem you'd need a tremendous corpus to catch up.

Not so.

You can apply the same generative pretraining strategy, using a smaller, domain specific corpus. In our case, this means:

• The vocabulary matches exactly the words we have rules for.
• Pretraining material exists to teach grammar and vocabulary.
• The context window accommodates the actual task.

Vocabulary

The actual word-palette contains 45,245 words, nearly all of which are present in the Haiku-333K dataset. Some words weren't, but the final vocabulary is 43,208 tokens, including the basics. I included a token for the line separator, and encoded the English possessive suffix ('s) as its own token. (In the rulebook, 's represents zero syllables.) A model trained on this vocabulary can emit, in a single forward pass, a distribution for nearly the entire menu.

Pretraining

Haiku 333K, with this vocabulary, gives only 4.8 million tokens. Is that enough? I expect it isn't, and wanted to give the model a bit more grammatical grounding first.

I chose the much larger TinyStories dataset (arxiv), and filtered samples according to our vocabulary. This filtered out a surprising amount of junk and Mojibake, as well as stories with personal names and words out of distribution. A little more than half of the samples did encode correctly, giving 195 million tokens.

Context & Parameters

The longest sample in Haiku 333K amounted to just 22 tokens.

I set the context window to 64 tokens, which easily encompasses all valid sequences, and gives enough space to accommodate paragraphs TinyStories, which I cropped to this limit. (An experiment to increase this limit to 128 only gave a slight increase in token volume, to 199 million. I left it at 64.)

Results

The resulting model delivers pleasing and usable recommendations in just 4 ms, using an AMD 5900X (released 2020), single-threaded. (Note: the benchmarks above were completed on the more modest EPYC-Genoa.)

Let's have a look at its output with a Boltzmann softmax sampler.

The pretraining on TinyStories completed in less than half an hour. Three epochs. The actual downstream objective does not require case or punctuation, so the tokenizer does not support it, but the model easily strings together parts of speech in a plausible order:

a man was walking in the woods he was a compassionate man who worked in the trees he had a scale in his hand he examined it carefully [eos]

From that checkpoint, I tuned the model with Haiku 333K. With random sampling, the output doesn't quite match the rules:

fireflies at dusk /
tiny lanterns flicker twice /
night lights what we have [eos]

stars reclaim the mind /
their light worn and barren /
hopes eternal guard [eos]

At this point, we can bring in constraint based sampling, simply masking out the logits for violating tokens.

morning's gentle kiss /
trickster of life and laughter /
seeking morning's glow

a mirror cracks bright /
reflections written in light /
beauty in the truth

mountain's shadow falls /
canyon whispers echo low /
universe in sound

Folding this into the service, I thus promote the model from "glorified autocomplete" to just "autocomplete".

You can try here, with no registration.

Thanks for reading!