From Seq2Seq to Transformer - Part.2

Transformer Model

The transformer model represents another significant leap in the NLP field after the Seq2Seq model and the attention mechanism. The main idea behind the transformer model is to remove the RNN structure from the Seq2Seq + attention model. Hence, the transformer is a model that exclusively uses attention operations (and FC operations) without relying on RNNs.

Why Remove RNNs?

Sequential data not only depends on the values of each item but also their order. For instance, "John loves Sarah" and "Sarah loves John" both contain the same words, but their order results in different meanings. Thus, a model that can handle sequential data must also be able to process this order information.

RNNs were suitable for handling sequential data since they accept inputs sequentially and update their hidden states. Models like Seq2Seq leveraged RNNs mainly because there weren't many alternatives at the time.

However, RNNs have several disadvantages:

1. Parallelization Issue: Due to its inherent sequential nature, RNNs cannot be parallelized. This makes training on large datasets impractical due to extended training times.

2. Long Distance Dependency Problem: Relationships between distant items in a sequence don't learn well due to the gradient vanishing/exploding problem.

The attention operation provides a solution to both of these issues. It allows for parallelization and directly compares each query with all keys, addressing the long-distance dependency issue. Theoretically, a model using only attention operations should learn faster than an RNN-based model, allowing for training on larger datasets and producing better performance.

How RNNs were Eliminated: Positional Encoding

While RNN structures naturally incorporated order information, attention operations do not consider it. Therefore, in the transformer model, a vector containing positional information, termed 'positional encoding,' of the same dimension as the input embedding is added.

To incorporate order information, the transformer model adds positional encoding to the input embedding.

The positional encoding ensures that even the same word will have different embedding vectors based on its position, successfully conveying positional information.

However, the method mentioned above isn't the only way to generate vectors with positional information.

Structure of the Transformer Model

Like the Seq2Seq model, the transformer model comprises an encoder and a decoder.

The transformer encoder consists of multiple stacked encoder layers. The output of a previous encoder layer serves as the input to the next encoder layer.

Each encoder layer is divided into a multi-head self-attention sub-layer (orange block) and a position-wise fully connected feed-forward sub-layer (blue block). The output of each sub-layer is passed through a residual connection, added to its input, and then normalized (yellow block) before being fed to the next sub-layer.

The transformer decoder consists of multiple stacked decoder layers. The output of a previous decoder layer serves as the input to the next decoder layer.

Each decoder layer comprises a masked multi-head self-attention sub-layer (first orange block), a multi-head cross-attention sub-layer (second orange block), and a position-wise fully connected feed-forward sub-layer (blue block). Similar to the encoder, the output of each sub-layer in the decoder undergoes a residual connection, is added to its input, and then normalized (yellow block) before being fed to the subsequent sub-layer.

Let's delve deeper into each operation.

Multi-head Attention Operation

Multi-head attention operation is a technique designed to execute the attention operation in parallel multiple times. Here, "head" refers to the entity executing the attention operation. Thus, "multi-head" signifies performing attention operations multiple times (on the same sequence). Performing attention operations $h$ times on the same source indicates the ability to focus with $h$ different "focus methods", or in more layman's terms, viewing the sequence from $h$ different "perspectives".

To be specific, when the sequence length is $L$ and the dimension of the input embedding is $d$, a multi-head attention operation with $h$ heads works as follows:

1. Inputs $Q$, $K$, and $V$ are given.

   • In the encoder layer and decoder layer using multi-head self attention,$Q = K = V = X$ (the output from the previous layer).

   • In the decoder layer using multi-head cross attention, $Q = X$ (the output from the previous layer) and $K = V = Y$ (the output from the encoder).

   • $X$, $Y$, and $Z$ are matrices of size$(L, d)$.

