~aniket

Recent Notes

5 min read

small language model

read - https://www.ibm.com/think/topics/small-language-models

  • the rise of LLM models

  • why we need SLMs?
    • LLMs are large and need more computational power to run and train
    • SLMs contain less than 10B⇐ params whereas LLMs have 100B+
    • SLMs are fast and efficient and can be deployed on small devices like phones, laptops, IoTs and can be fine tuned for specific tasks like healthcare, legal, customer support etc
    • On device privacy
    • ex - Microsoft Phi-3 Mini (3.8B), Google’s Gemma 2 (2B-9B) and Meta’s Llama 3.2 (1B-3B), DistilBERT, Mistral Ministral, Qwen2.5, SmolLM2

  1. Dataset

  • tiny stories dataset - https://huggingface.co/datasets/roneneldan/TinyStories generated by GPT-3.5 and 4 - https://arxiv.org/pdf/2305.07759

  • the goal of this dataset is to have as small dataset as possible while having all the concise information to teach something to the model

  • so our task is to use this dataset which contains many tiny stories in english for 3-5 year old children, can we successfully teach english to our SLM which is going to be only 15M params - if yes, then our goal is achieved!

    Goal: can we construct a language model with just 10-15M params and which produces coherent text?

  • the dataset contains 2.2M rows - every row is a story!

    • training - 2.1M
    • validation - 20k
  1. Data pre-processing

  • since LLMs cannot understand words like humans, we need to tokenize the stories in some way to convert them into numbers using tokenization:

    • types of tokenization:
      • character level tokenization - each character is a token
      • word level tokenization - each word is a token (english language has a big vocabulary so too many number of tokens, also spelling mistakes would not be a token!)
      • subword tokenization or byte pair encoding - characters and some basic words are tokens but some subwords are also tokens which can form new words like “tokenization” = “token” + “ization” ex- Byte Pair Encoding

  • so we follow a roadmap something like:

    • dataset → tokenizer → tokens → token id

  • the dataset contains rows of stories where each story is a row or a data sample and there are 2 million such rows for training and nearly 20000 such rows for validation!

  • iterating through the setup with a small example, we have to do 2 things:

    • tokenize the dataset → use gpt 2 subword tokenizer (BPE) for example 1 story is a collection of token ids with its length.

    • store all the IDs in a single .bin file! why?

  • we are storing the bin file as a memmap

  • split into batches (for faster loading in RAM) and add everything as a whole into train.bin and validation.bin for validation and test datasets!

  • creating input and output pairs from Dataset! why?

    • because LLM is a a next token prediction task and we want to take a sequence of tokens and predict the next token so we need ip and output pairs to calculate our loss!
    • context size - max len of words of tokens that the LLM needs to see to predict the next token ex- 4
    • batch size - the number of samples that the network processes to update the weights once!
    • when creating the input output pairs, we always stack the input tensor in batch_size x context_size matrix and output tensor in batch_size x context_size matrix but one token shifted meaning for each row, we have context_size number of tasks!
  1. Assembling the model architecture

  • first we get the tokenized text which is our input dataset but tokenized
  • each token is then transformed into a 768 dim vector using a dictionary which stores 768 dim vector for each token in the vocabulary of the LLM called as token embedding
  • finally we want to add positional embedding which gives the neural network the order of the tokens ex- “the dog eats the cat” and “the cat eats the dog” would have the same meaning in the n/w if there was no pos embedding!
  • the pos embedding is the same for all sample, as the number of pos would be at max → context size ex - 1024
  • the token embeddings and pos embeddings are intially randomized matrix
  • If your input has shape (B, T) = (8, 1024):
    1. tok_emb becomes (8, 1024, 768)
    2. pos_emb becomes (1024, 768) (broadcasted across the batch)
    3. x = tok_emb + pos_emb produces the final tensor of shape (8, 1024, 768), ready to enter the Transformer Blocks.
  • coming to the transformer block:

    • pre-layer norm - we normalize the tokens across the embedding dim of 768 so that the mean=0 and var=1

    • MultiHeadAttention

      • takes the embedding vectors and converts them into context vectors which contains the context of how a token relates to others!
      • so masked multihead attention takes the embedding vector of size 768 and converts it to context vector of the same size!
      • causality - a token should not attend to future tokens

  1. Setting up the SLM training pipeline

  2. Pre-training the SLM

  3. Running inference

202608280808

Graph View