I was recently playing around with training some small models from scratch. I was interested in whether I could make a decent one. I picked a size limit of 5 MB and tried a few different parameter counts and quantisation settings for a classic autoregressive transformer on the all-recipes dataset.
| Bits per weight | Parameters | Width, layers | Held-out loss | Training time |
|---|---|---|---|---|
| 2 | 18.6 million | 448, 7 | 1.684 | 20.6 min |
| 3 | 12.7 million | 384, 6 | 1.635 | 15.1 min |
| 4 | 9.6 million | 320, 7 | 1.654 | 16.1 min |
| 6 | 6.4 million | 288, 6 | 1.664 | 12.0 min |
| 8 | 4.8 million | 256, 5 | 1.704 | 8.6 min |
These models use quantisation-aware training (QAT), which is a cool technique where you basically keep two versions of each weight in the network: the full BF16 weight (the shadow weight) and the quantised weight used in the forward pass. On the backward pass, the shadow weight is updated and used as the source of the quantised value.
My takeaway is that you want as many parameters as possible, but there’s a limit to how far you can push quantisation. In this case, the sweet spot seemed to be 3 bits per weight.
I was having fun with this and getting some ~decent recipes out of the model:
Fish Terrine With Lemon, Garlic And Oregano
Ingredients:
- 1/4 cup extra virgin olive oil
- 4 garlic cloves, minced
- 1 cup diced tomato
- 2 teaspoons dried oregano
- 2 teaspoons chopped fresh thyme
- 4 1/2 - 5 lbs catfish fillets, cut in half horizontally
- lemon wedge, to serve
- black pepper, to taste
Directions:
- In a large frying pan heat oil on medium-high heat. Add garlic and cook for 1 minute. Stir in tomato, oregano and thyme. Simmer for 1 minute.
- Season catfish with salt and pepper.
- Arrange catfish on a baking tray and drizzle with the tomatoes. Drizzle with some of the tomato tomato sauce. Bake at 450 degrees F for 10 minutes, until fish is cooked through.
- Serve with lemon wedge and a green salad.
Sounds plausible, especially the “tomato tomato sauce” 😂
Turning this into a TTS model
I started thinking about what a TTS model would be like at this size. To turn this into a TTS model I’d need something to predict: some representation of the audio that isn’t too big. You basically want to compress the audio in each frame and have the model predict that, then feed the prediction back into the model (autoregressive) and get it to predict the representation for the next frame.
The first attempt was the codebook approach: splitting the audio into 20 ms frames, extracting a few features from each (such as loudness, pitch and the spectral envelope), and running k-means over those frames to find 4,096 clusters. Each cluster becomes a token, which the model is trained to predict. So this is essentially compressing each 20 ms of audio to a single token.
This is some pretty extreme compression and predictably loses quite a lot. Here’s what that sounds like.
The vocoder comparisons below use the same sentence from a clip held out of training, about 6.6 seconds long:
a scaffold was erected in front of that prison for the execution of several convicts named by the Recorder.
That codebook version includes the compression to one token per frame, so it’s the best the first approach could sound for this clip, even if the model predicted every token correctly.
This was going to need a much bigger vocab size or a different approach. Even if the model predicted the 19 numbers per frame exactly, it still isn’t great. That led to looking at FARGAN, which is a neural vocoder that recreates the audio from 20 numbers per 10 ms frame. You can hear a sample here of the features extracted from the original clip and then “uncompressed” with FARGAN.
To my ears that sounds pretty close to the original.
The next step was to run the LJ Speech dataset through the fixed, non-learned Opus feature extractor to get the 20 numbers per 10 ms frame that FARGAN can use to recreate the original, and then train the same autoregressive transformer design as above (8 layers, 384 wide), this time in plain float with no quantisation. It predicted the next 20 ms of speech, so two FARGAN frames or 40 numbers at a time.
The input was the whole sentence, then a separator, then each character followed by the speech for it. This needed a character-to-audio alignment computed beforehand. After each step the model also predicted whether to move on to the next character.
For example, for “Hi” the input could look roughly like this. Each audio bracket stands for 40 numbers, or 20 ms of speech:
H i ← the whole text
<separator>
H ← the current character
[audio step 1: 40 numbers]
[audio step 2: 40 numbers]
i ← the next character
[audio step 3: 40 numbers]
[audio step 4: 40 numbers]
...The number of audio steps under each character depends on the alignment; two each here is just an example. After each audio step, the model predicts whether to stay on the current character or move to the next one.
The idea is still to predict what comes next, similar to text, but the audio went in as 40 continuous numbers through a linear layer rather than tokens. For the output, there were two versions: predicting each number as one of 32 levels, and predicting the numbers directly.
Here’s what those two versions sounded like on the same sentence:
Obviously not great. I think generating audio is difficult this way because speech barely changes from one 20 ms step to the next. If the model is just trained to predict the next frame, I think it learns to just look at the previous one and make a slight variation of it. That gets low loss in training, where the previous frame comes from a real recording, but when it’s generating, the previous frame is its own slightly wrong output and the errors pile up. What you want is closer to a blank canvas, where you decide up front that there is an “h” sound here and an “i” sound there, and fill in every step from the text without the model ever looking at its own output.
In order to do this, I think there needs to be some model up front which takes the input characters and tries to learn some representation of their sound in that context, then stretches that out and runs it through another model that learns how the frames change from step to step so say your “h” sound lasts 200ms that would start as 10 input vectors which a model learns how they evolve throughout the 200ms.
This turns out to be the fastspeech approach except in that in that they predict mel-spectrogram frames.
So we trained a model along those lines, with a text encoder, a duration predictor and an acoustic decoder. This time the decoder predicts the speech features in parallel, rather than feeding its own predictions back in.
The encoder and decoder use the same transformer block: 384 wide, with 6 attention heads and a 1,024-unit feed-forward layer. The decoder just has more layers, so it ends up about 1.5 times the size of the encoder.
| Part | Layers | Parameters |
|---|---|---|
| Text encoder (characters → context vectors) | 4 | 7.1M |
| Acoustic decoder (stretched vectors → FARGAN features) | 6 | 10.6M |
| Embeddings and output layer | — | About 0.06M |
| Acoustic model total | 17.8M | |
| Duration predictor (a separate model with its own character embedding) | — | 3.2M |
Together that’s about 21 million parameters. To make these models small enough to download and run in the browser, I went back to the quantisation-aware training from the recipe experiment, storing the weights at 3 bits each. The acoustic model and duration predictor together came to 8.3 MB.
The quality loss from quantisation was pretty small. On our held-out test, the full-precision version had a character error rate of 9.1%, compared with 9.3% for the 3-bit version. They sounded very similar too.
Those parameter counts are just for the models we trained. FARGAN, which turns the features into audio, and the optional bandwidth extender both come from Opus. We didn’t train those, and they aren’t included in the 21 million parameters.
Try it
You can try the model here. It runs entirely in your browser, and this is what makes up the download:
| File | What it does | Size |
|---|---|---|
tts.bin | Acoustic model and duration predictor | 8.32 MB |
fargan.bin | Vocoder that turns the features into audio | 3.27 MB |
bwe.bin | Bandwidth extender for 48 kHz audio | 1.51 MB |
Supporting .js and .json files | Code to run the demo and describe the model layouts | 59 KB |
| Total | 13.16 MB |
Those are the uncompressed file sizes. Type something and press Speak to try it; the model weights only download when you first press the button.