Decoder-Based Large Language Models: A Complete Guide

Large Language Models (LLMs) have revolutionized the field of natural language processing (NLP) by demonstrating remarkable capabilities in generating human-like text, answering questions, and assisting with a wide range of language-related tasks. At the core of these powerful models lies the decoder-only transformer architecture, a variant of the original transformer architecture proposed in the seminal paper “Attention is All You Need” by Vaswani et al.

In this comprehensive guide, we will explore the inner workings of decoder-based LLMs, delving into the fundamental building blocks, architectural innovations, and implementation details that have propelled these models to the forefront of NLP research and applications.

The Transformer Architecture: A Refresher

Before diving into the specifics of decoder-based LLMs, it's essential to revisit the transformer architecture, the foundation upon which these models are built. The transformer introduced a novel approach to sequence modeling, relying solely on attention mechanisms to capture long-range dependencies in the data, without the need for recurrent or convolutional layers.

The original transformer architecture consists of two main components: an encoder and a decoder. The encoder processes the input sequence and generates a contextualized representation, which is then consumed by the decoder to produce the output sequence. This architecture was initially designed for machine translation tasks, where the encoder processes the input sentence in the source language, and the decoder generates the corresponding sentence in the target language.

Self-Attention: The Key to Transformer's Success

At the heart of the transformer lies the self-attention mechanism, a powerful technique that allows the model to weigh and aggregate information from different positions in the input sequence. Unlike traditional sequence models, which process input tokens sequentially, self-attention enables the model to capture dependencies between any pair of tokens, regardless of their position in the sequence.

The self-attention operation can be broken down into three main steps:

1. Query, Key, and Value Projections: The input sequence is projected into three separate representations: queries (Q), keys (K), and values (V). These projections are obtained by multiplying the input with learned weight matrices.
2. Attention Score Computation: For each position in the input sequence, attention scores are computed by taking the dot product between the corresponding query vector and all key vectors. These scores represent the relevance of each position to the current position being processed.
3. Weighted Sum of Values: The attention scores are normalized using a softmax function, and the resulting attention weights are used to compute a weighted sum of the value vectors, producing the output representation for the current position.

Multi-head attention, a variant of the self-attention mechanism, allows the model to capture different types of relationships by computing attention scores across multiple “heads” in parallel, each with its own set of query, key, and value projections.

Architectural Variants and Configurations

While the core principles of decoder-based LLMs remain consistent, researchers have explored various architectural variants and configurations to improve performance, efficiency, and generalization capabilities. In this section, we'll delve into the different architectural choices and their implications.

Architecture Types

Decoder-based LLMs can be broadly classified into three main types: encoder-decoder, causal decoder, and prefix decoder. Each architecture type exhibits distinct attention patterns.

Encoder-Decoder Architecture

Based on the vanilla Transformer model, the encoder-decoder architecture consists of two stacks: an encoder and a decoder. The encoder uses stacked multi-head self-attention layers to encode the input sequence and generate latent representations. The decoder then performs cross-attention on these representations to generate the target sequence. While effective in various NLP tasks, few LLMs, such as Flan-T5, adopt this architecture.

Causal Decoder Architecture

The causal decoder architecture incorporates a unidirectional attention mask, allowing each input token to attend only to past tokens and itself. Both input and output tokens are processed within the same decoder. Notable models like GPT-1, GPT-2, and GPT-3 are built on this architecture, with GPT-3 showcasing remarkable in-context learning capabilities. Many LLMs, including OPT, BLOOM, and Gopher, have widely adopted causal decoders.

Prefix Decoder Architecture

Also known as the non-causal decoder, the prefix decoder architecture modifies the masking mechanism of causal decoders to enable bidirectional attention over prefix tokens and unidirectional attention on generated tokens. Like the encoder-decoder architecture, prefix decoders can encode the prefix sequence bidirectionally and predict output tokens autoregressively using shared parameters. LLMs based on prefix decoders include GLM130B and U-PaLM.

All three architecture types can be extended using the mixture-of-experts (MoE) scaling technique, which sparsely activates a subset of neural network weights for each input. This approach has been employed in models like Switch Transformer and GLaM, with increasing the number of experts or total parameter size showing significant performance improvements.

Decoder-Only Transformer: Embracing the Autoregressive Nature

While the original transformer architecture was designed for sequence-to-sequence tasks like machine translation, many NLP tasks, such as language modeling and text generation, can be framed as autoregressive problems, where the model generates one token at a time, conditioned on the previously generated tokens.

Enter the decoder-only transformer, a simplified variant of the transformer architecture that retains only the decoder component. This architecture is particularly well-suited for autoregressive tasks, as it generates output tokens one by one, leveraging the previously generated tokens as input context.

