SMall LLM project

SMALL LLM Project

Project explanation

This project used a modified version of the transformer architecture, using a mixture of experts and an alternate version of multi head attention known as multi latent attention for more efficient training. The model has 115 million parameter model, the optimizer used to train the model is a Muon. The training and inference was all done on google colab with its A100 GPU that it provided.

github repo: https://github.com/Thomasche69/SmallLLM/tree/main

Training dataset

link: https://huggingface.co/datasets/HuggingFaceTB/smollm-corpus

The training data used was the smollm-corpus, a diverse dataset that could be combined into a text corpus. The dataset contains two main columns along with three other additional ones acting as descriptors. The first column was called 'prompt' while the second one was called 'text'. For my project, my model only went through the process of pre training, hence I only selected the text column for training.

Example row of dataset

This is the class responsibile for handling functions during the creation of batches. The strides are set to seq_len as deafult, this ensures that each new sequence length of text does not overlap with one and another which prevents the risk of overfitting on the same data, a mistake that I did originally by not including this feature.

This code is responsible for compiling rows from the text column of the dataset into a combined text chunk acting as a text corpus. The number of rows is specified by hte num_documents configuration

Model Architecture

Multi Latent Attention

The code for this implementation is too large, if you are curious you could simply check the github page and navigate to the model.py file.

In most LLMs today, they use this technique known as KV Cache. To understand this technique more in detail, let's use this example. Attention is responsible for capturing the relationship between different tokens in a sequence length. Using the captured relationship, it will then predict the next possible token, which the model in this case correctly does as 'blue'. This output is known as an attention score.

After it predicts the token 'blue' it is then passed back into the model. However when doing so, we are recomputing the previous scores that we have already computed, when we actually only need to compute the new relationship between blue and the previous tokens.

Credit: Welch Labs

To prevent this from happening, we can cache the previously computed rows and store them in memory, which saves us from redundant computation.

Credit: Welch Labs

This implementation essentially reduces the cost of compute during inference from O(N^2) to just O(N)

Credit: Welch Labs

However this technique comes at a memory cost. If we use the example of the Deepseek V3 architecture's configuration size and with an input of 100k tokens, the memory cost ends up being 400GB.

Credits: Welch Labs

This of course became a huge issue in model training for many front tier labs. LLMs such as Meta's Llama 3 used a variant of multi head attention known as Grouped query attention. While this solution did decrease the memory cost, it also came at a cost of performance.

However the Deepseek team created a new variant of Multi Head Attention known as Multi Latent Attention.

Credits: Welch Labs

Deepseek's implementation is essentially 57 times more efficient than standard multi head attention, yet it also achieves higher performance than it. What MLA does differently from standard MHA, is it learns to efficiently compress its own keys and values into a latent space. With a few optimizations in clever linear algebra, using this idea, they came up with this final version of MLA.

Credits: Welch Labs

This was the implementation I used for my Transformer architecture. However I had to make some key modifications. An important thing to remember is that KV cache is only applied during inference. So I had to add a system handing this key distinction, allowing me to turn KV cache on and off using a single variable. The idea of adding this system initially came from someone's feedback on a discord server where I posted my project. The username of that person was Rohan Khan, who gave me the idea for this implementation.

Rotary Positional Encoding

This function is responsible for establishing the relative position between two tokens in a sentence. The key idea behind Rotary Positional Encoding (RoPE) ensures that the attention module encodes the relative distance between tokens in a way so that it stays the same even if more tokens are added into the sentence. RoPE will ensure that the model's recognition between pig and dog stays the same even after more tokens are added. Take the example below, where the tokens pig and dog have the same relative distance even with the sentence being expanded.

Credits: Efficient NLP

This is the key equation used to capture this effect, it uses the angle theta to encode each token's position

Credits: Zhuiyi Technology Co., Ltd. Shenzhen

This function is only applied to the K and Q matrices, however in my implementation instead of applying this effect to all parts of the Q and K matrix, I only use a certain part where I apply RoPE, which splits the Q and K matrix along the model_dim dimension between a rope_dim and nope_dim.

Mixture of Experts

The code for this is also too large, look on github if you want to see the full code.

In previous iterations of transformer architectures, they used a simple Feed Forward Network (FFN). Nowadays, this module is now replaced by the Mixture of Experts (MoE) module.

