Practical: Custom Attention Layer

Why build what Keras already gives you

Keras ships MultiHeadAttention, so writing your own is never about production — it's about understanding. Attention has a reputation for being mystical, but the mechanism is a handful of matrix multiplies. Building it once, by hand, with nothing but keras.ops turns the equation in every Transformer paper into code you can read at a glance.

The three pieces of a custom layer

Subclassing keras.layers.Layer means filling in three methods. __init__ stores config (here, the projection width units). build(input_shape) creates the weights once the input shape is known — three projection matrices for query, key, and value, via add_weight. call(inputs) runs the forward pass. Splitting __init__ from build is what lets Keras infer weight shapes lazily from whatever feeds the layer.

Scaled dot-product attention, line by line

The call body is the famous formula. Project the input into Q, K, V. Score every query against every key with matmul(q, transpose(k)). Divide by sqrt(units) — the scaling that keeps the softmax from saturating as dimension grows. Softmax the scores into attention weights, then use those weights to take a weighted sum of the values. Because every op is keras.ops, the same layer runs unchanged on TensorFlow, PyTorch, and JAX.

Code

SimpleAttention — scaled dot-product attention from scratch·python

class SimpleAttention(keras.layers.Layer):
    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.W_q = self.add_weight(
            shape=(input_shape[-1], self.units), name="query_weight"
        )
        self.W_k = self.add_weight(
            shape=(input_shape[-1], self.units), name="key_weight"
        )
        self.W_v = self.add_weight(
            shape=(input_shape[-1], self.units), name="value_weight"
        )

    def call(self, inputs):
        q = keras.ops.matmul(inputs, self.W_q)
        k = keras.ops.matmul(inputs, self.W_k)
        v = keras.ops.matmul(inputs, self.W_v)

        # Scaled dot-product attention
        scale = keras.ops.sqrt(
            keras.ops.cast(self.units, dtype="float32")
        )
        scores = keras.ops.matmul(q, keras.ops.transpose(k)) / scale
        weights = keras.ops.nn.softmax(scores)
        return keras.ops.matmul(weights, v)

Exercise

Extend SimpleAttention into multi-head self-attention as a Subclassed Layer using only keras.ops: project to Q/K/V, reshape into num_heads heads, run scaled dot-product attention per head, then concatenate and apply a final output projection. Verify it produces the same output (within float tolerance) as keras.layers.MultiHeadAttention when you copy the same weights into both.

Hint

The reshape from (batch, seq, units) to (batch, num_heads, seq, head_dim) — and the transpose back after attention — is where most bugs live. Print shapes at every step and check against MultiHeadAttention's output one head at a time.