🤏Pico

Pico is my experimental language model decoder-only architecture designed to bypass tokenization by processing raw bytes directly.

To offset the computational cost of working at the byte level, the model employs sliding window attention over a small context. To reintroduce longer-range dependencies, intermediate transformer blocks operate on larger windows, but they are selectively applied over a small fixed fraction of the bytes using a Mixture-of-Depths router.

This architecture share similarities with SpaceByte, but instead of applying intermediate transformer blocks when encountering whitespaces, Pico learns by itself when to insert them in an end-to-end manner.

I also experimented with multi-byte prediction, which allows the model to not only predict the next byte but also the next few bytes. This has been shown to improve the performance of byte-level models, although probably not at the scale I run my experiments but it also allows for self-speculative decoding at inference time.

Other architectural choices include: - Sliding window, grouped-query attention - SwiGLU - ALiBi positional encoding (might try RoPE later) - SOAP optimizer for pretraining

Here is an example of the output of a ~10M parameters Pico model trained on the TinyStories dataset:

``` Once_ upon_ a time, there_ was a little_ girl named_ Lily._ She_ loved_ to_ play_ outside in the_ park._ One_ day,_ she_ saw a bee flying around_ her_ and said, "Hello, bee! Do_ you_ want to_ play_ with_ me?_"

The bee nodded_ and said, "Yes, I_ would love_ to_ join me,_ too!_" ```

The _ prefixes the bytes where the model decided to insert the intermediate transformer blocks. As we can see, the model has somehow chosen to insert them at word boundaries like a form of "tokenization".

When completly disabling the intermediate transformer blocks, the model can still generate well-formed words, but loose coherence at the sentence level:

Once upon a time in the wet borided in the safe in all the wonderful and thought again. He smilly came of the brave enough to listen that missides! He said whines waved as she funny thanks the weak down that is all the bricycle.

This suggests that the intermediate blocks may produces "higher-level" representations that helps the model to maintain coherence over longer ranges.

That being said, I also observed harder-to-interpret patterns on models trained on other datasets, so it's not clear that the Mixture-of-Depths is systematically used by the model as a form of tokenization.

Usage

```bash

Installation

uv sync

Initialize a new model

python picolm.py init models/my-model --dim 372 --latent-num-blocks 8 --fb-att-window-size 32

Train the model

python picolm.py train models/my-model my-train-run roneneldan/TinyStories --dataset-column-name text --batch-size 8

Generate text

python picolm.py run models/my-model my-train-run --temperature 0.8 --prompt "Once upon a time" ```