credit: https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts

The key idea between MoE is that a router creates a probability distribution that the model learns based on the input for which expert should be selected. The expert that is selected is just a simple Feed forward network, that will recieve and compute the input.

However, one important thing that MoE needs to watch out for, is so that no one expert is used more often than others. Ideally, all experts should be used roughly equally. In order to ensure this, a few extra additions such as noise that helps to normalize the probability distribution and an extra auxiliary loss that punishes the model if it overuses a few experts.

Training process

Muon Optimizer

Optimizers help adjust the model's parameters for better performance during training. For a long time, AdamW had been the most effective optimizer in deep learning. However recently, a new parameter known as Muon has been getting more popular, with frontier models like KimiK2 using it during their training process and delivering better performance in smaller LLMs while also being computationally less expensive.

The problem with AdamW is that the momentum factor only delivers a few small dominant directions that drive the updates. Muon optimizer solves this issue by increasing the effects delivered by less prominent directions through orthoganalization.

In linear algebra, the solution would be to use a technique known as Singular Value Decomposition or SVD. Say the matrix we use to update the weights is called M and let's just consider it to be a 2D matrix for now

To simply orthogonalize this matrix, we just set all the non zero values in the matrix S to one, giving us the orthogonal matrix.

Credits: Jia-Bin Huang

This is in fact not the solution we use for Muon, because performing SVD in practice is computationally expensive. But there is alternative method that uses the same idea, known as Newton-Schultz5.

Credits: Jia-Bin Huang

When applying this function iteratively, the input functions of the matrix slowly approach one, we can adjust the parameters a,b and c to control how quickly this happens

credits: Jia-Bin Huang

In practice however, the values don't need to exactly converge to 1, in fact they just need to be within a close and acceptable range like between 0.7 and 1.3. The values of a, b and c for Muon optimizer are 3.4445, -4.7750 and 2.0135 respectively in practice, as they achieve the desired outcome within only a few iterations

However in this project, we don't actually use Muon optimizer for the entire model. Instead we apply AdamW optimizer to the embedding and normalization layers and use Muon optimizer for the rest. Muon optimizer tends to be faster in convergence, while this is good for attention and feed forward network layers, it could lead to instability in embedding layers. Hence we use AdamW to ensure stable training in embedding layers.

Training loop

In the training loop, we use a cuda's autocast to dynamically uses FP16 and FP32 depending on how numerically sensitive a task is. Tasks such as matrix multiplication are represented in FP16 while tasks that re more numerically sensitive such as summing are done in FP32. FP16 and FP32 essentially describe how precise each number is represented using bits, with numbers in FP16 being represented by 16 bits and numbers in FP32 being represented by 32 bits.

We also apply a decay rate and warmup steps for the learning rate, with it starting from 0 and rising to the desired and set value, then slowly decaying towards 0. This is known as cosine decay, which ensures that training remains stable.

Credits: https://scorrea92.medium.com/cosine-learning-rate-decay-e8b50aa455b

Results

Training results

The model ended up with a final validation loss of 2.5998 and achieved an accuracy of 47.07% when predicting new tokens on the validation dataset. The auxiliary loss remained at the same value of 0.06 throughout training, which suggests that the use of experts had been relatively equal throughout training.

Inference

After training the model, I ran a few prompts to test out its knowledge and ability to form coherent sentences

I then decided to test its ability to write code, giving it the prompt "Using the numpy library". Surprisingly, it was actually able to write numpy code.

I then gave it the prompt "Algorithms". To my surprise again, it was able to form grammatically correct sentences and talk about different types of intelligence, though their factual validity is questionable. Nevertheless, it was even able to import the scikit learn library, which I did not expect it to do.

Evaluation and future work

The resulting model at the end of training delivered intelligent outputs as it formed grammatically correct sentences and was even able to write code. Nevertheless, there were still a few issues. For example, sometimes it just ended up repeating words over and over again when it was unsure of a response. This may be due to its limited training data, as the model was only run through 200 million tokens. Sufficiently well performing models usually need more than a billion tokens in training. Furthermore, the model only underwent pretraining, which means that it does not yet have the ability to hold a conversation. Future extensions of this project could explore fine tuning, in order to sharpen its conversational ability and perhaps RLHF, to make the responses more tuned to human aligned.

Page updated

Report abuse