Practical: Custom Transformer Block

Everything, in one block

This is the capstone: a Transformer encoder block that uses every idea in this track at once. It's a subclassed Layer (Lessons 1–3), it threads the training flag through its Dropout sub-layers (Lesson 4), and it's built to drop into a Functional model like any built-in (Lesson 5). The block holds four kinds of sub-layer — multi-head self-attention, a two-layer feed-forward network, two LayerNormalizations, and two Dropouts — and wires them in call() (see the Code block).

The two sub-blocks and their residuals

The shape is the original Transformer recipe: two sub-blocks, each wrapped in a residual connection. The first runs self-attention over the input, then adds the result back to the input (out1 = norm1(inputs + attn_output)) — the inputs + is the skip connection that lets gradients flow past the attention. The second runs the feed-forward network and adds its input back the same way. The two LayerNormalization layers stabilize each sub-block, and the residual adds are why a deep stack of these trains at all. This is post-norm (normalize after the add); many modern variants move the norm before the sub-layer (pre-norm) for steadier deep-stack training — a one-line change once you own the block.

Stack N of these and you have a Transformer encoder. Because you own the source, every research variant — RoPE, sparse attention, a LoRA adapter, grouped-query attention — is a small edit inside a class you control, not a fork of someone else's library.

Code

A Transformer encoder block as a subclassed Layer·python

class TransformerBlock(keras.layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1, **kwargs):
        super().__init__(**kwargs)
        self.att = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim
        )
        self.ffn = keras.Sequential([
            layers.Dense(ff_dim, activation="relu"),
            layers.Dense(embed_dim),
        ])
        self.norm1 = layers.LayerNormalization(epsilon=1e-6)
        self.norm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)

    def call(self, inputs, training=False):
        attn_output = self.att(inputs, inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.norm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.norm2(out1 + ffn_output)

Practical: Custom Transformer Block

Everything, in one block

The two sub-blocks and their residuals

Code

External links

Exercise

Progress

Comments 0