Pluto


Part 1. I decided to see how far I can go with making my own model. This is an ongoing essay that tracks Pluto’s progress across various improvement. Consider this to be v1. Pluto is a small GPT style language model. It is built using a GPT-2 style tokenizer with 50,257 tokens. The context length is 256 tokens. 6 transformer layers, 6 attention heads with an embedding size of 384. The dropout is set at 0.1. It’s currently around 30 million parameters. The architecture contains token embeddings, learned positional embeddings, six casual transformer blocks, pre-layer normalization, multi-head self-attention, GELU feed forward networks, a final layer normalization and a language modelling output head. The model uses casual attention so each token can attend to previous tokens. This allows it to learn text generation by predicting the next token.

The token embedding and output projection share weights through weight tying reducing the number of parameters. It uses the GPT-2 byte pair encoding tokenizer from tiktoken. The pretraining data is prepared by streaming documents from fineweb-edu. Because this was the first model, and I was experimenting, it only collects 100 million tokens in total with 1 million tokens for validation. It uses the remaining tokens for training. The training steps, pretraining and supervised fine-tuning were done via on Google Colab using a T4 due to a lack of decent compute. It would’ve taken days on my laptop.

Each document is tokenized and terminated with the GPT-2 end of text token. The base model training task is standard casual language modelling. More precisely, for every 256 tokens, the model receives the sequence shifted by one token as the target. The loss is cross-entropy over the next token predictions.

The default training settings include a batch size of 8, gradient accumulation of 4 steps. The effective batch size is 32 sequences with a sequence length of 256 tokens. The learning rate occurs at  3 × 10⁻⁴. It uses an AdamW optimizer with a weight decay of 0.1, gradient clipping of 1.0 and goes up to 12,000 training steps. It uses mixed precision training.

After pretraining, the model is fine-tuned to behave more like an instruction following an assistant. Initially, I was getting responses like this: alt text It spits out a random sentence it learned from the data. To avoid this, the tatsu-lab/alpaca dataset was applied. Each example is converted to this format:

### Instruction:
Explain how photosynthesis works.

### Response:
Photosynthesis is the process by which...
<|endoftext|>

The data is tokenized with the same GPT-2 tokenizer. It starts from the pre-trained base checkpoint and uses a much simpler learning rate of 1 × 10⁻⁵. However, the model ran into some issues. At first we trained on 3000 steps. This was the output. alt text

As a result, the step size was increased to 10,000. But this caused overfitting. alt text
So, as an attempt to reduce overfitting, I added early stopping and reduced the step size to 6000. This was the result alt text

Improvements are needed all across. The aim of Pluto is to be as reliable as possible. The next obvious steps are to improve pre training. Part 2 dropping soon.

You can find the current version of Pluto here: https://huggingface.co/spaces/bennymach/plutochat