TensorFlow Transformer Model

The Transformer is the architecture behind GPT, BERT, and virtually every powerful language model released in the past several years. Introduced in the 2017 paper "Attention Is All You Need," the Transformer abandoned recurrent layers entirely and introduced a mechanism called self-attention that allows every word in a sentence to relate directly to every other word simultaneously. Understanding Transformers is essential for working with modern natural language processing in TensorFlow.

Why Recurrent Networks Were Not Enough

Before Transformers, RNNs and LSTMs processed text one word at a time, left to right. This created two significant problems:

  • Long-range dependencies: In the sentence "The bank by the river where I used to fish has now been rebuilt," an RNN struggles to connect "bank" to "river" because many words separate them. The connection weakens as it passes through each time step.
  • Sequential processing: RNNs cannot process words in parallel — each word depends on the previous one. This makes training on large datasets very slow.

Transformers solve both problems with self-attention: every word attends to every other word in one single parallel computation.

Self-Attention: The Core Idea

Self-attention calculates a relevance score between every pair of words in a sentence. For each word, it asks: "Which other words should I pay attention to?" and creates a weighted combination of all words based on those scores.

Sentence: "The cat sat on the mat because it was tired"

For the word "it", self-attention scores:
Word      Score
───────────────────────
The       0.01
cat       0.72  ← HIGH (it = cat)
sat       0.10
on        0.02
the       0.01
mat       0.05
because   0.03
it        0.01
was       0.03
tired     0.02

The model learns that "it" refers to "cat" by assigning
a high attention score to "cat" when processing "it".

Queries, Keys, and Values

Self-attention uses three learned matrices to produce three vectors for each word:

Each word's embedding vector gets multiplied by three different weight matrices:

Word embedding
      │
      ├──[Weight Matrix Q]──► Query (Q) — "What am I looking for?"
      ├──[Weight Matrix K]──► Key   (K) — "What do I contain?"
      └──[Weight Matrix V]──► Value (V) — "What information do I carry?"

Attention Score = softmax(Q × Kᵀ / √d_k) × V

Where d_k = dimension of the key vectors
      √d_k = scaling factor to prevent very large dot products

For each word's query vector, TensorFlow computes the dot product with every other word's key vector. This produces a raw score. Softmax converts these scores into probabilities (they sum to 1). Finally, each word's value vector gets multiplied by its probability and the results are summed, producing a context-aware representation of the original word.

Multi-Head Attention

Instead of running self-attention once, the Transformer runs it multiple times in parallel with different weight matrices. Each "head" learns to attend to different types of relationships.

[Input]
   │
   ├──► [Attention Head 1] — focuses on syntactic relationships
   ├──► [Attention Head 2] — focuses on coreference (what "it" refers to)
   ├──► [Attention Head 3] — focuses on semantic similarity
   └──► [Attention Head 4] — focuses on positional proximity
            │
            ▼
     [Concatenate all heads]
            │
            ▼
     [Linear projection]
            │
            ▼
     [Multi-head attention output]

The Full Transformer Architecture

ENCODER (processes input text)
──────────────────────────────────
Input Tokens
      │
[Token Embedding] + [Positional Encoding]
      │
[Multi-Head Self-Attention]
      │
[Add & LayerNorm]
      │
[Feed-Forward Network]
      │
[Add & LayerNorm]
      │
(Repeat N times)
      │
Encoder Output

DECODER (generates output text)
──────────────────────────────────
Target Tokens (shifted right)
      │
[Token Embedding] + [Positional Encoding]
      │
[Masked Multi-Head Self-Attention]  ← Can only see past tokens
      │
[Add & LayerNorm]
      │
[Cross-Attention] ← Attends to Encoder Output
      │
[Add & LayerNorm]
      │
[Feed-Forward Network]
      │
[Add & LayerNorm]
      │
(Repeat N times)
      │
[Linear + Softmax]
      │
Output Probabilities (next token)

Positional Encoding

Self-attention processes all words in parallel and has no built-in sense of order. "Dog bites man" and "Man bites dog" would look identical without position information. Positional encoding adds a unique pattern to each word's embedding based on its position in the sequence.

