The Transformer: Unveiling the Secrets of Modern AI

The Transformer: Unveiling the Secrets of Modern AI

In the ever-evolving world of Artificial Intelligence, one architecture has emerged as a true game-changer: the Transformer. This article dives deep into the inner workings of this powerful tool, breaking down complex concepts into easily digestible explanations. Get ready to understand how the Transformer is revolutionizing everything from language translation to image recognition.

Attention: The Core of the Transformer

At the heart of the Transformer lies the concept of attention. Imagine reading a long article; you don't focus equally on every word. Instead, you pay closer attention to the words and phrases that are most relevant to understanding the meaning. The Transformer's attention mechanism works similarly, allowing it to selectively focus on different parts of the input when processing information.

How Attention Works

Let's break down how this attention mechanism actually functions. Think of it as a way for a model to represent a token (a single element, like a word) by intelligently gathering context from other related tokens that came before it. In simple terms, the model looks at previous tokens and integrates only the most vital pieces to construct a new representation for the current token.

For language models, especially those working from left to right, the context is essentially anything the model has already processed. This means that when the model is understanding a word, it can use all previous words to inform its understanding, usually thousands of tokens. As you will see later, the "attention" part can also look to future words.

Imagine you are trying to understand the sentence, "The cat sat on the mat because it was tired." To understand what "it" refers to, you need to pay attention to "cat". The attention mechanism helps the model make this connection.

A Simplified View of Attention

To get a clearer picture, let's simplify things. Essentially, attention is a weighted sum of all the context vectors. The weights determine how much each context vector contributes to the final output. If we let ai stand for the attention output at token position i, and xj represent the tokens, and use αij to indicate how much token xj should contribute to token i.
Then, the simplified formula for attention output is:
ai = Σ αijxj
where the sum is calculated only for j ≤ i

Each αi j value is a number, indicating how important the input token xj is when we compute ai. We figure out the weighting α by comparing how similar the current token is to previous ones. The more similar a token is, the more weight it gets. This similarity is calculated by taking the dot product of the two vectors. Finally, we use softmax function to normalize the scores, ensuring the weights add up to 1.

Here's the formula breakdown:
score(xi,xj) = xi ⋅ xj
αij = softmax(score(xi,xj)) ∀j ≤ i

In essence, we are comparing each token to its predecessors, normalizing these scores into probabilities, and then using these probabilities to calculate a weighted sum of the preceding vectors.

Query, Key, and Value: The Attention Head

Now, let's move beyond the simplified intuition and introduce the concept of the attention head, a key component of the Transformer architecture. The attention head allows us to distinguish between three different roles that each input embedding can play:

• Query: The current element being compared to the preceding inputs.
• Key: A preceding input that is being compared to the current element to determine a similarity weight.
• Value: A preceding element that gets weighted and summed up to compute the output for the current element.

To capture these distinct roles, the Transformer uses weight matrices: WQ, WK, and WV. These matrices project each input vector xi into a representation of its role as a key, query, or value:

qi = xiWQ
ki = xiWK
vi = xiWV

To determine the similarity between the current element xi and a prior element xj, we use the dot product between the current element's query vector qi and the preceding element's key vector kj. To maintain numerical stability during training, this dot product is scaled by dividing by the square root of the dimensionality of the query and key vectors (dk). The resulting softmax calculation produces the weights αij. Finally, the output ai is computed as a weighted sum of the value vectors v.

Putting it all together, here are the equations for computing self-attention:

qi = xiWQ; kj = xjWK; vj = xjWV
score(xi,xj) = qi ⋅ kj / √dk
αij = softmax(score(xi,xj)) ∀j ≤ i
ai = Σ αijvj

Essentially, this means that the "attention" part can compare with itself to come up with more accurate values.

Multi-Head Attention: Capturing Diverse Relationships

Transformers don't just use a single attention head; they employ multi-head attention. The idea is that each head can focus on different aspects of the relationships between context elements and the current token. Some heads might specialize in identifying linguistic relationships, while others look for specific patterns in the context.

So, in multi-head attention, we have h separate attention heads working in parallel. Each head has its own set of key, query, and value matrices (WKi, WQi, and WVi), allowing it to project the inputs into separate key, value, and query embeddings.

To generate multiple heads the formula becomes:
qc i = xiWQc; kc j = xjWKc; vc j = xjWVc; ∀c 1 ≤ c ≤ h
scorec(xi,xj) = qc i ⋅ kc j / √dk
αc ij = softmax(scorec(xi,xj)) ∀j ≤ i
headc i = Σ αc ijvc j
ai = (head1 ⊕ head2...⊕ headh)WO

In conclusion MultiHeadAttention(xi,[x1,··· ,xN]) = ai

The output of each head is concatenated and then projected down to the original dimensionality, producing an output that captures a richer understanding of the context.

Transformer Blocks: The Building Blocks of Powerful Models

The self-attention mechanism is at the core of what is called a Transformer block. In addition to the self-attention layer, there are three other types of layers that make it up:

• A feedforward layer
• residual connections
• normalizing layers (colloquially called “layer norm”)

Imagine a stream of water flowing. This is a similar concept to the residual stream, in which a token passes through a transformer block in a representation dimension. The residual stream begins with the input vector, and the various components get read from the stream and have their outputs added to it.

