LLMs (Part-04): Inside the Decoder Stack

A Deep Dive Into the 6-block Decoder Stack in the Transformer Neural Network (2017)

Greetings from the author! If this is the first story that you have come across on transformers, don’t miss out on the entire Large Language Models (LLMs) series — explaining in great detail the transformer architecture, encoder & decoder stacks, blocks and layers, embedding & un-embedding layers, mathematical operations inside LLMs, feature extraction, feed forward mechanism — and a lot more. We also have a great multi-part Deep Learning series as well; don’t forget to check it out too!

Part-01: Decoder’s Internal Working Overview

The decoder block processes information as follows. Each of the 3 main layers in each of the six decoder blocks performs these actions in the exact same sequence. The decoder has three main layers, written below. Notice that after each of the mai layers generates its output, we apply add (residual connections) & norm (normalization).

• Masked Multi-head Self-attention => Add & Norm operations
• Multi-head Cross-attention => Add & Norm operations
• Feed-Forward Network => Add & Norm operations

Part-02: Masked Multi-head Self-attention

( a ) Self-attention

Here, self-attention means that the decoder considers the “attention scores” among all the tokens of the entire generated output sequence so far (so far = masked). Attention has the same meaning in this context as in the context of the encoder stack: figuring out the relationships or relevance scores among all the generated output tokens so far.

( b ) Masking During Inference

The term “masked” means that the decoder is going to pay attention to only those tokens that it has already generated in the output sequence so far — not the future ones.

For example, “the cat sat down on the mat, because it was tired” can be our full output sequence. If the model is generating the word “sat”, it can only attend to the words “the” and “cat”, not anything else that has not yet been generated.

( c ) Masking During Training

During the training, however, the decoder has access to the full training data, which means knowing ahead-of-time what and how many tokens will be generated in the entire given output sequence of embeddings (hidden states). So, we “mask” or hide the output tokens from the decoder that are in front of the x-th position (where x is the position at which the token is being predicted), and we keep on shifting the mask to the right. So, during training, the decoder learns to pay attention to x — 1 tokens of the output sequence (where x is the token at the x-th position being predicted; x <= n), without looking at the next tokens. This, in production, helps the decoder to generate output tokens without even having any future tokens available.

Masking happens during each forward pass phase of the transformer for predicting every single token, i.e., in each forward pass, we only predict one token. This essentially means that all the self-attention layers in all the six blocks of the decoder stack can attend only to the token at the current (x-th) position and the positions below (previous to) it.

( d ) Time Complexity

The time complexity for decoder’s masked self-attention turns out to be O(n²). But …. the process is not finished yet; the decoder does not predict next token based on masked multi-head self-attention alone; it also pays attention to the input sequence as well — or, more precisely, all the encoder output states (high-dimensional vector embeddings) to figure out which parts of the input are relevant for generating the next token in the output sequence.

( e ) Decoding is Expensive at Inference Time

Remember, during inference, the transformer only goes over the encoder stack once and finalizes its output state. The decoder now refers to it again and again for predicting each token. To predict a single token, we loop over the entire decoder stack once, and each time, we refer to the output state of the encoder from the cache (variable) stored in the V-RAM (RAM of the GPU / TPU, not the system RAM).

Part-03: Multi-head Cross-attention

In the original 2017 paper, multi-head cross-attention is referred to as “Multi-head Encoder-Decoder” attention.

( a ) Cross-attention

The core idea, intuitively, is simple: the decoder asks the encoder, “which parts of the input are the most relevant for me to be able to successfully predict the next output token?” The answer to this question for the decoder comes from the encoder’s output state of the processed user’s input — keeping in view the decoder’s own hidden state after self-attention (cross-attention). Now, keeping in view both these states (decoder’s hidden state and encoder’s output state and FFN), the decoder predicts the next output token.

( b ) Multi-head

Again, the decoder computes attention scores 8 times in parallel, one in each of the 8 heads over the output state of the encoder. The purpose is to compute scores for the output states (vectors) generated by the encoder to capture which output states of the encoder (a vector representing a token) the decoder should pay the most attention to in order to generate the next token. The reason we have 8

( c ) Decoder’s Query & Encoder’s KV Cache

The Full Process looks like this:

1. The Decoder produces a Query ($Q$) representing “what I am looking for right now.”
2. The Encoder Output provides the Keys ($K$) and Values ($V$) from the system VRAM (GPU’s / TPU’s memory).
3. The system calculates attention scores by comparing $Q$ and $K$.
4. These scores are applied to $V$ to extract the relevant context.

Now, after multi-head self-attention and multi-head cross-attention (and FFN), we predict the next token in the decoder block.

Part-04: Feed Forward Network

In the FFN layer of any of the six given transformer decoder blocks, we perform three mathematical calculations:

( a ) First Linear Combination

*z = (W1 X) + b1**

( b ) Non-linear Activation

h = activation(z);

The activation function is usually RELU or GELU, not Sigmoid in transformers.

( c ) Second Linear Combination

*y = (W2 h) + b2**

Part-05: Output (Token) Prediction

Once masked multi-head self-attention and multi-head cross-attention and FFN are done in each of the six blocks of the entire decoder stack, we are able to successfully make a prediction to generate the next token in the sequence that is well-aligned with the tone, emotionality, logic and coherence of the previously generated text as well as keeping in view the user’s input.

However, keep in mind that this output state of the decoder is not an actual word; it is merely a latent (hidden) state vector, which is passed through the un-embedding layer and a softmax layer to be converted to a natural language word.

( a ) Decoding is a Multi-class Classification Problem

Remember, the decoder is actually trying to solve a multi-class classification problem. Yes, it is a classification problem. Whe the decoder performs self and cross attentions combined with FFN, add and norm to predict the next token, it is actually generating a hidden state (latent representation) that must be within the vocabulary of the LLM / transformer so that it can be converted to a token and then a natural language word.

In other words:

• Each of the neurons in the LM head is “representing” a single token of the vocabulary, and
• Each token is actually a “class” that the latent state of the decoder must fit into — or resemble with high level of similarity.

Part-06: Summary

Because the input sequence does not change and is already considered final during the auto-regressive generation of the output, the encoder’s computation is performed only once at the beginning of the inference process, and its results are reused for all subsequent token predictions until an end-of-sequence token is generated. The encoder’s final output state is stored in the VRAM (the memory of the GPU / TPU), and, at the level of software, it is represented in the form of a variable.

Until Next Time,

0s & 1s