How does the (decoder-only) transformer architecture work?

Question

How does the (decoder-only) transformer architecture work which is used in impressive models such as GPT-4?

Answer 1

Introduction

Large-language models (LLMs) have gained tons of popularity lately with the releases of ChatGPT, GPT-4, Bard, and more. All these LLMs are based on the transformer neural network architecture. The transformer architecture was first introduced in the paper "Attention is All You Need" by Google Brain in 2017. LLMs/GPT models use a variant of this architecture called de' decoder-only transformer'.

The most popular variety of transformers are currently these GPT models. The only purpose of these models is to receive a prompt (an input) and predict the next token/word that comes after this input. Nothing more, nothing less.

Note: Not all large-language models use a transformer architecture. However, models such as GPT-3, ChatGPT, GPT-4 & LaMDa use the (decoder-only) transformer architecture.

Overview of the (decoder-only) Transformer model

It is key first to understand the input and output of a transformer:

• The input is a prompt (often referred to as context) fed into the transformer as a whole. There is no recurrence.
• The output depends on the goal of the model. For GPT models, the output is a probability distribution of the next token/word that comes after the prompt. It outputs one prediction for the complete input.

Next, it is essential to understand the key components that make up the decoder-only transformer architecture:

1. The embedding: the input of the transformer model is a prompt. This prompt needs to be embedded into something that the model can use.
2. The block(s): This is the main source of complexity. Each block contains a masked multi-head attention submodule, a feedforward network, and several layer normalization operations. Blocks are put in sequence to make the model deeper.
3. The output: the output of the last block is fed through one more linear layer to obtain the final output of the model (a classification, a next word/token etc.)

The following visualization gives an overview of the transformer architecture.

Self-Attention Mechanism

Self-attention makes the transformer powerful. The intuition of self-attention is that the mechanism allows the model to focus on (attend to) the most relevant parts of the input. A single self-attention mechanism is called a head.

The head works as follows. First, the input is fed into three separate linear layers. Two of those (the queries Q and the keys K) are multiplied, scaled, and turned into a probability distribution using a softmax activation function. Think of this probability distribution as describing which indices matter most for the output (i.e. which words in the prompt matter for the next word to be predicted). Finally, the output is multiplied with values V. This thus gives V * the importance of each of the tokens in V. A key observation is that the learnable parameters in the head are the three linear layers.

The following figure gives an overview of the operations done in a head and an overview of how multi-head attention works.

Multi-Head Attention

Multi-head attention is nothing more than several individual heads stacked on top of one another. The input to all heads is equivalent. However, each head has its own weights. After forwarding the input through all the heads, the output of the heads is concatenated and passed through a linear layer which brings the dimensionality back to the dimension of the initial input.

Masked Self-Attention

In the decoder-only transformer, masked self-attention is nothing more than sequence padding. The 'masking' term is a left-over of the original encoder-decoder transformer module in which the encoder could see the whole (original language) sentence, and the decoder could only see the first part of the sentence which was already translated. As such, they called it 'masking'.

Block

Each block contains a multi-head attention submodule, a feedforward network, 2 layer-normalization operations, and 2 skip connections.

The feedforward network is simply a multi-layer perceptron. In the original paper, the proposed feedforward module consisted of (1) a fully connected layer; (2) a ReLU activation; (3) another fully connected layer; and (4) a dropout layer.

The 'add & norm' blocks get the output from the multi-head attention/feedforward submodule and add it to the input into those modules. After that, a layer normalization operation is performed. Adding the input and output of a submodule together is known as a skip-connection. As blocks can be put in sequence, the skip connections help tremendously reduce the problem of vanishing or exploding gradients. In other words, skip connections are necessary to ensure proper backpropagation of the gradients.

Positional Embedding

Transformers take in a complete prompt at once (in contrast to RNNs) and embed this as one big Tensor. As such, transformers do not know which word is at what position in the sentence. This is problematic as the following two sentences mean entirely different things, only dependent on the order of the words:

The boy chased the bird with a butterfly net.

The bird chased the boy with a butterfly net.

To that end, a positional embedding is added to allow the model to deduce which word is where. Positional embeddings can be learned using any embedding layer from your favourite AI library. However, the original authors proposed a much more complicated method, which does not require learning any parameters. Please find an elaborate explanation (not mine) on the original positional embedding here.

Output

After the prompt is forwarded through all the blocks sequentially, the output is forwarded through one final linear layer. This final linear layer maps the output of the model back to the size of the 'vocabulary'. I.e. if you want to predict the next letter in a message, it would map to 26 (letters) + additional stuff (such as .,-!? etc.).

The output of the model is a probability distribution. For GPT models, the output is the probability of each token being the next token in the sequence.

Training

Basic training

The basic training process consists of self-supervised learning. Simply put, you gather lots of text, strip the last word from that text, feed it as input into the transformer, check if the prediction matches the word you cut off and backpropagate the error.

Every text/sentence/book/webpage can be separated into several samples.

sample = [
    ["This"],
    ["This", "is"],
    ["This", "is", "a"]] # padding is added until the max-sequence length is reached.
targets = ["is", "a", "sample"]

Fine-tuning / transfer-learning

After the first stage of training is completed, the model is now a large-language model. As in, it can predict the next word based on a context. However, through fine-tuning/transfer-learning the model can be adapted to better suit the needs of the final application.

One of the key reasons why ChatGPT & GPT-4 seem so ridiculously impressive is because of this second stage of training. In this stage, the following process is executed many times:

1. The model is given a prompt and generates different answers
2. The different answers are ranked by a human from best to worst.
3. The scores of the different answers are backpropagated.

However, transformer models can also be used for different tasks than language generation. They can, for example, be used for sentiment analysis. After doing the basic training, a transformer can be fine-tuned for sentiment analysis by removing the outgoing linear layer and replacing it with a different layer suitable for the task to be executed. Consequently, it can be trained in a supervised fashion on a custom (sentiment analysis) dataset.

Inference (answer generation)

Doing inference with a transformer is just like training. You insert a prompt and out comes the next word/classification/other.

For GPT models, this means that the prompt is extended one word at a time. You insert the prompt, and out comes the first word of the answer. The first word of the answer is now added to the prompt, creating a new, slightly different prompt. This prompt is again forwarded through the model, giving the prediction of a new word.

As the output is the probability for each token to be the next one, you can do several things during inference. For one, you can sample from the probability distribution. This induces some randomness into the algorithm. You can also take the token that has the highest probability; then, the model becomes deterministic.

How can a transformer do X?

The transformer can do X because it has seen enough examples of similar sentences in its training to give a satisfactory output to your prompt. However, for the question 'Why can a transformer not do X?' the answers are vastly more diverse.