2. Creating Query, Key, Value for Each Head:

   • If all heads use the same query, key, and value, it won't provide a multi-perspective multi-head attention operation. Thus, weight matrices $W^Q_i$ , $W^K_i$, and $W^V_i$ are multiplied with $Q$, $K$, and $V$ respectively to generate the query, key, and value for the $i$ $-th$ head: $Q_i$, $K_i$, and $V_i$.

   • These matrices are of dimensions: $(d, d_k)$, $(d, d_k)$, and $(d, d_v)$respectively. The resulting matrices $Q_i$, $K_i$, and $V_i$ are of dimensions$(L, d_k)$, $(L, d_k)$, and $(L, d_v)$ respectively.

   • $d_k$ is the dimension of the query and key for each head, while $d_v$ is the dimension of the value.

   • During actual implementation, rather than multiplying $h$ matrices $W_i^Q$, $W_i^K$, and $W_i^V$ individually, matrices $W^Q$, $W^K$, and $W^V$ of size $(d, h \times d_k)$ are used. After multiplication, the result is split into the desired sizes.

3. Calculate Attention:

   • Each head performs the scaled dot-product attention operation using $Q_i$, $K_i$, and $V_i$.
   $$head_i = Attention(Q_i,K_i,V_i)=softmax(\frac{Q_iK_i^T}{\sqrt{d_k}})V_i$$

   • The resultant attention value is a matrix of size $(L, d_v)$.

   • Typically, during implementation, vector operations are used to calculate all at once rather than separately for each head.

4. Concatenate Each Head's Attention Value:

   • All attention values from each head are concatenated, then multiplied with the weight matrix $W^O$. The result is the final output of the multi-head attention.
   $$MultiHead(Q,K,V)=Concat(head_1,...,head_h)W^O$$
   • $W^O$ is a matrix of size $(d_v, d)=(d,d)$.

It should be noted that the original Transformer paper used $d=512,h=8,d_v=64$.

Thanks to extensive vector operations, the cost of the multi-head attention operation is roughly similar to the cost of a single-head attention operation

Masked Attention Operation

In the decoder, a specialized form of attention operation called the masked multi-head self attention operation is used. This masking ensures that during attention calculations, the model cannot refer to future tokens in the sequence, thus mimicking the real-time generation of sequences.

Several methods can implement masking, but in the Transformer model, a mask matrix is utilized. The mask matrix $M$ is an $(L \times L)$ matrix of $0$s and negative infinity $(-\infty )$, where $L$ is the length of the input sequence. The values in the matrix are defined as:

• $M_{i,j} = -\infty$ if the $ith$ word (query) should not attend to the jth word (key). Specifically, if the ith word (from the query) comes after the jth word (from the key) in the sequence.

• $M_{i,j} = 0$ otherwise.

Given this mask matrix, the procedure for masked attention is:

1. Compute Attention Scores: For the provided query $Q$, key $K$, and value $V$, compute the scaled dot product attention scores $S$:

Where $d$ is the dimension of the query and key vectors.

1. Apply Mask: To the attention scores, add the mask matrix $M$:

   $$S_{masked} = S + M$$

2. Compute Attention Distribution: Calculate the softmax of the masked attention scores $S_{masked}$ to obtain the attention distribution $\alpha$.

   $$\alpha=softmax(S_{masked})$$

3. Calculate Attention Value: Multiply the attention distribution $\alpha$ with the value matrix $V$ to obtain the attention value $V_{attention}$.

Position-wise Fully Connected Feed-Forward Operation

The position-wise fully connected feed-forward operation in the Transformer model can be defined using linear operations followed by a ReLU activation. Specifically:

Where:

• $W_1$ and $W_2$ are weight matrices.

• $b_1$ and $b_2$ are bias vectors.

• $d$ is the dimension of the input embedding.

• $d_{ff}$ is the dimension of the feed-forward operation's hidden layer.

Layer Normalization Operation

In both the Transformer's encoder and decoder layers, the output of each sub-layer undergoes a residual connection followed by layer normalization. Mathematically, if $X$ is the input to a sublayer with output $Y$, then the output $O$ of the layer normalization is:

Layer normalization is a type of normalization that operates across the features (as opposed to batch normalization, which operates across batches) and is commonly used in NLP tasks and the Transformer architecture. It helps in stabilizing the activations in the network and speeds up training.