A language model for how India types

A keyboard is a language model with a tiny screen attached. Every suggestion above your keys is a guess about what comes next.

For English, those guesses are good because the models have seen billions of English sentences. For "kal office jaana hai yaar", they've seen almost nothing. And for the way you write it, they've seen exactly nothing.

So Keygram needed its own model. Small enough to run inside an iOS keyboard, trained on romanized Indian text, and cheap enough to build on a hobby budget (my total compute ceiling was around $30 to $50 on Modal). This post is how that went: the data, the model, and how we squeezed it.

Why from scratch

The obvious move is to take a pretrained model and fine-tune it. I didn't, for three reasons.

• Size. A keyboard extension on iOS gets a small memory budget and has to answer in about 20ms. Even "small" pretrained LLMs are hundreds of millions of parameters. I needed a few million.
• Tokenizers. Pretrained tokenizers are built for English and Devanagari. Romanized Hindi gets chopped into odd fragments, so the model spends its capacity on spelling instead of meaning.
• Licensing. Everything had to be commercial-friendly open source. A lot of small models and weights come with strings attached.

Training from scratch also meant I controlled the whole stack, from the vocabulary to the export format. For a product whose whole point is "your spelling is not a typo", that mattered.

The data problem

There is very little clean, romanized, code-mixed Indian text out there. Most public datasets are either pure English, native script, or formal translation pairs. Nobody writes "bhai kal milte h" in a parallel corpus.

So we generated a lot of it. We built two separate data pipelines:

• keygram-datagen for Indo-Aryan languages, mainly Hindi and Hinglish
• dravidian-datagen for Tamil, Telugu and Kannada

They had to be separate because the languages behave differently. Dravidian words are long and agglutinative: one word can carry what Hindi spreads over three or four. Vocabulary, word lengths and spelling variation all look different, so one pipeline tuned for Hindi produced bad Tamil.

Sources included public datasets like Dakshina (native script plus romanizations written by real people), plus other openly licensed text.

An idea the data killed

Early on I wanted to group spelling variants into one "canonical class": bahut, bhot and bahout all counted as one word. It would have made the model much more efficient.

We set a gate first: the grouping had to collapse the vocabulary by at least 3x to be worth it. Dakshina came in around 1.77x. An independent phonetic method got 1.99x, and that number was inflated by false merges. Gate failed, twice. The idea isn't dead for real typing data, but we didn't have that data, so we dropped it.

Setting the number before running the test is the only reason I let this one go.

ATLAS: the model

The main model is called ATLAS: 13.44M parameters, perplexity around 19.4, trained from scratch.

It's a hybrid built on GLA, gated linear attention (Yang et al., ICML 2024). The reason is simple. A normal transformer looks back over everything you've typed, and that cost grows with every word. GLA can run like a recurrent network: it keeps a fixed-size state and updates it one token at a time. For a keyboard that runs on every keystroke, constant cost per step is exactly what you want.

One honest note on that perplexity: it's measured on our own romanized test set with our own tokenizer. It's useful for comparing our versions against each other, not against other people's models.

What we tried along the way

Most of the learning came from the earlier, smaller (~5M) versions. Each idea got a gate before it went in.

The ranking one taught me the most. Top-1 fell because summing log-probabilities quietly favours short candidates, and the new objective made that bias worse. Dividing the score by the candidate's length fixed it. A one-line change that was a free win everywhere.

Squeezing it to 1.12M

13M parameters is small for a language model and still big for a keyboard. So the next step was GLA-Micro: the same idea, cut down to 1.12M parameters.

The bigger trick wasn't making the neural model smaller. It was not asking it to do everything.

n-gram first, neural second

Keygram splits the job in two:

1. A personal n-gram model proposes. It's fast, lives entirely on your phone, and learns your spellings directly from what you type. If you write "nhi", it knows.
2. The neural model reranks. It only scores a short list of candidates instead of the whole vocabulary, which keeps each keystroke cheap.

Each side covers the other's weakness. The n-gram is personal from day one but has no real sense of context. The neural model understands context but has never seen your spellings. Together you get suggestions that fit the sentence and look like you wrote them.

It also made the system easier to debug. When a suggestion was wrong, we could tell whether the candidate was never proposed or proposed and ranked badly. Those are two very different bugs.

What I took away

• Set the number before the experiment. The canonical-class idea felt right. A gate set in advance is the only thing that made me drop it.
• Small models are a design problem, not a compression problem. The biggest wins came from splitting the job (n-gram plus reranker) and fixing scoring, not from shaving layers.
• Data beats architecture, and real data beats generated data. Our ceiling was always the lack of real typing data, not the model.
• A tight budget is a good teacher. $30 to $50 of compute forces you to ask which experiment actually answers a question.

Next post: getting all of this to run inside an iOS keyboard in 20ms.