Transformer Architecture

Putting everything together to assemble a full transformer architecture

In the previous articles, I described the attention mechanism and multi-head attention in details. Multi-head attention (MHA) is the main building block of Transformer architecture. There is yet, one more component that we need to discuss, and that is the position encoding. In this article, we first start with position embedding, and then we will put everything together to build a full transformer according to the original architecture as proposed in “Attention Is All You Need”.

Position Encoding

First, let’s discuss why we need position encoding, and how it works.

Why we need position encoding?

As we saw in the previous article, the role of self-attention is to learn the attention weights, that is given an input sequence of tokens, self-attention provides a weight matrix that determines the relationship among the tokens, the degree at which one token is related to another token.

However, self-attention is permutation invariant. That means self-attention by itself does not inherently account for the order of tokens in a sequence. It computes the attention weights based on the pairwise similarity of tokens, regardless of their positions in the sequence. In otehr words, pure self-attention treats the input as a set of tokens without any inherent sequence. But in a sequence data, order matters.

So we need some way to make the model consider the order of tokens in the input sequence. Transformer uses position encoding to make attention aware of the order of tokens.

How positional encoding works?

The position information can be encoded using absolute or relative position. Furthermore, the embedding can be fixed or learnable. The original transformer opted for using absolute position with fixed sinusoidal embeddings.

Denoting the position in the sequence as pos and using i to denote the index of the each embedding vector, we have the following position embedding for even and odd indices:

• Even dimensions

\(PE(pos, 2i)=sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)\)

Odd dimensions:

\(PE(pos, 2i+1)=cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)\)

Assuming d=128, we can visualize the resulted embeddings with a heatmap as shown below:

Heatmap displaying the sinusoidal position embedding; (image credit: image by author)

The position embeddings will be added to the token embeddings, and the result form the input to the first layer of the transformer.

Putting Everything Together to Build a Transformer

The transformer is composed of a stack of encoder blocks and a stack of decoder blocks that are interconnected together. The following diagram displays the detailed architecture, with the repeating encoder and decoder blocks.

Full Transformer architecture; (Image credit: redrawn figure, originally adapted from “Attention Is All You Need”)

The first encoder block takes the sum of token-embedding (TE) and position-embedding (PE) as input. The subsequent encoder blocks takes the output of the preceding encoder block as input.

The first decoder layer takes the sum of token-embedding and position-embedding (TE + PE) from the target sequence as input. The next decoder receives the output of the first decoder as input, and this continues to the Nth block.

Besides the output from the previous decoder layer, each decoder also receives the output of the encoder on the same layer. The following figure illustrates the input/output of each encoder/decoder block as well as the cross-connection between encoder and decoder blocks.

Expanded view of transformer architecture; detailed breakdown of the transformer model, showcasing each layer from 1 to N within the encoder and decoder stacks. (image credit: by author)

Transformer Encoder Stack

The encoder is composed of N repeating blocks. Each block contains two residual sub-network:

1. Multi-head self-attention

2. Position-wise feed forward network (FFN)

The following figure shows the architecture of the transformer encoder, and zooming into the MHA sub-layer.

Transformer encoder architecture: From left to right, full encoder architecture, zoomed-in multi-head attention, zoomed-in scaled dot-production attention. (image credit: redrawn based on the source form “Attention Is All You Need”)

We have already covered the multi-head attention (MHA) in details in the previous article. But as a brief recap, for multi-head attention, breaks the query, key and value matrices into h heads and performs scaled dot-product attention on each individual head. The results of the heads are then concatenated:

\(MHA(Q, K, V) = \text{Concat}\left(head^{(1)}, head^{(2)}, ..., head^{(h)}\right)\)

where

\(head^{(i)} = \text{Softmax}\left(\frac{Q^{(i)}{K^{(i)}}^\top}{\sqrt{d}}\right)V^{(i)}\)

The feed-forward network is composed of two linear layer with a ReLU activation in between.

Encoder architecture with its zoomed-in view of Feed-Forward Network (FFN) sub-network; (image credit: redrawn based on image in “Attention Is All You Need”)

The output of both MHA and FFN sub-networks are added to their input in a residual connection. The residual connection can be represented as follows:

\(\begin{array}{c|c} \text{MHA sublayer:} & \text{FFN sublayer:} \\ \hline X + \text{MHA}(X) & X + \text{FFN}(X) \\ \end{array} \)

Lastly, the result is normalized using a Layer Normalization.

Transformer Decoder Stack

The decoder block contains three sub-networks:

1. Multi-head self-attention

2. Multi-head cross attention

3. Position-wise feed-forward network (FFN)

Transformer decoder diagram; (image credit: adapted and redrawn by author, originally based on “Attention Is All You Need“)

The multi-head self-attention and the feed-forward sub-network is similar to the encoder block, so instead we focus on the multi-head cross-attention. For cross-attention, the query vector comes from the previous decoder layer, and the keys and values come from the encoder outputs in the same layer. With this cross-attention, the decoder is also able to attend to all the positions in the input sequence, and learn relationships based on the queries that are derived from the target sequence.

Masking the future

As you saw in the previous section, the decoder receives the target sequence as input, which means the decoder has access to the entire target sequence during training. In order to preserve the auto-regressive propert of the model, we need to ensure that the model does not attend to future information.

Masking the future elements of a sequence; (image credit: image by author)

Setting the upper-trinagular elements in this matrix will ensure that their values is suppressed to zero after applying the Softmax() function.

Training a Transformer Model

After compiling the transformer architecture, the next step is to train the model. When training this transformer, we may face some stability issues. To avoid such stability issues, we have to use an extremely small learning rate (LR) during the early stage of training. But, using a small learning rate, for the entire training is not practical. Therefore, the authors applied a warmup stage to gradually increase the learning-rate from its small initial value. After the warmup stage, we can then apply a learning-rate scheduler to decrease the learning-rate, as we do in a normal training.

The proposed warmup process is shown in the following figure:

Warmup curve, adjusting the learning rate during training, starting with very small learning-rate, and increasing the LR during the warmup stage (image credit: image by author)

Conclusion

In this article, we put the indivudal building blocks such as multi-head attention, feed-forward network, and position embeddings together to form a full transformer architecture with a stack of encoders and a stack of decoders that are cross-linked together. We also looked at how to mask the attention scores during training in order to preserve the auto-regressive property of the model. Lastly, we saw a way to adjust the learniong rate to facilitate the training process.

In future articles, we will dive deeper into the cause of the stability issues raised during training, and will cover other variants of the transformer that resolve this issue without the need for the warmup stage.