The key difference between the decoder-only transformer and the original transformer decoder lies in the self-attention mechanism. In the decoder-only setting, the self-attention operation is modified to prevent the model from attending to future tokens, a property known as causality. This is achieved through a technique called “masked self-attention,” where attention scores corresponding to future positions are set to negative infinity, effectively masking them out during the softmax normalization step.

Architectural Components of Decoder-Based LLMs

While the core principles of self-attention and masked self-attention remain the same, modern decoder-based LLMs have introduced several architectural innovations to improve performance, efficiency, and generalization capabilities. Let's explore some of the key components and techniques employed in state-of-the-art LLMs.

Input Representation

Before processing the input sequence, decoder-based LLMs employ tokenization and embedding techniques to convert the raw text into a numerical representation suitable for the model.

Tokenization: The tokenization process converts the input text into a sequence of tokens, which can be words, subwords, or even individual characters, depending on the tokenization strategy employed. Popular tokenization techniques for LLMs include Byte-Pair Encoding (BPE), SentencePiece, and WordPiece. These methods aim to strike a balance between vocabulary size and representation granularity, allowing the model to handle rare or out-of-vocabulary words effectively.

Token Embeddings: After tokenization, each token is mapped to a dense vector representation called a token embedding. These embeddings are learned during the training process and capture semantic and syntactic relationships between tokens.

Positional Embeddings: Transformer models process the entire input sequence simultaneously, lacking the inherent notion of token positions present in recurrent models. To incorporate positional information, positional embeddings are added to the token embeddings, allowing the model to distinguish between tokens based on their positions in the sequence. Early LLMs used fixed positional embeddings based on sinusoidal functions, while more recent models have explored learnable positional embeddings or alternative positional encoding techniques like rotary positional embeddings.

Multi-Head Attention Blocks

The core building blocks of decoder-based LLMs are multi-head attention layers, which perform the masked self-attention operation described earlier. These layers are stacked multiple times, with each layer attending to the output of the previous layer, allowing the model to capture increasingly complex dependencies and representations.

Attention Heads: Each multi-head attention layer consists of multiple “attention heads,” each with its own set of query, key, and value projections. This allows the model to attend to different aspects of the input simultaneously, capturing diverse relationships and patterns.

Residual Connections and Layer Normalization: To facilitate the training of deep networks and mitigate the vanishing gradient problem, decoder-based LLMs employ residual connections and layer normalization techniques. Residual connections add the input of a layer to its output, allowing gradients to flow more easily during backpropagation. Layer normalization helps to stabilize the activations and gradients, further improving training stability and performance.

Feed-Forward Layers

In addition to multi-head attention layers, decoder-based LLMs incorporate feed-forward layers, which apply a simple feed-forward neural network to each position in the sequence. These layers introduce non-linearities and enable the model to learn more complex representations.

Activation Functions: The choice of activation function in the feed-forward layers can significantly impact the model's performance. While earlier LLMs relied on the widely-used ReLU activation, more recent models have adopted more sophisticated activation functions like the Gaussian Error Linear Unit (GELU) or the SwiGLU activation, which have shown improved performance.

Sparse Attention and Efficient Transformers

While the self-attention mechanism is powerful, it comes with a quadratic computational complexity with respect to the sequence length, making it computationally expensive for long sequences. To address this challenge, several techniques have been proposed to reduce the computational and memory requirements of self-attention, enabling efficient processing of longer sequences.

Sparse Attention: Sparse attention techniques, such as the one employed in the GPT-3 model, selectively attend to a subset of positions in the input sequence, rather than computing attention scores for all positions. This can significantly reduce the computational complexity while maintaining reasonable performance.

Sliding Window Attention: Introduced in the Mistral 7B model , sliding window attention (SWA) is a simple yet effective technique that restricts the attention span of each token to a fixed window size. This approach leverages the ability of transformer layers to transmit information across multiple layers, effectively increasing the attention span without the quadratic complexity of full self-attention.

Rolling Buffer Cache: To further reduce memory requirements, especially for long sequences, the Mistral 7B model employs a rolling buffer cache. This technique stores and reuses the computed key and value vectors for a fixed window size, avoiding redundant computations and minimizing memory usage.

Grouped Query Attention: Introduced in the LLaMA 2 model, grouped query attention (GQA) is a variant of the multi-query attention mechanism that divides attention heads into groups, each group sharing a common key and value matrix. This approach strikes a balance between the efficiency of multi-query attention and the performance of standard self-attention, providing improved inference times while maintaining high-quality results.