Introduction

Transformers have revolutionized the field of Natural Language Processing (NLP) by enabling models to understand context over long sequences of text. In this guide, we will create a basic implementation of a Transformer model using TensorFlow. By the end, you should have a grasp of the architecture and how to build it in code.

---

Pre-requisites

Before diving into the code, ensure you have the following installed:

• Python 3.x
• TensorFlow 2.x
• NumPy

You can install TensorFlow using pip:

pip install tensorflow

---

Step 1: Import Necessary Libraries

First, we need to import the necessary libraries:

import tensorflow as tf
from tensorflow.keras.layers import Layer, Dense, Embedding, Dropout, LayerNormalization

• tensorflow: The main library for building and training models.
• keras.layers: Contains components to create neural network layers.

---

Step 2: Create the Multi-Head Attention Layer

At the core of a Transformer is the Multi-Head Attention mechanism. It allows the model to focus on different parts of the input sequence simultaneously.

class MultiHeadAttention(Layer):
    def __init__(self, num_heads, key_dim):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.key_dim = key_dim
        self.depth = key_dim // num_heads

        self.wq = Dense(key_dim)

# query layer
        self.wk = Dense(key_dim)

# key layer
        self.wv = Dense(key_dim)

# value layer

        self.dense = Dense(key_dim)

# final linear layer

    def split_heads(self, x, batch_size):
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
        return tf.transpose(x, perm=[0, 2, 1, 3])

# (batch_size, num_heads, seq_len, depth)

    def call(self, v, k, q, mask):
        batch_size = tf.shape(q)[0]

        q = self.wq(q)

# (batch_size, seq_len, key_dim)
        k = self.wk(k)

# (batch_size, seq_len, key_dim)
        v = self.wv(v)

# (batch_size, seq_len, key_dim)

        q = self.split_heads(q, batch_size)

# (batch_size, num_heads, seq_len, depth)
        k = self.split_heads(k, batch_size)

# (batch_size, num_heads, seq_len, depth)
        v = self.split_heads(v, batch_size)

# (batch_size, num_heads, seq_len, depth)

 

# Scaled Dot-Product Attention
        scaled_attention_logits = tf.matmul(q, k, transpose_b=True)

# (batch_size, num_heads, seq_len, seq_len)
        scaled_attention_logits /= tf.math.sqrt(tf.cast(self.depth, tf.float32))

        if mask is not None:
            scaled_attention_logits += (mask * -1e9)

        attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)

# (batch_size, num_heads, seq_len, seq_len)
        output = tf.matmul(attention_weights, v)

# (batch_size, num_heads, seq_len, depth)

        output = tf.transpose(output, perm=[0, 2, 1, 3])

# (batch_size, seq_len, num_heads, depth)
        output = tf.reshape(output, (batch_size, -1, self.key_dim))

# (batch_size, seq_len, key_dim)

        return self.dense(output)

# (batch_size, seq_len, key_dim)

Explanation:

1. Layer Initialization: The __init__ method initializes the number of heads and dimension of keys.
2. Splitting Heads: The split_heads method reshapes and transposes the input so that the model can attend to multiple heads.
3. Call Method: The call method computes the attention weights and returns the output of the multi-head attention using scaled dot-product attention.

---

Step 3: Building the Transformer Block

The Transformer block consists of multi-head attention layers followed by feed-forward layers and normalization.

class TransformerBlock(Layer):
    def __init__(self, num_heads, key_dim, ff_dim, rate=0.1):
        super(TransformerBlock, self).__init__()
        self.attention = MultiHeadAttention(num_heads, key_dim)
        self.ffn = tf.keras.Sequential([
            Dense(ff_dim, activation='relu'),

# Feed Forward Network
            Dense(key_dim)  
        ])
        self.layernorm1 = LayerNormalization(epsilon=1e-6)
        self.layernorm2 = LayerNormalization(epsilon=1e-6)
        self.dropout1 = Dropout(rate)
        self.dropout2 = Dropout(rate)

    def call(self, x, training):
        attn_output = self.attention(x, x, x, None)

# Self-attention
        out1 = self.layernorm1(x + self.dropout1(attn_output, training=training))

# Residual connection
        ffn_output = self.ffn(out1)
        return self.layernorm2(out1 + self.dropout2(ffn_output, training=training))

# Residual connection

Explanation:

1. Self-Attention: The attention layer computes self-attention transposed to the original input.
2. Feed Forward Network: The ffn applies a dense layer followed by ReLU activation and another dense layer.
3. Layer Normalization: The layernorm layers stabilize the learning by normalizing inputs.

---

Step 4: Create the Complete Transformer Model

We will build the full Transformer model, which stacks multiple Transformer blocks.

class Transformer(tf.keras.Model):
    def __init__(self, num_layers, num_heads, key_dim, ff_dim, input_vocab_size, max_position_encoding, rate=0.1):
        super(Transformer, self).__init__()
        self.encoder = tf.keras.Sequential([
            Embedding(input_vocab_size, key_dim),
            tf.keras.layers.Dropout(rate),
            *[TransformerBlock(num_heads, key_dim, ff_dim, rate) for _ in range(num_layers)],
        ])
        self.final_layer = Dense(input_vocab_size)

    def call(self, x, training):
        x = self.encoder(x, training=training)
        return self.final_layer(x)

# Output logits

Explanation:

1. Embedding Layer: Transforms input into dense vectors.
2. Stacking Layers: We stack multiple instances of TransformerBlock.
3. Final Layer: Outputs the final prediction logits.

---

Step 5: Training the Model

Now that we have our Transformer model, let's set it up for training:

def compile_model(model):
    model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])
    return model

Explanation:

The compile_model function configures the model with an optimizer (Adam) and loss function (Sparse Categorical Crossentropy).

---

Final Note:

To train the model, you would typically provide it with input data and call the fit method. However, having the entire pipeline torn out in full detail can be quite lengthy. Thus, following these steps, you can create and modify your Transformer model as needed!

Armed with this knowledge, dive into the exciting world of Transformers and build your own machine learning applications.