Training a Small Language Model

TL;DR: I’ve built and trained an ex­tremely small lan­guage model from scratch, specif­i­cally de­signed for short-form sen­tence rewrit­ing tasks that are highly re­source con­strained.

In my last post, I dis­cussed how I plan to tackle the dif­fi­cult prob­lem of help­ing Harper’s users take ad­van­tage of the ac­tive voice. The ac­tive voice is seen as more con­fi­dent and au­thor­i­ta­tive, which is of­ten de­sired in pro­fes­sions like cus­tomer ser­vice. I have ex­pe­ri­ence in that area, which is how I know the ex­tent to which help with the ac­tive voice can make the job eas­ier.

Naturally, we would only want to pro­vide help with the ac­tive voice when specif­i­cally re­quested by the user. That means it’s an opt-in fea­ture that is­n’t en­abled by de­fault. If we suc­ceed in im­ple­ment­ing the func­tion­al­ity, and peo­ple seem to like it, the Harper pro­ject can re­con­sider this stance. For now, the plan is for it to be dis­abled by de­fault.

As I it­er­ated in my last post, the state-of-the-art method for con­vert­ing from the pas­sive voice to the ac­tive voice is to use an au­tore­gres­sive trans­former model. This is the same fun­da­men­tal tech­nol­ogy be­hind apps like Google Translate and ChatGPT. While the ac­tual model I in­tend use for this fea­ture is sim­i­lar in na­ture to those used in these apps, my im­ple­men­ta­tion will be en­tirely dif­fer­ent in scale.

You see, since our do­main is so lim­ited in scope, we can use an ex­tremely small model. Google’s small­est T5 model is a great point of ref­er­ence. A quan­tized copy of that model runs about 60 megabytes. I be­lieve we can train our own trans­former model from scratch and pack it in an un­quan­tized con­tainer mea­sur­ing no larger than 30 megabytes.

Why do I be­lieve this to be pos­si­ble? Be­cause I’ve done it.

Our Model Architecture

Over the past few days, I’ve im­ple­mented a ba­sic lan­guage model in Rust us­ing Burn. For the cu­ri­ous, the code is open source. It’s pretty messy, so please don’t look at the com­mit his­tory. Once I have a solid proof-of-con­cept work­ing, I’ll pol­ish it up and find a good home for it in the main Harper repos­i­tory.

The key here is that our use-case only needs a few hun­dred char­ac­ters of con­text. Since the com­pu­ta­tional cost of an LLM can be ap­prox­i­mated as the square of the con­text size, we can train highly ef­fec­tive mod­els our­selves that are small and can run any­where. If we want to pro­vide ad­di­tional in­for­ma­tion to the model, such as style or do­main, we can sup­ple­ment this con­text with cus­tom em­bed­dings.

Using a small model al­lows us to main­tain our pri­vacy guar­an­tees with­out sig­nif­i­cant cap­i­tal ex­pense. In fact, once we’ve trained the mod­els, they’re free to run in­def­i­nitely.

I’ve cho­sen to use a char­ac­ter-level to­k­enizer. That means that each char­ac­ter is its own to­ken, and the con­text size of the model is equiv­a­lent to the num­ber of char­ac­ters it can read at once. The hope here is that it al­low us to min­i­mize Harper’s fi­nal bi­nary size and give the model greater free­dom to change the in­flec­tion of words how­ever it pleases.

For now, here are some sim­ple sen­tences gen­er­ated with the ba­sic model I trained for just un­der an hour. It’s pretty amaz­ing what can be done with just un­der a dol­lar’s worth of com­put­ing power.

The forest rested in a speed of the night.

Feels… in­tro­spec­tive, no?

It was the today of the world.

Moving Forward

My fo­cus this week was to get a ba­sic au­tore­gres­sive model trained on generic English text. My fo­cus start­ing next week will be to fine-tune this model for our sen­tence rewrit­ing task. Once we have a very solid base model, I see no rea­son we could­n’t fine-tune sim­i­lar mod­els for other fea­tures, like chang­ing tone or im­prov­ing clar­ity.

My cur­rent im­ple­men­ta­tion, while fast, could also use some im­prove­ment. To that end, I’d also like to im­ple­ment a KV cache. That will al­low us to scale up to larger con­text sizes with­out a drop in qual­ity or speed.

P.S. Today is my birth­day! I’m so grate­ful for the years I’ve been on this earth, and I can’t wait for more. Thanks for read­ing! Here’s to 21.