Transformer basics in TensorFlow

Introduction

Transformers help computers understand and generate language by focusing on important words in a sentence.

When you want to translate sentences from one language to another.

When you want to summarize a long article into a short paragraph.

When you want to answer questions based on a text.

When you want to generate text like writing stories or emails.

When you want to analyze the meaning of sentences for chatbots.

Syntax

TensorFlow

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

class TransformerBlock(tf.keras.layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super(TransformerBlock, self).__init__()
        self.att = MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = tf.keras.Sequential([
            Dense(ff_dim, activation="relu"),
            Dense(embed_dim),
        ])
        self.layernorm1 = LayerNormalization(epsilon=1e-6)
        self.layernorm2 = LayerNormalization(epsilon=1e-6)
        self.dropout1 = Dropout(rate)
        self.dropout2 = Dropout(rate)

    def call(self, inputs, training):
        attn_output = self.att(inputs, inputs, training=training)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)

This code defines a Transformer block with attention and feed-forward parts.

LayerNormalization helps keep the data stable during training.

Examples

This creates a Transformer block with 32 features, 2 attention heads, and a feed-forward layer of size 64. It processes a batch of 1 sequence with 10 tokens.

TensorFlow

transformer_block = TransformerBlock(embed_dim=32, num_heads=2, ff_dim=64)
output = transformer_block(tf.random.uniform((1, 10, 32)), training=False)

This shows how to wrap the Transformer block inside a simple model for training.

TensorFlow

class SimpleTransformer(tf.keras.Model):
    def __init__(self):
        super(SimpleTransformer, self).__init__()
        self.transformer = TransformerBlock(embed_dim=64, num_heads=4, ff_dim=128)

    def call(self, x):
        return self.transformer(x, training=True)

Sample Model

This program creates a Transformer block and runs it on random data. It prints the shapes to show input and output sizes match.

TensorFlow

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

class TransformerBlock(tf.keras.layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super(TransformerBlock, self).__init__()
        self.att = MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = tf.keras.Sequential([
            Dense(ff_dim, activation="relu"),
            Dense(embed_dim),
        ])
        self.layernorm1 = LayerNormalization(epsilon=1e-6)
        self.layernorm2 = LayerNormalization(epsilon=1e-6)
        self.dropout1 = Dropout(rate)
        self.dropout2 = Dropout(rate)

    def call(self, inputs, training):
        attn_output = self.att(inputs, inputs, training=training)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)

# Create dummy input: batch size 2, sequence length 5, embedding size 16
input_data = tf.random.uniform((2, 5, 16))

# Create Transformer block
transformer_block = TransformerBlock(embed_dim=16, num_heads=2, ff_dim=32)

# Run the block
output = transformer_block(input_data, training=False)

print("Input shape:", input_data.shape)
print("Output shape:", output.shape)

Important Notes

Transformers use attention to focus on important parts of the input.

Embedding size must stay the same before and after the Transformer block.

Multi-head attention splits attention into parts to learn different relationships.

Summary

Transformers help models understand sequences by focusing on important words.

A Transformer block has attention and a small neural network inside.

TensorFlow makes it easy to build Transformer blocks with built-in layers.

Practice

1. What is the main purpose of the attention mechanism in a Transformer model?

easy

A. To reduce the number of layers

B. To increase the size of the model

C. To focus on important parts of the input sequence

D. To generate random outputs

Solution

1. Step 1: Understand attention in Transformers

   The attention mechanism helps the model look at important words or parts in the input sequence.

2. Step 2: Identify the main goal of attention

   Its goal is to give more weight to relevant parts, improving understanding of context.

3. Final Answer:

   To focus on important parts of the input sequence -> Option C

4. Quick Check:

   Attention = Focus on important parts [OK]

Hint: Attention means focusing on key input parts [OK]

Common Mistakes:

• Thinking attention increases model size
• Confusing attention with layer count
• Believing attention creates random outputs

2. Which of the following is the correct way to import the MultiHeadAttention layer in TensorFlow?

easy

A. from tensorflow.keras.layers import MultiHeadAttention

B. import MultiHeadAttention from tensorflow.layers

C. from tensorflow import MultiHeadAttention

D. import MultiHeadAttention from keras.layers

Solution

1. Step 1: Recall TensorFlow Keras layer imports

   TensorFlow's Keras layers are imported from tensorflow.keras.layers module.

2. Step 2: Check correct import syntax

   The correct syntax is: from tensorflow.keras.layers import MultiHeadAttention.

3. Final Answer:

   from tensorflow.keras.layers import MultiHeadAttention -> Option A

4. Quick Check:

   Correct import = from tensorflow.keras.layers import MultiHeadAttention [OK]

Hint: Use tensorflow.keras.layers for Keras layers [OK]

