Transformer basics in TensorFlow - Model Pipeline Trace

Model Pipeline - Transformer basics

This pipeline shows how a Transformer model processes text data to learn relationships between words and make predictions, such as translating or understanding sentences.

Data Flow - 7 Stages

1Input Text

1000 sentences x variable length→Raw text sentences collected for training→1000 sentences x variable length

"The cat sat on the mat."

↓

2Tokenization

1000 sentences x variable length→Convert words to integer tokens→1000 sentences x 20 tokens

[12, 45, 78, 9, 34, 56, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

↓

3Embedding Layer

1000 sentences x 20 tokens→Convert tokens to dense vectors→1000 sentences x 20 tokens x 64 features

[[0.12, -0.05, ..., 0.33], ..., [0.01, 0.22, ..., -0.11]]

↓

4Positional Encoding

1000 sentences x 20 tokens x 64 features→Add position info to embeddings→1000 sentences x 20 tokens x 64 features

Embedding vector + position vector for token 5

↓

5Multi-Head Attention

1000 sentences x 20 tokens x 64 features→Model relationships between tokens→1000 sentences x 20 tokens x 64 features

Attention weights highlight 'cat' related to 'sat'

↓

6Feed Forward Network

1000 sentences x 20 tokens x 64 features→Process each token independently→1000 sentences x 20 tokens x 64 features

Non-linear transformation of token features

↓

7Output Layer

1000 sentences x 20 tokens x 64 features→Predict next word probabilities→1000 sentences x 20 tokens x vocabulary size (10000)

[0.01, 0.05, ..., 0.10, ..., 0.02]

Training Trace - Epoch by Epoch

Loss:
2.3 |****
1.8 |******
1.5 |********
1.2 |**********
1.0 |************
     1  2  3  4  5  Epochs

Epoch	Loss ↓	Accuracy ↑	Observation
1	2.30	0.25	Model starts learning basic word patterns
2	1.85	0.40	Attention helps model understand word relations
3	1.50	0.55	Model improves predictions with context
4	1.20	0.65	Loss decreases steadily, accuracy rises
5	1.00	0.72	Model converges well on training data

Prediction Trace - 6 Layers

Layer 1: Tokenization

Layer 2: Embedding Layer

Layer 3: Positional Encoding

Layer 4: Multi-Head Attention

Layer 5: Feed Forward Network

Layer 6: Output Layer (Softmax)

Model Quiz - 3 Questions

What does the positional encoding add to the model?

A Random noise to embeddings

B Extra vocabulary words

C Information about word order

D Loss values for training

Key Insight

Transformers use attention to understand how words relate in a sentence, improving predictions by focusing on important parts of the input. Positional encoding helps the model know word order, which is crucial for meaning.

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