Transformer: mixing tokens, then transforming them

A Transformer processes a sequence as vectors. Its repeating block combines two operations: attention mixes information between token positions, and a feed-forward network transforms the vector at each position. Stacking blocks lets later layers work with increasingly contextual representations.

The feed-forward weights are shared across positions, but each position receives its own input vector. Residual connections add each operation's result back to its input. Normalization controls the scale of those representations.

A block you can inspect

This Keras block uses normalization before each operation. Its causal mask prevents a position from attending to later positions, as required when training a model to predict the next token. The input has shape batch × tokens × width; choose a width divisible by the number of heads.

import tensorflow as tf

class TransformerBlock(tf.keras.layers.Layer):
    def __init__(self, width, heads):
        super().__init__()
        self.norm1 = tf.keras.layers.LayerNormalization()
        self.attention = tf.keras.layers.MultiHeadAttention(
            num_heads=heads, key_dim=width // heads)
        self.norm2 = tf.keras.layers.LayerNormalization()
        self.ffn = tf.keras.Sequential([
            tf.keras.layers.Dense(4 * width, activation="gelu"),
            tf.keras.layers.Dense(width),
        ])

    def call(self, x, training=False):
        normalized = self.norm1(x)
        x = x + self.attention(
            normalized, normalized, use_causal_mask=True,
            training=training)
        return x + self.ffn(self.norm2(x), training=training)

A complete language model also needs token embeddings, position information and an output projection to vocabulary scores. Position information can enter through embeddings or through the attention computation itself. The block alone is the repeating computational core.

Decoder-only models use causal attention to generate continuations. Encoders commonly allow attention in both directions; encoder–decoder models connect the two. Our Small Language Model builds a working decoder from these ingredients and runs the result in your browser.