import tensorflow as tf
import numpy as np

def positional_encoding(max_length, d_model):
    positions = np.arange(max_length)[:, np.newaxis]
    dims = np.arange(d_model)[np.newaxis, :]

    angles = positions / np.power(10000, (2 * (dims // 2)) / d_model)
    angles[:, 0::2] = np.sin(angles[:, 0::2])  # Sine for even indices
    angles[:, 1::2] = np.cos(angles[:, 1::2])  # Cosine for odd indices

    return tf.cast(angles[np.newaxis, :, :], dtype=tf.float32)

Implementing a Transformer Block in TensorFlow

import tensorflow as tf

class TransformerBlock(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff, dropout_rate=0.1):
        super().__init__()
        self.attention = tf.keras.layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=d_model // num_heads
        )
        self.ffn = tf.keras.Sequential([
            tf.keras.layers.Dense(dff, activation='relu'),
            tf.keras.layers.Dense(d_model)
        ])
        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = tf.keras.layers.Dropout(dropout_rate)
        self.dropout2 = tf.keras.layers.Dropout(dropout_rate)

    def call(self, x, training=False):
        # Multi-head self-attention with residual connection
        attn_output = self.attention(x, x, x)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(x + attn_output)  # Residual + Norm

        # Feed-forward network with residual connection
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layernorm2(out1 + ffn_output)  # Residual + Norm

        return out2

Building a Text Classification Model with Transformers

import tensorflow as tf

vocab_size = 10000
max_len = 200
d_model = 128
num_heads = 4
dff = 256

inputs = tf.keras.Input(shape=(max_len,))
x = tf.keras.layers.Embedding(vocab_size, d_model)(inputs)

# Add positional encoding (simplified)
pos_enc = positional_encoding(max_len, d_model)
x = x + pos_enc

# Stack two Transformer blocks
x = TransformerBlock(d_model, num_heads, dff)(x)
x = TransformerBlock(d_model, num_heads, dff)(x)

# Classify
x = tf.keras.layers.GlobalAveragePooling1D()(x)
x = tf.keras.layers.Dropout(0.3)(x)
outputs = tf.keras.layers.Dense(2, activation='softmax')(x)  # Binary: positive/negative

model = tf.keras.Model(inputs, outputs)
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

Key Transformer Hyperparameters

  • d_model — the dimension of all embeddings and internal representations (typically 128–1024)
  • num_heads — number of attention heads; must divide evenly into d_model
  • dff — the inner dimension of the feed-forward network (typically 4× d_model)
  • num_layers — how many Transformer blocks to stack (6 in the original paper; BERT uses 12 or 24)
  • dropout_rate — probability of dropping connections during training (typically 0.1)

Where Transformers Are Used Today

  • NLP — language translation, summarization, question answering, text generation
  • Vision Transformers (ViT) — image classification by treating image patches as tokens
  • Speech — Whisper (OpenAI's transcription model) is a Transformer
  • Protein structure — AlphaFold2 uses attention mechanisms to predict 3D protein shapes
  • Code generation — GitHub Copilot and similar tools use Transformer-based language models
  • Multi-modal models — models that understand both images and text (CLIP, Flamingo, GPT-4V)

Transformer vs. LSTM Summary

Feature               LSTM              Transformer
─────────────────────────────────────────────────────
Processing order      Sequential         Parallel
Long-range deps.      Weak               Strong
Training speed        Slow               Fast (GPU-friendly)
Memory usage          Low                High
Scales with data      Limited            Excellent
State of the art      Outdated for NLP   Current standard
─────────────────────────────────────────────────────

Mastering the Transformer architecture positions you at the frontier of modern machine learning. The attention mechanism, residual connections, and layer normalization that you implemented here appear in every major model — from BERT and GPT to Vision Transformers and multimodal AI systems. This foundation makes every advanced TensorFlow project more approachable, from fine-tuning pre-trained language models to building your own domain-specific AI from scratch.

Previous lesson
Back to course

Leave a Comment

Your email address will not be published. Required fields are marked *