← All posts

Architecture7 min read

Transformers: how generative language models work

The Transformer architecture behind large language models, and the generative pre-training that brings them to life.

Parth Shastri

The age of AI !! The future is here !! This is what I hear everywhere I go now-a-days, because the impact of the Generative AI has been that high. The Language models that are able to generate and respond to content and interact with users have proven to be the next big thing. Let us go into the working or rather decoding of the Transformer decoder based Large Language Models!

An abstract image representing generative AI and large language models.

The parent: the Transformer#

The base behind the Large Language Models is the Transformer. Transformer is a type of architecture introduced first in 2017 by Vaswani et al. in their seminal paper “Attention Is All You Need”. It was introduced for Seq-2-seq tasks like Neural machine translation, but since its conception it has evolved into something big and is used around in many modalities apart from Natural Language Processing (NLP).

The Transformer architecture consists of an encoder and a decoder as shown in Figure 1. The encoder is used to model the context and the decoder is used for the auto-regressive decoding of the output. So if your input is a sentence in English and you want it translated to Hindi, then the encoder will encompass the information about the English sentence and this information will be used in the decoding phase by the decoder auto-regressively.

The encoder-decoder Transformer architecture.
The encoder-decoder Transformer architecture, as depicted in "Attention is all you need".

So this was the context set by the paper “Attention is all you need!” (a great read, I recommend you to read it). After this paper a lot of improvements and experiments were done around the transformer architecture, with encoder only models (BERT etc.) and decoder only models (GPT, Gopher etc.), also scaling of the encoder-decoder architecture (T5 etc.) The reason that the transformer architecture was widely researched was, they were computationally efficient to train than their counterparts, they leveraged self-attention, they assumed no inductive biases making them widely adaptive and they were data-hungry giving great results with more data !! Hush… that was too much information in a single line ! I recommend you to read this excellent blogpost by Jay Alammar to understand the Transformer better !

The Transformer that can generate!#

As we saw earlier the main objective behind the release of the vanilla Transformer was improving the existing models that can perform, Neural Machine Translation (NMT). The evolution of this architecture led to the recent developments in LLMs. Moving on, the Decoder architecture was the one that was able to produce amazing results in the world of LLMs. The LLM training is based on a technique called Generative Pretraining. As the name suggests we use generative pre-training. Don’t get intimidated by the word, it means teaching the model to generate words/tokens on its own by showing it a lot of raw data. If we think about it it kind of resembles how we humans learn to generate/create our own words, we listen or read a lot of text/voice and try to formulate our own understanding and then find out a way of putting out what we think or simply put forth new ideas.

So how is this “Generative Pre-training” exactly done? It is done by teaching the model to predict the probability of the next word in a sentence given a context sentence. It is also commonly known as language modelling in ML and NLP. The idea of Generative Pretraining using transformer decoders was proposed by Radford et al. in the original GPT paper. This paper proposed the use of Transformer decoder as their main architecture.

Generative pretraining#

The task of generative pretraining or predicting next word given a context can be mathematically represented as a multiplication of conditional probability distributions. The below equation depicts the typical left-to-right probability objective of Language Models.

P(x)=∏i=1nP(xi∣x1,…,xi−1)\boxed{P(x) = \prod_{i=1}^n P(x_i|x_1, \dots, x_{i-1})}
The training phase of a decoder, with input tokens shifted by one.
The training phase of the decoder. The input tokens are shifted by one token so that at each step the decoder learns to predict the next token.

In practice however, the decoder is given a start token/symbol and is then auto-regressively sampled for new text generations. Auto-regressively meaning that at each timestep the context and the output of the previous timestep is concatenated and given as the input. This means, that the decoder is iterated through the number of maximum tokens to be generated, passing the initial input plus the previous generated token at each iteration (Figure 3). These models are typically trained to minimize the cross-entropy loss between the predicted tokens and the labels. Although some kind of regularizers and additional loss terms are added in some cases, typically for better stability and performance.

Auto-regressive sampling: the model predicts a token, appends it, and predicts again.
Auto-regressive sampling of the decoder. Given an input the model predicts the next token, then the combined input goes in at the next step. Credits: Jay Alammar, The Illustrated GPT-2.

You can play the same game yourself below. Pick greedy a few times, then reset & sample at a high temperature to see why the same prompt can give different answers.

Interactive

The cat sat on the

    Decoding one token at a time. The bars are softmax(logits / T) for the next token. Greedy always takes the top bar. Sampling rolls the dice. The numbers come from a toy model, not a real one.

    So this is a high-level overview of generative pretraining! Once these models are trained for a fairly simple objective, on a large corpus of data, they are able to generalize well on language and any downstream tasks like summarization, sentiment classification, writing essays. These models can be made to generalize using simple language and also recently have shown to follow instructions. We can finetune these models to enhance performance using techniques like prompt tuning, traditional finetuning etc.

    The recent models that we are able to see are the large versions of the Decoder-only transformers fine-tuned to follow instructions using a technique called RLHF (Reinforcement Learning using Human Feedback). This was popularized by a paper called InstructGPT, by OpenAI. Later models like Gopher were later trained using a similar objective to follow instructions. These models are now popularized as Chat LLMs which are proving to be a pivotal discovery in helping to increase productivity and taking us a step closer to AGI !

    The architecture#

    The architecture of recent LLMs is similar to the “decoder” part of the transformer architecture, depicted in Figure 1. The architecture is described in detail in Figure 5. The input text is first tokenized into discrete token mappings of a fixed length using algorithms like Sentence Piece, WordPiece Tokenization. If you want to explore pre-processing for NLP, I recommend you to read through the Hugging Face course for NLP. This tokenized text is then embedded into a continuous vector using the embedding layer which is simply a matrix that maps the integer numbers to a learnable vector. Position information is added to this using positional embeddings (learnable or fixed). These obtained embeddings are then passed through the transformer block that employs Multiheaded Self Attention followed by a linear layer, to obtain a probability distribution over all the possible outcomes for the next token.

    The decoder-only architecture and its downstream tasks.
    The overall view of the decoder-only architecture, along with applications of the model on various downstream tasks.

    The recent versions of the LLM model family, GPT-3, GPT-3.5, GPT-4, PaLM etc. use a similar architecture while increasing the parameters in the model and also some optimization to increase the token generation capability. These models have gradually shown to exhibit great few-shot and zero-shot learning capabilities with less number of training examples or no examples at all.

    Conclusion#

    We saw the fuss around what exactly are LLMs or large language models made of, decoding the fairly simple idea behind the construction of these models and also saw the recent developments around these models. Further we also saw how are they trained and what objective do they use, also highlighting their use cases in various downstream tasks.

    References#

    1. Radford, Alec, et al. “Improving Language Understanding by Generative Pre-Training.” (2018).
    2. Alammar, Jay. “The Illustrated Transformer.” (2018).
    3. Rae, Jack W., et al. “Scaling Language Models: Methods, Analysis & Insights from Training Gopher.” arXiv preprint arXiv:2112.11446 (2021).
    4. Alammar, Jay. “The Illustrated GPT-2 (Visualizing Transformer Language Models).” (2019).
    5. Vaswani, Ashish, et al. “Attention is all you need.” Advances in neural information processing systems 30 (2017).
    6. Ouyang, Long, et al. “Training language models to follow instructions with human feedback.” arXiv preprint arXiv:2203.02155 (2022).