An embedding with dimensionality d is used as input to the stream, which is then conveyed via residual connections and added to by other components of the transformer, the attention layer, and the feedforward layer.

Feedforward Layer

The feedforward layer is a two-layer fully connected network, meaning it has two weight matrices and one hidden layer. Even though the same weights are used for each token position i, each layer is different. The hidden layer dimension is usually bigger than the model d, for instance d = 512 and dff = 2048.

The output is calculated as:
FFN(xi) = ReLU(xiW1 + b1)W2 + b2

Layer Norm

At two stages in the transformer block, the vector is normalized. Normalization, also called layer normalization, facilitates gradient-based training by keeping the values of the hidden layer within a range that enables efficient training.

The term layer norm is confusing because it's used on a single token embedding vector. The first step in normalization is calculating the mean µ, and standard deviation σ.

For an input vector with dimensionality d, these values are:
µ = 1/d Σ di=1 xi
σ = √1/d Σ di=1 (xi − µ)2

The new vectors are normalized by dividing subtracting the mean and dividing by the standard deviation to result in a new vector with a mean of 0 and standard deviation of 1.
x̂ = (x − µ) / σ

Gain and offset values are represented by the two learned parameters γ and β in layer normalization.

LayerNorm(x) = γ(x − µ) / σ + β

Putting it all Together

A transformer block can be expressed by breaking it down by one equation for each computation, and using t to stand for transformer:

t1 i = LayerNorm(xi)
t2 i = MultiHeadAttention(t1 i, [x1 1, ···, x1 N])
t3 i = t2 i + xi
t4 i = LayerNorm(t3 i)
t5 i = FFN(t4 i)
hi = t5 i + t3 i

Multi-head attention looks at all neighboring tokens in the context. But the attention is added to this token's embedding stream.

Transformer blocks are stacked to enhance power and the input/output dimensions are matched. Stacking is a crucial way to build larger language models with between 12 to 96 layers.

Parallelizing Computation with a Single Matrix

Attention and computation are performed on the token to compute ai. It's independent of each token, and can be parallelized through matrix manipulation.

By packing the input embeddings for the tokens into a single matrix X of size [N x d], each row represents the token embedding.

Parallelizing Attention

X is multiplied by value, query, and key matrices to produce matrices with key, query, and value vectors that have a corresponding shape.
Q = XWQ; K = XWK; V = XWV

After this, comparisons are made in a single matrix multiplication. After getting the QKTmatrix, the result is multiplied by V resulting in a representation for each input. The self-attention step has now been reduced to the following computation:
A = softmax(mask(QKT / √dk))V

Masking the Future

During the calculation in QK, each query value results in a score, including those that follow the query, which is inappropriate in language modeling. To resolve this, the elements in the upper-triangular portion are zeroed out, eliminating knowledge of the words that follow in the sequence.

The Input: Token and Position Embeddings

The input to the transformer comes from a sequence of N tokens, where the matrix X has an embedding for each word. The transformer does this by calculating two embeddings: input positional and input token embeddings.

A token embedding, which is a vector dimension d, is our initial representation for an input token. To select token embeddings from the matrix is representing the tokens as one-hot vectors. The dimensions represent each word, where one dimension equals the words index.

These token embeddings are not position-dependent, so the position of each token has to be represented. To represent the position of each token, the token embeddings are combined with the positional embeddings.

The simplest method is absolute position, which starts with initial embeddings corresponding to each input position. To make an input embedding that captures position, the word embedding for each input is combined with its positional embedding.
The matrix X is an [N x d] matrix in which each row is the representation of the i token, calculated by embedding the position to the positional embedding.

The Language Modeling Head: Predicting the Next Word

The final component is the language modeling head, which adds circuitry to the basic transformer when applying to various tasks. The language model provides the circuitry to do language modeling.

Language models are word predictors that can use the context to predict each possible word. From the preceding context, we would be able to compute:
P(fish|Thanks for all the)

The language modeling head takes the last N output from the last transformer layer and uses it to predict the upcoming word at position N+1.

The first module in the language modeling head is the logit vector or score vector, which uses a linear layer, and has a single score for each possible word in the vocabulary V. The linear layer is common and can be learned, although the matrix can be tied to the embedding matrix E. This mapping back to an embedding to a vector is called unembedding

Using a softmax layer turns logits into probabilities.
u = hL NET
y = softmax(u)

These probabilities can be used to help assign text to a given probability, but also to generate text, which we do by sampling from these probabilities.

As a final note, a transformer used for unidirectional language models is called the decoder-only model. The name comes from being used to create the encoder-decoder model for transformers for machine translation.

In Summary

• Transformers are a non-recurrent network based on multi-head attention, and self-attention. The multi-head attention takes an input and maps it to an output, while adding vectors from prior tokens.
• A transformer block has a residual stream where the input is passed to the next layer. This has a multi-head attention layer, followed by a feedforward layer, each preceded by layer normalizations.
• Input is calculated through an embedding that is added to a positional encoding, which represents the token's position.
• A language model can be built by stacking the transformer blocks with the language model head at the top. This generates word probabilities.
• Language models based on transformers have context windows and can predict upcoming words.