Transformers &amp; DSPy: The Perfect Combo to Start with LLMs
Source: tai_blog

IntroductionWho has never used ChatGPT? Probably every single one of us! However  we do not face one of the latest and most promising developments in artificial intelligence only when we use ChatGPT. Large Language Models (LLMs) have been implemented across different companies from different domains and we are likely exposed to them every day.   For example  customer service teams use this technology to quickly handle basic queries and let agents focus on more demanding issues. Marketing agencies use it to support their creative side when building campaigns or to understand customer sentiment in social media posts. Or  Spotify could have used this technology to create the lyrics through audio transcription.   With so many possible use cases and the level of exposure that we have  this article aims to provide a simple but detailed explanation of how the backbone architecture of LLMs works and what novel concepts companies like Meta  Mistral AI and Google introduced to this architecture with their own models  LLaMA  Mixtral and Gemma.   Finally  we provide a practical implementation in python using the library DSPy of these LLMs to tackle different use cases such as sentiment analysis  summarization  and RAG systems.   As always  the code is available on Github.   Transformers: from tokenization to text generationFirst things first: TokenizationTokenization is the first step in the process of text generation. It is not part of the Transformer architecture but it is crucial to transform the raw input text into a suitable format  tokens  so that Transformers can process text.   SentencePiece [1] developed by Google  is one of the most used tokenizers for LLMs. The way it works is the following:   1. Splits the words into individual characters Imagine that the training data contains the following set of words and the frequency of each word has been determined  for example  word hug appears 10 times  pug 5 times  pun 12 times  and so on.   Based on this frequency  the words need to be split into individual characters as shown below:   2. Iteratively merges the most frequent character pairs into subwords until a predefined vocabulary size is reached From the example above  we can see that u followed by g appears 20 times ( hug 10 times + pug 5 times + hugs 5 times)  which is the most frequent symbol pair  therefore we merge them:   We repeat this step until it reaches the predefined vocabulary size which  for example  in our case is 9 tokens:   3. Once the vocabulary size is reached  we are ready to tokenize new data With a trained Transformer on the vocabulary created previously  it is time to generate text based on a new input.   Imagine that the word bug is part of the input text  then it will be tokenized as [b  ug] because the whole word bug is not present in the vocabulary but the subwords b and ug are.   However  if the input data has a word like mug since  m is not part of the vocabulary  then it will be encoded as [<unk>  ug].   4. Encoding Just like every machine learning model  Transformers do not handle textual data  therefore  the vocabulary is encoded through a token ID. For example  the token hug becomes the ID 1  p becomes the ID 2  ug becomes the ID 3 and so on and so forth.   And that is how this tokenizer works!   Transformer: How does it work?The Transformer architecture [2] was developed in 2017 to perform language translations. It can be split into 5 different components: word embeddings  positional embeddings  encoder  decoder and next word prediction  as shown in Figure 2. Each of these components will be explained in detail in the following subsections.   1. Word Embedding The first component in a Transformer is the word embedding. This component is responsible for converting each token of the input text into a d-dimensional vector  but how?   Lets consider the example in Figure 3  where we want to translate the sentence Large Language Models from English to Portuguese. After tokenization  the sentence is converted into three different tokens each one representing a word from the input text.   These tokens go through a linear layer  in our example with 4 nodes  that converts the tokens into a 4-dimensional vector that will be consumed by the remaining components of the architecture to predict the next work in the sequence. During the training phase  the weights of this layer are optimised through backpropagation in order to improve the vector representation of a token which  consequently  will improve how well the model is able to predict the next word in the sequence.   2. Positional Embedding After the vector representation has been generated  the second component  Positional Embedding  kicks in.   Since the Encoder and Decoder are order invariant which means they cannot distinguish the connotation of two sentences with the same words but different meanings  for example LLaMA is better than Mistral and Mistral is better than LLaMA  Positional Embedding is used to solve this issue and add a sense of order to the model.   The Positional Embedding consists in a new vector to encode positions based on sinusoidal functions (sine and cosine) and the token position. It also depends on the model dimension since the positional vector is added to the word embedding  therefore  it has to have the same dimension of the word embedding.   Considering the previous example  the positional vector to be added to the word embedding needs to have a dimension of 4. For that  four sinusoidal functions are considered and each one of them will return a value based on the token position to generate the positional vector as shown in Figure 4. This vector is added to the word embedding before feeding the Encoder.   3. Encoder and Self-Attention Self-Attention is the key part of the Transformer architecture since it uses the similarity between words in a sequence in order to encode the current word in the sequence. It is based on three main components: Query  Key and Value.   3.1 Query Queries are used to represent each word in a sequence through a new d-dimensional vector. This vector is created by feeding the output of the word and positional embedding through a layer with multiple nodes  in our case 4 nodes  as shown in Figure 5. Just like in the word embedding layer  the weights linked to each node are updated through backpropagation in order to improve next word prediction.   3.2 Key Keys follow the exact same logic as Queries  where there is a layer that receives the combined output of word and positional embedding to generate a new d-dimensional vector  as shown in Figure 6. Once again the weights of this layer are updated through backpropagation to improve model performance.   3.3 Value Just like before  another independent layer from the ones used in Queries and Keys  is used to generate a new d-dimensional vector based on the combined output of word and positional embedding to represent the Value of a token  as shown in Figure 7.   3.4 Combining Query  Key and Value Once all the three main components were calculated  it is time to combine all of them to generate the encoded vector of a token. Continuing with the same sentence as before  Large Language Models  the first step is to calculate the similarity between tokens.   For that  we calculate the dot product between the Query of the token that is being processed and the Key of the remaining tokens in the sequence and its own Key. Based on the output of the dot product  a.k.a. similarity  a Softmax function is applied to return the weight that each token will have to create the encoded vector of the token Large  as shown in Figure 8.   These weights are going to be multiplied with the Values of each token to generate the contextual vector based on the surrounded tokens for the token Large.   Finally  this contextual vector is summed to the combined output of the word and positional embedding through a skip connection before feeding a Feed Forward Neural Network (FNN) that generates the final encoded vector.   4. Decoder The Decoder is responsible for generating the vector used in the next word prediction component based on the output of the encoder and all the tokens in the current sequence.   The first token to be processed is <SOS> which means Start of Sentence. Just like before  this token will be encoded through word embedding  positional embedding and self-attention.   After that  a skip connection is applied to sum the combined output of the word and position embedding to the output of the self-attention cell. The outcome of this process will go through another self-attention cell to be combined with the output of the Encoder  followed again by a skip connection. Just like in the Encoder  this final vector goes through a FNN to generate the vector that feeds the final component of the architecture  next word prediction.   5. Next Word Prediction The final output of the decoder goes through a simple neural network that will convert this vector into a new one with the same dimension as the vocabulary size. After that  a Softmax function is applied to determine which word should come next  as shown in Figure 13. The process stops when a <EOS> token (end of sentence) is predicted.   Final remarks about Self-Attention:   Query  Keys and Values are calculated in parallel  i.e.  Q  K and V are calculated at the same time for each word.We can have several Self-Attention Cells to train different weights and capture different relationships between words (the original had 8 cells that are combined through concatenation followed by a linear layer).There are multiple stacked layers (N=6) of self-attention cells + FFN which means the second layer input is the output of the first layer  the third layer input is the second layer output and so on and so forth.After each self-attention cell and FFN block  a normalization step is applied.LLaMA 3  Mixtral and Gemma: Whats new?This section will cover what are the main novel concepts  introduced by Meta  Mistral AI and Google with its respective models  LLaMA [3] [4] [5]  Mixtral [6] and Gemma [7].   Vocabulary Size  Context Length and ArchitectureAs we have seen previously in this article  the vocabulary size depends on the tokenizer adopted to feed the transformer. The tokenizer used in the original Transformer architecture had a vocabulary size of 32 000 tokens which is the same as in Mixtral. However  Gemma and LLaMA 3 have different tokenizers with 256 000 and 128 000 tokens  respectively.   Apart from the vocabulary size  the context length has been increasing over time since it has been demonstrated that a bigger context length  i.e. more tokens processed per sequence  yield to improved model performance. The original Transformer had a context length of 512 tokens  while Gemma and LLaMA 3 have 8192 and Mixtral 4096.   Also  on contrary to the original architecture that is an encoder-decoder Transformer for language translation  these models are decoder-only since their primary goal is to generate text.   Positional EmbeddingsLLaMA 3 and Gemma use Rotary Positional Embedding (RoPE) instead of the original version. This approach brings benefits such as modelling the relative position between tokens which means that the tokens in position 1 and 2 will be more similar than the tokens in position 1 and 500.   For example  lets consider the sentence Gemma is better than LLaMA and a 2D word embedding. The positional embedding of the word better will be given by a rotation of the original vector based on position 3 and a constant .   If two new words are added to the beginning of the sentence  then the angle between better and than will keep the same  as shown in Figure 15.   Grouped Query Attention instead of Multi Head AttentionLLaMA 3  Gemma and Mixtral replaced the traditional Multi Head Attention with Grouped Query Attention for faster decoding  hence  faster inference.   GQA-G divides query values into G groups that share a single key and value head (GQA-1 = MQA while a GQA-H = MHA). This approach reduces the number of keys and values heads into a single key and value per query group  accelerating the inference speed and reducing the memory requirements during decoding with a quality closer to MHA.   Activation functionLLaMA 3 and Gemma replaced the traditional ReLU activation function in the FFN block with SwiGLU and GeGLU  respectively. These functions  unlike ReLU that converts all negative values to 0  have a parameter that smooths this conversion where the probability of setting negative values to 0 increases as these values are closer to 0.   SMoE: Sparse Mixture of ExpertsMixtral differs from the other architectures by using Mixture of Experts rather than stacking a FFN on top of the different attention cells. Each Expert is responsible for processing a type of token  for example  one can be a punctuation expert  a visual description expert  or a number expert.   The Expert(s) that is going to process the token is chosen by a Gated Network trained to perform this allocation. This approach bring benefits such as efficiency by activating less model parameters and more accurate predictions because each expert is focused on a specific task.   DSPy: an easy to use library to get started with LLMsDSPy is an interesting library built around Signatures to test LLMs in different contexts or problems. The Signature allows to ask a LLM to perform different tasks without much prompt engineering or changes in the code  for example  we can perform sentiment analysis by using the signature sentence -> sentiment or summarization by using document -> summary or a RAG system using context  question -> answer.   Besides that  you can create personalized Signatures with just a few lines of code and  also  retrieve the reasoning that led the LLM to provide a certain answer by using ChainOfThought class. But  lets see it on practice.   First step is to import the libraries and setting up an env file with the HuggingFace (HF) token and  if you have a Open AI key  you can also use it with DSPy.   %load_ext autoreload %autoreload 2 import os import dspy from dotenv import load_dotenv load_dotenv('env/var.env')After that  we can load any model in HF repository or you can call ChatGPT by running the following lines of code.   In our case  we used mistralai/Mistral-7B-Instruct-v0.1 from HF.   chatgpt=dspy.OpenAI(api_key=os.getenv(OPENAI_KEY)  model='gpt-3.5-turbo-1106') lm = dspy.HFModel(model='mistralai/Mistral-7B-Instruct-v0.1'  token=os.getenv(HF_TOKEN)) dspy.configure(lm=lm)Sentiment Analysis   The Signature that allows to perform sentiment analysis is sentence -> sentiment. In our example  we created three different sentences about how AI is helpful  how uncertain it is if it brings more advantages than disadvantages and  finally  how threatening it is in order to try to capture three different sentiments.   As you can see below  with three lines of code we managed to extract sentiment from sentences using Mistral.   # positive sentiment sentence = AI has the potential to enhance human capabilities and automate tedious tasks  thereby improving our productivity and quality of life. classify = dspy.Predict('sentence -> sentiment') result = classify(sentence=sentence) print(result.sentiment) # neutral sentiment sentence = Poorly designed or uncontrolled AI systems pose risks such as job displacement  privacy violations  and even existential threats if advanced AI becomes misaligned with human values. result = classify(sentence=sentence) print(result.sentiment) # negative sentiment sentence = AI can bring existential threats if it becomes misaligned with human values. result = classify(sentence=sentence) print(result.sentiment)Positive Neutral Negative   Personalized Signature   In this case  we will create a personalized signature that aims to make Mistral classify product return reasons based on a list of reasons and the customer explanation for the return.   For that  we create a class called ReasonList that inherits a DSPy Signature and  in the output field  we define which classes the LLM must use to classify the input sentence. The huge benefit of this approach is that with just one line of code  we can make the LLM provide a formatted answer.    # an example below of the a custom signature and with a defined output class ReasonList(dspy.Signature): Classify reason among no_need  too_large  too_small  does_not_look_good sentence = dspy.InputField() reason = dspy.OutputField(desc=A list of values of any one of no_need  too_large  too_small or does_not_look_good  format=list[str]) sentence = I'm returning this item because my sister offered me a similar one classify = dspy.Predict(ReasonList) result = classify(sentence=sentence) print(result.reason) sentence = I'm returning this item because it is not to my taste result = classify(sentence=sentence) print(result.reason)no_need does_not_look_good   Summarization   The Signature that allows to perform summarization is document -> summary. In our example  we provide a huge text about the advantages and disadvantages of AI  and we ask Mistral to summarize it for us without providing any prompt and letting DSPy to handle that for us.   # an example below of the same signature but with different modules document = AI technologies hold great promise to enhance and augment human capabilities in a wide range of domains. One of the key benefits of AI is its ability to automate routine and repetitive tasks  freeing up human workers to focus on more complex  creative  and strategic work. AI-powered automation can drive gains in productivity  efficiency  and cost savings across industries. Additionally  AI systems excel at quickly processing and analyzing large datasets  identifying patterns and insights that may elude human cognition. This data-driven decision making can lead to more informed  evidence-based choices in fields like healthcare  finance  transportation  and scientific research. For example  AI algorithms can assist doctors in early disease detection  optimize logistics and supply chains  and accelerate drug discovery. Furthermore  AI has transformative potential in enhancing human experiences. AI-powered personal assistants  chatbots  and recommender systems can provide personalized assistance  information  and content tailored to individual needs and preferences. This can improve customer service  education  and quality of life. Advancements in natural language processing  computer vision  and robotic technologies also hold promise for assisting the elderly and disabled  improving accessibility  and extending human physical capabilities. Lastly  AI may play a vital role in addressing global challenges  such as climate change  resource scarcity  and disease outbreaks. AI systems can help optimize energy grids  model climate patterns  accelerate scientific research  and coordinate disaster response efforts more efficiently than human-only approaches. Of course  the development and deployment of AI also come with important ethical considerations and potential risks that must be carefully navigated. But overall  the transformative potential of AI to augment and empower humanity is profound and worth continued responsible exploration and investment. summarize = dspy.Predict('document -> summary') response = summarize(document=document) print(response)AI technologies hold great promise to automate routine tasks  drive productivity and efficiency gains  and enable data-driven decision making across industries. They also have the potential to enhance human experiences  address global challenges  and extend human physical capabilities. However  the development and deployment of AI also come with important ethical considerations and potential risks that must be carefully navigated.   RAG Systems   The Signature that allows to implement a RAG system is context  question -> answer. In our example  we provide a context about how the stock market and company valuation has been fluctuating in the last few weeks and we ask the LLM to retrieve which is the most valuable company in the end of the fluctuation period which it did correctly.    # an example below of the a custom signature using the basic RAG structure. context =  Context: Nvidia's stock has nearly tripled so far this year  compared with a rise of about 19% in Microsoft shares  with demand for its top-of-the-line processors outpacing supply. Tech giants Microsoft  Meta Platforms (META.O)  opens new tab and Google-owner Alphabet (GOOGL.O)  opens new tab are competing to build out their AI computing capabilities and add the technology to their products and services. An insatiable appetite for Nvidia's AI processors  viewed as far superior to competitors' offerings  has left them in tight supply  and many investors view Nvidia as the greatest winner to date from surging AI development. Although in the last few weeks we have been seing NVIDIA reaching the top of the most valued companies in the world  it was surpassed by Apple after the news about their heavy investment in AI.  question = What is the most valuable company? qa = dspy.ChainOfThought('context  question -> answer'  temperature=0.7) response = qa(context=context  question=question) print(response.answer)As of the recent developments  Apple has surpassed Nvidia and is currently the most valuable company  following their heavy investment in AI.   Since  we used ChainOfThought  we can also retrieve what was the reasoning for the LLM to get to that answer by running the following code:   print(lm.inspect_history(n=1))Reasoning: Lets think step by step in order to produce the answer. We need to consider the recent news and developments in AI technology and its impact on the value of companies.   Answer: As of the recent developments  Apple has surpassed Nvidia and is currently the most valuable company  following their heavy investment in AI.   ConclusionLarge Language Models have been around since the release of ChatGPT in 2022 and it is important to understand how these models work in order to extract their full potential to improve the way we live.   In this article  we went from the basics to the advanced concepts of the Transformer architecture to deeply understand how it works and how it is able to generate text in such accurate manner. Apart from that  we also explored how LLMs have been evolving since 2017  with big companies like Meta and Google releasing their own models leveraging the Transformer architecture with novel concepts.   Finally  the practical implementation of these models have been facilitated by packages like DSPy that remove the overhead of developing specific solutions for each LLM and allow to quickly perform several experiments to determine which LLM is more suitable for our use case.   Keep in touch: LinkedIn  Medium   References[1] Kudo  Taku  and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint arXiv:1808.06226 (2018).   [2] Vaswani et al.  Attention Is All You Need  arXiv preprint arXiv:1706.03762 (2017).   [3] Touvron et al.  LLaMA: Open and Efficient Foundation Language Models  arXiv preprint arXiv:2302.13971 (2023).   [4] Touvron et al.  Llama 2: Open Foundation and Fine-Tuned Chat Models  arXiv preprint arXiv:2307.09288 (2023).   [5] https://ai.meta.com/blog/meta-llama-3/   [6] Albert et al.  Mixtral of Experts  arXiv preprint arXiv:2401.04088 (2024).   [7] Gemma Team  Google DeepMind  Gemma: Open Models Based on Gemini Research and Technology  arXiv preprint arXiv:2403.08295 (2023).