Transformer basics in TensorFlow - Practice Problems & Coding Challenges

Challenge - 5 Problems

🧠 Conceptual

intermediate

What is the main purpose of the self-attention mechanism in a Transformer?

In simple terms, what does the self-attention mechanism help a Transformer model do?

A It helps the model focus on different parts of the input sequence to understand context better.

B It generates random noise to improve model robustness during training.

C It reduces the size of the input data by compressing it into a smaller vector.

D It sorts the input tokens in order of importance before processing.

💡 Hint

Think about how the model understands relationships between words in a sentence.

Self-attention allows the model to weigh the importance of each word relative to others, capturing context effectively.

❓ Predict Output

intermediate

Output shape of Transformer encoder layer

Given the following TensorFlow code snippet creating a Transformer encoder layer, what is the shape of the output tensor?

TensorFlow

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

class SimpleTransformerEncoder(tf.keras.layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim):
        super().__init__()
        self.mha = 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)

    def call(self, inputs):
        attn_output = self.mha(inputs, inputs)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        return self.layernorm2(out1 + ffn_output)

sample_input = tf.random.uniform((32, 10, 64))  # batch=32, seq_len=10, embed_dim=64
encoder = SimpleTransformerEncoder(embed_dim=64, num_heads=4, ff_dim=128)
output = encoder(sample_input)
output.shape

A (32, 10, 64)

B (32, 64, 10)

C (10, 32, 64)

D (32, 10, 128)

💡 Hint

Check the input shape and whether the embedding dimension changes after the encoder layer.

The Transformer encoder keeps the sequence length and embedding dimension the same, so output shape matches input shape.

❓ Hyperparameter

advanced

Effect of increasing number of attention heads

What is the most likely effect of increasing the number of attention heads in a Transformer model while keeping the total embedding dimension fixed?

A The embedding dimension per head increases, making each head more powerful.

B The model will have fewer parameters overall, reducing training time significantly.

C Each head attends to a smaller subspace, allowing the model to capture diverse features but with less information per head.

D The model will ignore positional information and treat all tokens equally.

💡 Hint

Think about how the embedding dimension is split among heads.

Increasing heads splits the embedding dimension into smaller parts per head, enabling the model to focus on different aspects of the input.

❓ Metrics

advanced

Choosing the right metric for Transformer language model

You trained a Transformer-based language model to predict the next word in a sentence. Which metric best measures how well your model predicts the sequence?

A Accuracy, because it counts how many words are predicted exactly right.

B Perplexity, because it measures how well the model predicts a sequence of words.

C Mean Squared Error, because it measures the difference between predicted and actual word embeddings.

D F1 Score, because it balances precision and recall for word prediction.

💡 Hint

Consider metrics commonly used in language modeling tasks.

Perplexity is the standard metric for language models, showing how surprised the model is by the actual words.

🔧 Debug

expert

Identifying the cause of shape mismatch error in Transformer

Consider this error during training a Transformer model: ValueError: Shapes (None, 10, 64) and (None, 10, 128) are incompatible. What is the most likely cause?

A The positional encoding was not added to the input embeddings.

B The batch size is inconsistent between input and output tensors.

C The number of attention heads is set higher than the embedding dimension.

D The feed-forward network output dimension does not match the embedding dimension expected by the residual connection.

💡 Hint

Check the dimensions of layers involved in residual connections.

Residual connections require matching dimensions; if the feed-forward layer outputs a different size, this error occurs.

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