Common Mistakes:

• Using incorrect import paths
• Confusing tensorflow and keras imports
• Using wrong import syntax

3. Given the following TensorFlow code snippet, what is the shape of the output tensor?

import tensorflow as tf
mha = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=4)
query = tf.random.uniform((1, 3, 8))
value = tf.random.uniform((1, 4, 8))
output = mha(query=query, value=value)

medium

A. (1, 3, 4)

B. (1, 4, 8)

C. (3, 1, 8)

D. (1, 3, 8)

Solution

1. Step 1: Understand MultiHeadAttention input shapes

   Query shape is (batch=1, query_len=3, features=8), value shape is (1, 4, 8).

2. Step 2: Determine output shape

   Output shape matches query length and feature dimension: (1, 3, 8).

3. Final Answer:

   (1, 3, 8) -> Option D

4. Quick Check:

   Output shape = query shape except last dim [OK]

Hint: Output shape matches query's first two dims [OK]

Common Mistakes:

• Confusing value length with output length
• Mixing batch and sequence dimensions
• Assuming output feature dims change

4. Identify the error in this Transformer block code snippet:

import tensorflow as tf
class SimpleTransformer(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.mha = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=4)
    def call(self, x):
        attn_output = self.mha(query=x)
        return attn_output + x

medium

A. MultiHeadAttention requires 'value' argument

B. The addition attn_output + x causes shape mismatch

C. The call method should return only attn_output

D. Missing positional encoding before attention

Solution

1. Step 1: Check MultiHeadAttention required arguments

   MultiHeadAttention needs query and value arguments; value is missing here.

2. Step 2: Understand why value is needed

   Without value, the layer cannot compute attention properly, causing an error.

3. Final Answer:

   MultiHeadAttention requires 'value' argument -> Option A

4. Quick Check:

   Missing value argument = MultiHeadAttention requires 'value' argument [OK]

Hint: Always provide query and value arguments to MultiHeadAttention [OK]

Common Mistakes:

• Assuming value defaults when only query is provided
• Ignoring required arguments in call
• Thinking positional encoding is mandatory here

5. You want to build a Transformer block in TensorFlow that includes MultiHeadAttention followed by a feed-forward network with ReLU activation. Which code snippet correctly implements this?

hard

A. class TransformerBlock(tf.keras.layers.Layer):
def __init__(self):
super().__init__()
self.mha = tf.keras.layers.MultiHeadAttention(num_heads=4, key_dim=8)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(8, activation='relu'),
tf.keras.layers.Dense(32)
])
def call(self, x):
attn_output = self.mha(query=x, key=x, value=x)
ffn_output = self.ffn(attn_output)
return ffn_output + x

B. class TransformerBlock(tf.keras.layers.Layer):
def __init__(self):
super().__init__()
self.mha = tf.keras.layers.MultiHeadAttention(num_heads=4, key_dim=8)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(32, activation='relu'),
tf.keras.layers.Dense(8)
])
def call(self, x):
attn_output = self.mha(query=x, key=x, value=x)
ffn_output = self.ffn(attn_output)
return ffn_output + x

C. class TransformerBlock(tf.keras.layers.Layer):
def __init__(self):
super().__init__()
self.mha = tf.keras.layers.MultiHeadAttention(num_heads=8, key_dim=4)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(8, activation='relu'),
tf.keras.layers.Dense(32)
])
def call(self, x):
attn_output = self.mha(query=x, key=x, value=x)
ffn_output = self.ffn(attn_output)
return ffn_output + x

D. class TransformerBlock(tf.keras.layers.Layer):
def __init__(self):
super().__init__()
self.mha = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=8)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(8, activation='relu'),
tf.keras.layers.Dense(32)
])
def call(self, x):
attn_output = self.mha(query=x, key=x, value=x)
ffn_output = self.ffn(attn_output)
return ffn_output + x

Solution

1. Step 1: Check MultiHeadAttention parameters

   num_heads=4, key_dim=8 produces output_dim=32, valid for Transformer blocks.

2. Step 2: Verify feed-forward network structure

   FFN uses Dense(32, activation='relu') to expand, then Dense(8) to match input dim.

3. Step 3: Confirm call method logic

   Applies self-attention, FFN on output, residual connection (shapes match: 8).

4. Final Answer:

   Code with num_heads=4, key_dim=8, FFN Dense(32, ReLU) then Dense(8) -> Option B

5. Quick Check:

   MHA(4 heads, key_dim=8) + FFN expand(32 ReLU)-contract(8) + residual [OK]

Hint: Feed-forward expands then reduces features after attention [OK]

Common Mistakes:

• Swapping Dense layer sizes in feed-forward network
• Using wrong num_heads or key_dim values
• Omitting key argument in MultiHeadAttention