Transformer (intro)

import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers
import numpy as np

# 1) Scaled dot-product attention — the core operation
class MultiHeadAttention(layers.Layer):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_model   = d_model
        assert d_model % num_heads == 0
        self.depth     = d_model // num_heads
        self.wq = layers.Dense(d_model)
        self.wk = layers.Dense(d_model)
        self.wv = layers.Dense(d_model)
        self.dense = layers.Dense(d_model)

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

    def call(self, q, k, v, mask=None):
        batch = tf.shape(q)[0]
        q = self.split_heads(self.wq(q), batch)
        k = self.split_heads(self.wk(k), batch)
        v = self.split_heads(self.wv(v), batch)

        scores = tf.matmul(q, k, transpose_b=True)
        scores = scores / tf.math.sqrt(tf.cast(self.depth, tf.float32))
        if mask is not None:
            scores += (mask * -1e9)
        attn = tf.nn.softmax(scores, axis=-1)
        out = tf.matmul(attn, v)
        out = tf.transpose(out, perm=[0, 2, 1, 3])
        out = tf.reshape(out, (batch, -1, self.d_model))
        return self.dense(out)

# 2) Transformer encoder block — attention + FFN with residual + norm
class EncoderBlock(layers.Layer):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attn  = MultiHeadAttention(d_model, num_heads)
        self.ffn   = keras.Sequential([
            layers.Dense(d_ff, activation='gelu'),
            layers.Dense(d_model),
        ])
        self.ln1 = layers.LayerNormalization(epsilon=1e-6)
        self.ln2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(dropout)
        self.dropout2 = layers.Dropout(dropout)

    def call(self, x, mask, training=False):
        a = self.attn(x, x, x, mask)
        a = self.dropout1(a, training=training)
        x = self.ln1(x + a)                            # residual + norm
        f = self.ffn(x)
        f = self.dropout2(f, training=training)
        return self.ln2(x + f)

