Transformer is a neural network architecture that has fundamentally changed the approach to
Artificial Intelligence. Transformer was first introduced in the seminal paper
"Attention is All You Need"
in 2017 and has since become the go-to architecture for deep learning models, powering text-generative
models like OpenAI's GPT, Meta's Llama, and Google's
Gemini. Beyond text, Transformer is also applied in
audio generation,
image recognition,
protein structure prediction, and even
game playing, demonstrating its versatility across numerous domains.
Fundamentally, text-generative Transformer models operate on the principle of next-token prediction: given a text prompt from the user, what is the
most probable next token (a word or part of a word) that will follow this input? The core
innovation and power of Transformers lie in their use of self-attention mechanism, which allows
them to process entire sequences and capture long-range dependencies more effectively than previous
architectures.
GPT-2 family of models are prominent examples of text-generative Transformers. Transformer
Explainer is powered by the
GPT-2
(small) model which has 124 million parameters. While it is not the latest or most powerful Transformer
model, it shares many of the same architectural components and principles found in the current
state-of-the-art models making it an ideal starting point for understanding the basics.
Every text-generative Transformer consists of these three key components:
Let's say you want to generate text using a Transformer model. You add the prompt like this
one: “Data visualization empowers users to”. This input needs to be converted
into a format that the model can understand and process. That is where embedding comes in: it
transforms the text into a numerical representation that the model can work with. To convert a
prompt into embedding, we need to 1) tokenize the input, 2) obtain token embeddings, 3) add
positional information, and finally 4) add up token and position encodings to get the final
embedding. Let’s see how each of these steps is done.
Tokenization is the process of breaking down the input text into smaller, more manageable
pieces called tokens. These tokens can be a word or a subword. The words "Data"
and "visualization" correspond to unique tokens, while the word
"empowers"
is split into two tokens. The full vocabulary of tokens is decided before training the model:
GPT-2's vocabulary has 50,257 unique tokens. Now that we split our input text into
tokens with distinct IDs, we can obtain their vector representation from embeddings.
GPT-2 (small) represents each token in the vocabulary as a 768-dimensional vector; the
dimension of the vector depends on the model. These embedding vectors are stored in a matrix
of shape (50,257, 768), containing approximately 39 million parameters! This
extensive matrix allows the model to assign semantic meaning to each token, in the sense
that tokens with similar usage or meaning in language are placed close together in this
high-dimensional space, while dissimilar tokens are farther apart.
The Embedding layer also encodes information about each token's position in the input
prompt. Different models use various methods for positional encoding. GPT-2 trains its own
positional encoding matrix from scratch, integrating it directly into the training process.
Finally, we sum the token and positional encodings to get the final embedding
representation. This combined representation captures both the semantic meaning of the
tokens and their position in the input sequence.
The core of the Transformer's processing lies in the Transformer block, which comprises
multi-head self-attention and a Multi-Layer Perceptron layer. Most models consist of multiple
such blocks that are stacked sequentially one after the other. The token representations
evolve through layers, from the first block to the last one, allowing the model to build up an
intricate understanding of each token. This layered approach leads to higher-order
representations of the input. The GPT-2 (small) model we are examining consists of 12 such blocks.
The self-attention mechanism enables the model to capture relationships among tokens in a
sequence, so that each token’s representation is influenced by the others. Multiple attention
heads allow the model to consider these relationships from different perspectives; for
example, one head may capture short-range syntactic links while another tracks broader
semantic context. In the following section, we will walk through how multi-head self-attention
is computed step by step.
Each token's embedding vector is transformed into three vectors:
Query (Q),
Key (K), and
Value (V). These vectors are derived by multiplying the input
embedding matrix with learned weight matrices for
Q,
K, and
V. Here's a web search analogy to help us build some intuition
behind these matrices:
By using these QKV values, the model can calculate attention scores, which determine how
much focus each token should receive when generating predictions.
Query, key, and
Value
vectors are split into multiple heads—in GPT-2 (small)'s case, into
12 heads. Each head processes a segment of the embeddings independently, capturing
different syntactic and semantic relationships. This design facilitates parallel learning of
diverse linguistic features, enhancing the model's representational power.
In each head, we perform masked self-attention calculations. This mechanism allows the model
to generate sequences by focusing on relevant parts of the input while preventing access to
future tokens.
The model uses the masked self-attention scores and multiplies them with the
Value matrix to get the
final output
of the self-attention mechanism. GPT-2 has 12 self-attention heads, each capturing
different relationships between tokens. The outputs of these heads are concatenated and passed
through a linear projection.
After the multiple heads of self-attention capture the diverse relationships between the input
tokens, the concatenated outputs are passed through the Multilayer Perceptron (MLP) layer to
enhance the model's representational capacity. The MLP block consists of two linear
transformations with a GELU activation function in between.
The first linear transformation expands the dimensionality of the input four-fold from 768
to
3072. This expansion step allows the model to project the token representations
into a higher-dimensional space, where it can capture richer and more complex patterns that
may not be visible in the original dimension.
The second linear transformation then reduces the dimensionality back to the original size of 768.This compression step brings the representations back to a manageable size while retaining
the useful nonlinear transformations introduced in the expansion step.
Unlike the self-attention mechanism, which integrates information across tokens, the MLP
processes tokens independently and simply maps each token representation from one space to
another, enriching the overall model capacity.
After the input has been processed through all Transformer blocks, the output is passed
through the final linear layer to prepare it for token prediction. This layer projects the
final representations into a 50,257
dimensional space, where every token in the vocabulary has a corresponding value called
logit. Any token can be the next word, so this process allows us to simply rank
these tokens by their likelihood of being that next word. We then apply the softmax function
to convert the logits into a probability distribution that sums to one. This will allow us to
sample the next token based on its likelihood.
The final step is to generate the next token by sampling from this distribution The temperature
hyperparameter plays a critical role in this process. Mathematically speaking, it is a very simple
operation: model output logits are simply divided by the
temperature:
In addition, the sampling process can be further refined using top-k
and
top-p parameters:
By tuning temperature, top-k, and top-p, you can
balance between deterministic and diverse outputs, tailoring the model's behavior to your
specific needs.
There are several auxiliary architectural features that enhance the performance of Transformer
models. While important for the model's overall performance, they are not as important for
understanding the core concepts of the architecture. Layer Normalization, Dropout, and
Residual Connections are crucial components in Transformer models, particularly during the
training phase. Layer Normalization stabilizes training and helps the model converge faster.
Dropout prevents overfitting by randomly deactivating neurons. Residual Connections allows
gradients to flow directly through the network and helps to prevent the vanishing gradient
problem.
Layer Normalization helps to stabilize the training process and improves convergence. It
works by normalizing the inputs across the features, ensuring that the mean and variance of
the activations are consistent. This normalization helps mitigate issues related to internal
covariate shift, allowing the model to learn more effectively and reducing the sensitivity
to the initial weights. Layer Normalization is applied twice in each Transformer block, once
before the self-attention mechanism and once before the MLP layer.
Dropout is a regularization technique used to prevent overfitting in neural networks by
randomly setting a fraction of model weights to zero during training. This encourages the
model to learn more robust features and reduces dependency on specific neurons, helping the
network generalize better to new, unseen data. During model inference, dropout is
deactivated. This essentially means that we are using an ensemble of the trained
subnetworks, which leads to a better model performance.
Residual connections were first introduced in the ResNet model in 2015. This architectural
innovation revolutionized deep learning by enabling the training of very deep neural
networks. Essentially, residual connections are shortcuts that bypass one or more layers,
adding the input of a layer to its output. This helps mitigate the vanishing gradient
problem, making it easier to train deep networks with multiple Transformer blocks stacked on
top of each other. In GPT-2, residual connections are used twice within each Transformer
block: once before the MLP and once after, ensuring that gradients flow more easily, and
earlier layers receive sufficient updates during backpropagation.
Transformer Explainer is built to be interactive and allows you to explore the inner workings
of the Transformer. Here are some of the interactive features you can play with:
Transformer Explainer features a live GPT-2 (small) model running directly in the browser.
This model is derived from the PyTorch implementation of GPT by Andrej Karpathy's
nanoGPT project
and has been converted to
ONNX Runtime
for seamless in-browser execution. The interface is built using JavaScript, with
Svelte
as a front-end framework and
D3.js
for creating dynamic visualizations. Numerical values are updated live following the user input.
Transformer Explainer was created by
Aeree Cho,
Grace C. Kim,
Alexander Karpekov,
Alec Helbling,
Jay Wang,
Seongmin Lee,
Benjamin Hoover, and
Polo Chau
at the Georgia Institute of Technology.
Source: Hacker News
You must be logged in to post a comment.