Pattern Roller

Parking lot pattern

This project started because I wanted to code a transformer from scratch using PyTorch. Instead of simply copying a tutorial or re-implementing nanoGPT or similar, I wanted to try something else. So instead of embedding discrete tokens, I wanted to feed pixels directly into the model.

Training then happens with standard MSE loss where the output is interpreted as pixels as well, which then produces images with a fixed height but theoretically unlimited width. Hence the “pattern roller” analogy where we can “roll” the trained input image indefinitely. The model is just trained on one image to reproduce it as efficiently as possible. Since we just train only a single image, efficiency here meant parameter count/model size.

To generate the training data, the input image was pre-processed to stitch together the image and its horizontally flipped copy and the model was trained to reconstruct the next column of pixels as discussed already.

Let’s get one thing out of the way: this is utterly useless!

Unless… the model size is smaller than the single image it has been trained from. Which - spoiler alert - was not the case. So you would always be better off to just cut and flip the original image.

Everything started with a “hand-written” transformer with attention written in PyTorch as well. I then applied several architecture and training optimizations to increase parameter efficiency. This part started 1+ year later and was LLM-assisted.

After the optimization the model has 230,272 parameters. Assuming fp8 (assuming almost no loss in performance) the model weights are around 230kB. This is around 60% larger than a PNG of the input and around 500% larger than a JPEG image.

Reducing the parameter count further might be a fun project for future me. The goal should then be to at least beat the PNG file size. For now these are the optimizations that landed here:

1. input/output embedding weight tying (Press and Wolf 2017) in model/embedding.py:104–115, model/transformer.py:226–233
2. reusing a single transformer block for all layers with sinusoidal step encodings (Dehghani et al. 2019) in model/transformer.py:217–220
3. fused SwiGLU FFN layers (Shazeer 2020) in model/transformer.py:16–36
4. multi-query attention (MQA) (Shazeer 2019) in model/attention.py:40–81
5. use AdamW with weight decay of 0.01 (Loshchilov and Hutter 2019) in model/transformer.py:170–189
6. cosine decay of learning rate to 10% after 5% epochs warmup (Loshchilov and Hutter 2017) in model/transformer.py:188–212

Beach pattern

Cloud pattern

Forrest pattern

References

Dehghani, Mostafa, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser. 2019. Universal Transformers. https://arxiv.org/abs/1807.03819.

Loshchilov, Ilya, and Frank Hutter. 2017. SGDR: Stochastic Gradient Descent with Warm Restarts. https://arxiv.org/abs/1608.03983.

Loshchilov, Ilya, and Frank Hutter. 2019. Decoupled Weight Decay Regularization. https://arxiv.org/abs/1711.05101.

Press, Ofir, and Lior Wolf. 2017. Using the Output Embedding to Improve Language Models. https://arxiv.org/abs/1608.05859.

Shazeer, Noam. 2019. Fast Transformer Decoding: One Write-Head Is All You Need. https://arxiv.org/abs/1911.02150.

Shazeer, Noam. 2020. GLU Variants Improve Transformer. https://arxiv.org/abs/2002.05202.