# 3) Positional encoding — sin/cos so the model knows position
class PositionalEncoding(layers.Layer):
    def __init__(self, max_len, d_model):
        super().__init__()
        pos = np.arange(max_len)[:, None]
        i   = np.arange(d_model)[None, :]
        angle_rates = 1 / (10000 ** ((2 * (i // 2)) / d_model))
        angles = pos * angle_rates
        pe = np.zeros((max_len, d_model), dtype=np.float32)
        pe[:, 0::2] = np.sin(angles[:, 0::2])
        pe[:, 1::2] = np.cos(angles[:, 1::2])
        self.pe = tf.constant(pe[None, ...])
    def call(self, x):
        return x + self.pe[:, :tf.shape(x)[1], :]

# 4) Full encoder model — input -> embedding -> N blocks -> pool -> dense
VOCAB = 20_000
MAX_LEN = 128
D_MODEL = 128
N_HEADS = 4
N_LAYERS = 4
D_FF = 256
NUM_CLASSES = 4

inputs = keras.Input(shape=(MAX_LEN,), dtype='int32')
mask   = tf.cast(tf.equal(inputs, 0), tf.float32)[:, None, None, :]

x = layers.Embedding(VOCAB, D_MODEL, mask_zero=True)(inputs)
x = PositionalEncoding(MAX_LEN, D_MODEL)(x)
for _ in range(N_LAYERS):
    x = EncoderBlock(D_MODEL, N_HEADS, D_FF)(x, mask)
x = layers.GlobalAveragePooling1D()(x)
outputs = layers.Dense(NUM_CLASSES)(x)

model = keras.Model(inputs, outputs)
model.compile(
    optimizer=keras.optimizers.AdamW(learning_rate=3e-4, weight_decay=1e-5),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['sparse_categorical_accuracy'],
)
model.summary()

# 5) Tokenisation + dataset
tokenizer = keras.layers.TextVectorization(max_tokens=VOCAB, output_sequence_length=MAX_LEN)
tokenizer.adapt(train_texts)

def make_ds(texts, labels, batch=32):
    ds = tf.data.Dataset.from_tensor_slices((texts, labels))
    ds = ds.map(lambda t, l: (tokenizer(t), l)).batch(batch).prefetch(tf.data.AUTOTUNE)
    return ds

model.fit(make_ds(train_texts, train_labels), validation_data=make_ds(val_texts, val_labels), epochs=5)

# 6) Learning rate schedule — warmup then decay
class TransformerSchedule(keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, d_model, warmup_steps=4000):
        super().__init__()
        self.d_model = tf.cast(d_model, tf.float32)
        self.warmup_steps = warmup_steps
    def __call__(self, step):
        step = tf.cast(step, tf.float32)
        return tf.math.rsqrt(self.d_model) * tf.minimum(tf.math.rsqrt(step), step * (self.warmup_steps ** -1.5))

# Use as: keras.optimizers.AdamW(learning_rate=TransformerSchedule(D_MODEL))

# 7) Pre-trained transformers — Hugging Face wins for production
from transformers import TFAutoModelForSequenceClassification, AutoTokenizer
tok = AutoTokenizer.from_pretrained('bert-base-uncased')
mod = TFAutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=4)

inputs = tok(['Hello world', 'How are you?'], padding=True, return_tensors='tf')
logits = mod(**inputs).logits

# Fine-tune
mod.compile(optimizer=keras.optimizers.AdamW(learning_rate=2e-5),
            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True))
mod.fit(x=dict(inputs), y=labels, epochs=3, batch_size=16)

# 8) Attention visualisation — interpretable but tricky
# Inspect attention weights from a captured layer; visualise as a heatmap.
# Useful for debugging on small models; not a definitive interpretability tool.

# 9) Decoder-style (GPT-like) — causal mask
# Add a TRIANGULAR mask in the attention layer:
# mask[i, j] = 0 if j <= i else 1
# Predict next token from prior tokens only. Loss = sparse cross-entropy shifted by one.

# 10) Encoder-decoder (T5-style)
# Encoder block + decoder block (extra cross-attention layer attending to encoder outputs).
# Used for translation, summarisation, question answering.

# 11) Mixed precision + accelerator tips
keras.mixed_precision.set_global_policy('mixed_float16')
# 2-3x speedup on Ampere+ GPUs and TPUs. Don't use on small models — overhead can dominate.

# 12) Long-context tricks
# • Sparse attention (Longformer, BigBird)
# • Flash Attention (TF: tfa or custom kernel)
# • Sliding window + global tokens
# • Memory-efficient retraining via gradient checkpointing

# 13) When to use a transformer
# • You have data + compute (or can fine-tune a pretrained model)
# • Sequences are reasonably long (50-512 tokens for vanilla; longer needs sparse attention)
# • Quality matters more than millisecond latency
# For tiny budgets or short sequences (CRF-style sequence labeling), LSTMs may suffice.

# 14) Common bugs
# • Forgot positional encoding → model treats input as bag-of-tokens; can't learn order
# • Mask shape mismatch → softmax assigns weight to padding positions
# • LayerNorm placed AFTER the residual sum (post-LN) — possible; pre-LN (norm before sublayer) trains more stably
# • Vanishing gradients → wrong LR schedule; use warmup + adamW + lower LR than CNNs
// • Inference slower than expected — disable training-time dropout, use cached attention if generating
// • OOM on long sequences → reduce batch size or use gradient checkpointing
// • Fine-tuning pretrained without LR warmup → catastrophic forgetting
// • Generation looping → temperature too low; or KV cache reset mistakes

from tensorflow.keras import layers
# Multi-head self-attention block
x = layers.MultiHeadAttention(num_heads=8, key_dim=64)(q, k)
x = layers.LayerNormalization()(x + q)
x = layers.Dense(256, activation='gelu')(x)