How to Implement Multi-Head Attention from Scratch in TensorFlow and Keras

By Stefania Cristina on January 6, 2023 in Attention 27

We have already familiarized ourselves with the theory behind the Transformer model and its attention mechanism. We have already started our journey of implementing a complete model by seeing how to implement the scaled-dot product attention. We shall now progress one step further into our journey by encapsulating the scaled-dot product attention into a multi-head attention mechanism, which is a core component. Our end goal remains to apply the complete model to Natural Language Processing (NLP).

In this tutorial, you will discover how to implement multi-head attention from scratch in TensorFlow and Keras.

After completing this tutorial, you will know:

The layers that form part of the multi-head attention mechanism.
How to implement the multi-head attention mechanism from scratch.

Kick-start your project with my book Building Transformer Models with Attention. It provides self-study tutorials with working code to guide you into building a fully-working transformer model that can
translate sentences from one language to another...

Let’s get started.

How to implement multi-head attention from scratch in TensorFlow and Keras
Photo by Everaldo Coelho, some rights reserved.

Tutorial Overview

This tutorial is divided into three parts; they are:

Recap of the Transformer Architecture
- The Transformer Multi-Head Attention
Implementing Multi-Head Attention From Scratch
Testing Out the Code

Prerequisites

For this tutorial, we assume that you are already familiar with:

Recap of the Transformer Architecture

Recall having seen that the Transformer architecture follows an encoder-decoder structure. The encoder, on the left-hand side, is tasked with mapping an input sequence to a sequence of continuous representations; the decoder, on the right-hand side, receives the output of the encoder together with the decoder output at the previous time step to generate an output sequence.

The encoder-decoder structure of the Transformer architecture
Taken from “Attention Is All You Need“

In generating an output sequence, the Transformer does not rely on recurrence and convolutions.

You have seen that the decoder part of the Transformer shares many similarities in its architecture with the encoder. One of the core mechanisms that both the encoder and decoder share is the multi-head attention mechanism.

The Transformer Multi-Head Attention

Each multi-head attention block is made up of four consecutive levels:

On the first level, three linear (dense) layers that each receive the queries, keys, or values
On the second level, a scaled dot-product attention function. The operations performed on both the first and second levels are repeated h times and performed in parallel, according to the number of heads composing the multi-head attention block.
On the third level, a concatenation operation that joins the outputs of the different heads
On the fourth level, a final linear (dense) layer that produces the output

Multi-head attention
Taken from “Attention Is All You Need“

Recall as well the important components that will serve as building blocks for your implementation of the multi-head attention:

The queries, keys, and values: These are the inputs to each multi-head attention block. In the encoder stage, they each carry the same input sequence after this has been embedded and augmented by positional information. Similarly, on the decoder side, the queries, keys, and values fed into the first attention block represent the same target sequence after this would have also been embedded and augmented by positional information. The second attention block of the decoder receives the encoder output in the form of keys and values, and the normalized output of the first decoder attention block as the queries. The dimensionality of the queries and keys is denoted by $d_k$, whereas the dimensionality of the values is denoted by $d_v$.

The projection matrices: When applied to the queries, keys, and values, these projection matrices generate different subspace representations of each. Each attention head then works on one of these projected versions of the queries, keys, and values. An additional projection matrix is also applied to the output of the multi-head attention block after the outputs of each individual head would have been concatenated together. The projection matrices are learned during training.

Let’s now see how to implement the multi-head attention from scratch in TensorFlow and Keras.

Implementing Multi-Head Attention from Scratch

Let’s start by creating the class, MultiHeadAttention, which inherits from the Layer base class in Keras and initialize several instance attributes that you shall be working with (attribute descriptions may be found in the comments):

class MultiHeadAttention(Layer):
    def __init__(self, h, d_k, d_v, d_model, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.attention = DotProductAttention()  # Scaled dot product attention 
        self.heads = h  # Number of attention heads to use
        self.d_k = d_k  # Dimensionality of the linearly projected queries and keys
        self.d_v = d_v  # Dimensionality of the linearly projected values
        self.W_q = Dense(d_k)  # Learned projection matrix for the queries
        self.W_k = Dense(d_k)  # Learned projection matrix for the keys
        self.W_v = Dense(d_v)  # Learned projection matrix for the values
        self.W_o = Dense(d_model)  # Learned projection matrix for the multi-head output
        ...

class MultiHeadAttention(Layer):

def __init__(self, h, d_k, d_v, d_model, **kwargs):

super(MultiHeadAttention, self).__init__(**kwargs)

self.attention = DotProductAttention() # Scaled dot product attention

self.heads = h # Number of attention heads to use

self.d_k = d_k # Dimensionality of the linearly projected queries and keys

self.d_v = d_v # Dimensionality of the linearly projected values

self.W_q = Dense(d_k) # Learned projection matrix for the queries

self.W_k = Dense(d_k) # Learned projection matrix for the keys

self.W_v = Dense(d_v) # Learned projection matrix for the values

self.W_o = Dense(d_model) # Learned projection matrix for the multi-head output

...

Here note that an instance of the DotProductAttention class that was implemented earlier has been created, and its output was assigned to the variable attention. Recall that you implemented the DotProductAttention class as follows:

from tensorflow import matmul, math, cast, float32
from tensorflow.keras.layers import Layer
from keras.backend import softmax

# Implementing the Scaled-Dot Product Attention
class DotProductAttention(Layer):
    def __init__(self, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)

    def call(self, queries, keys, values, d_k, mask=None):
        # Scoring the queries against the keys after transposing the latter, and scaling
        scores = matmul(queries, keys, transpose_b=True) / math.sqrt(cast(d_k, float32))

        # Apply mask to the attention scores
        if mask is not None:
            scores += -1e9 * mask

        # Computing the weights by a softmax operation
        weights = softmax(scores)

        # Computing the attention by a weighted sum of the value vectors
        return matmul(weights, values)

from tensorflow import matmul, math, cast, float32

from tensorflow.keras.layers import Layer

from keras.backend import softmax

# Implementing the Scaled-Dot Product Attention

class DotProductAttention(Layer):

def __init__(self, **kwargs):

super(DotProductAttention, self).__init__(**kwargs)

def call(self, queries, keys, values, d_k, mask=None):

# Scoring the queries against the keys after transposing the latter, and scaling

scores = matmul(queries, keys, transpose_b=True) / math.sqrt(cast(d_k, float32))

# Apply mask to the attention scores

if mask is not None:

scores += -1e9 * mask

# Computing the weights by a softmax operation

weights = softmax(scores)

# Computing the attention by a weighted sum of the value vectors

return matmul(weights, values)

Next, you will be reshaping the linearly projected queries, keys, and values in such a manner as to allow the attention heads to be computed in parallel.

The queries, keys, and values will be fed as input into the multi-head attention block having a shape of (batch size, sequence length, model dimensionality), where the batch size is a hyperparameter of the training process, the sequence length defines the maximum length of the input/output phrases, and the model dimensionality is the dimensionality of the outputs produced by all sub-layers of the model. They are then passed through the respective dense layer to be linearly projected to a shape of (batch size, sequence length, queries/keys/values dimensionality).

The linearly projected queries, keys, and values will be rearranged into (batch size, number of heads, sequence length, depth), by first reshaping them into (batch size, sequence length, number of heads, depth) and then transposing the second and third dimensions. For this purpose, you will create the class method, reshape_tensor, as follows:

def reshape_tensor(self, x, heads, flag):
    if flag:
        # Tensor shape after reshaping and transposing: (batch_size, heads, seq_length, -1)
        x = reshape(x, shape=(shape(x)[0], shape(x)[1], heads, -1))
        x = transpose(x, perm=(0, 2, 1, 3))
    else:
        # Reverting the reshaping and transposing operations: (batch_size, seq_length, d_model)
        x = transpose(x, perm=(0, 2, 1, 3))
        x = reshape(x, shape=(shape(x)[0], shape(x)[1], -1))
    return x

def reshape_tensor(self, x, heads, flag):

if flag:

# Tensor shape after reshaping and transposing: (batch_size, heads, seq_length, -1)

x = reshape(x, shape=(shape(x)[0], shape(x)[1], heads, -1))

x = transpose(x, perm=(0, 2, 1, 3))

else:

# Reverting the reshaping and transposing operations: (batch_size, seq_length, d_model)

x = transpose(x, perm=(0, 2, 1, 3))

x = reshape(x, shape=(shape(x)[0], shape(x)[1], -1))

return x

The reshape_tensor method receives the linearly projected queries, keys, or values as input (while setting the flag to True) to be rearranged as previously explained. Once the multi-head attention output has been generated, this is also fed into the same function (this time setting the flag to False) to perform a reverse operation, effectively concatenating the results of all heads together.

Hence, the next step is to feed the linearly projected queries, keys, and values into the reshape_tensor method to be rearranged, then feed them into the scaled dot-product attention function. In order to do so, let’s create another class method, call, as follows:

def call(self, queries, keys, values, mask=None):
    # Rearrange the queries to be able to compute all heads in parallel
    q_reshaped = self.reshape_tensor(self.W_q(queries), self.heads, True)
    # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

    # Rearrange the keys to be able to compute all heads in parallel
    k_reshaped = self.reshape_tensor(self.W_k(keys), self.heads, True)
    # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

    # Rearrange the values to be able to compute all heads in parallel
    v_reshaped = self.reshape_tensor(self.W_v(values), self.heads, True)
    # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

    # Compute the multi-head attention output using the reshaped queries, keys and values
    o_reshaped = self.attention(q_reshaped, k_reshaped, v_reshaped, self.d_k, mask)
    # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)
    ...

def call(self, queries, keys, values, mask=None):

# Rearrange the queries to be able to compute all heads in parallel

q_reshaped = self.reshape_tensor(self.W_q(queries), self.heads, True)

# Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

# Rearrange the keys to be able to compute all heads in parallel

k_reshaped = self.reshape_tensor(self.W_k(keys), self.heads, True)

# Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

# Rearrange the values to be able to compute all heads in parallel

v_reshaped = self.reshape_tensor(self.W_v(values), self.heads, True)

# Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

# Compute the multi-head attention output using the reshaped queries, keys and values

o_reshaped = self.attention(q_reshaped, k_reshaped, v_reshaped, self.d_k, mask)

# Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

...

Note that the reshape_tensor method can also receive a mask (whose value defaults to None) as input, in addition to the queries, keys, and values.

Recall that the Transformer model introduces a look-ahead mask to prevent the decoder from attending to succeeding words, such that the prediction for a particular word can only depend on known outputs for the words that come before it. Furthermore, since the word embeddings are zero-padded to a specific sequence length, a padding mask also needs to be introduced to prevent the zero values from being processed along with the input. These look-ahead and padding masks can be passed on to the scaled-dot product attention through the mask argument.

Once you have generated the multi-head attention output from all the attention heads, the final steps are to concatenate back all outputs together into a tensor of shape (batch size, sequence length, values dimensionality) and passing the result through one final dense layer. For this purpose, you will add the next two lines of code to the call method.

...
# Rearrange back the output into concatenated form
output = self.reshape_tensor(o_reshaped, self.heads, False)
# Resulting tensor shape: (batch_size, input_seq_length, d_v)

# Apply one final linear projection to the output to generate the multi-head attention
# Resulting tensor shape: (batch_size, input_seq_length, d_model)
return self.W_o(output)

...

# Rearrange back the output into concatenated form

output = self.reshape_tensor(o_reshaped, self.heads, False)

# Resulting tensor shape: (batch_size, input_seq_length, d_v)

# Apply one final linear projection to the output to generate the multi-head attention

# Resulting tensor shape: (batch_size, input_seq_length, d_model)

return self.W_o(output)

Putting everything together, you have the following implementation of the multi-head attention:

from tensorflow import math, matmul, reshape, shape, transpose, cast, float32
from tensorflow.keras.layers import Dense, Layer
from keras.backend import softmax

# Implementing the Scaled-Dot Product Attention
class DotProductAttention(Layer):
    def __init__(self, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)

    def call(self, queries, keys, values, d_k, mask=None):
        # Scoring the queries against the keys after transposing the latter, and scaling
        scores = matmul(queries, keys, transpose_b=True) / math.sqrt(cast(d_k, float32))

        # Apply mask to the attention scores
        if mask is not None:
            scores += -1e9 * mask

        # Computing the weights by a softmax operation
        weights = softmax(scores)

        # Computing the attention by a weighted sum of the value vectors
        return matmul(weights, values)

# Implementing the Multi-Head Attention
class MultiHeadAttention(Layer):
    def __init__(self, h, d_k, d_v, d_model, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.attention = DotProductAttention()  # Scaled dot product attention
        self.heads = h  # Number of attention heads to use
        self.d_k = d_k  # Dimensionality of the linearly projected queries and keys
        self.d_v = d_v  # Dimensionality of the linearly projected values
        self.d_model = d_model  # Dimensionality of the model
        self.W_q = Dense(d_k)  # Learned projection matrix for the queries
        self.W_k = Dense(d_k)  # Learned projection matrix for the keys
        self.W_v = Dense(d_v)  # Learned projection matrix for the values
        self.W_o = Dense(d_model)  # Learned projection matrix for the multi-head output

    def reshape_tensor(self, x, heads, flag):
        if flag:
            # Tensor shape after reshaping and transposing: (batch_size, heads, seq_length, -1)
            x = reshape(x, shape=(shape(x)[0], shape(x)[1], heads, -1))
            x = transpose(x, perm=(0, 2, 1, 3))
        else:
            # Reverting the reshaping and transposing operations: (batch_size, seq_length, d_k)
            x = transpose(x, perm=(0, 2, 1, 3))
            x = reshape(x, shape=(shape(x)[0], shape(x)[1], self.d_k))
        return x

    def call(self, queries, keys, values, mask=None):
        # Rearrange the queries to be able to compute all heads in parallel
        q_reshaped = self.reshape_tensor(self.W_q(queries), self.heads, True)
        # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

        # Rearrange the keys to be able to compute all heads in parallel
        k_reshaped = self.reshape_tensor(self.W_k(keys), self.heads, True)
        # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

        # Rearrange the values to be able to compute all heads in parallel
        v_reshaped = self.reshape_tensor(self.W_v(values), self.heads, True)
        # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

        # Compute the multi-head attention output using the reshaped queries, keys and values
        o_reshaped = self.attention(q_reshaped, k_reshaped, v_reshaped, self.d_k, mask)
        # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)

        # Rearrange back the output into concatenated form
        output = self.reshape_tensor(o_reshaped, self.heads, False)
        # Resulting tensor shape: (batch_size, input_seq_length, d_v)

        # Apply one final linear projection to the output to generate the multi-head attention
        # Resulting tensor shape: (batch_size, input_seq_length, d_model)
        return self.W_o(output)

from tensorflow import math, matmul, reshape, shape, transpose, cast, float32

from tensorflow.keras.layers import Dense, Layer

from keras.backend import softmax

# Implementing the Scaled-Dot Product Attention

class DotProductAttention(Layer):

def __init__(self, **kwargs):

super(DotProductAttention, self).__init__(**kwargs)

def call(self, queries, keys, values, d_k, mask=None):

# Scoring the queries against the keys after transposing the latter, and scaling

scores = matmul(queries, keys, transpose_b=True) / math.sqrt(cast(d_k, float32))

# Apply mask to the attention scores

if mask is not None:

scores += -1e9 * mask

# Computing the weights by a softmax operation

weights = softmax(scores)

# Computing the attention by a weighted sum of the value vectors

return matmul(weights, values)

# Implementing the Multi-Head Attention

class MultiHeadAttention(Layer):

def __init__(self, h, d_k, d_v, d_model, **kwargs):

super(MultiHeadAttention, self).__init__(**kwargs)

self.attention = DotProductAttention() # Scaled dot product attention

self.heads = h # Number of attention heads to use

self.d_k = d_k # Dimensionality of the linearly projected queries and keys

self.d_v = d_v # Dimensionality of the linearly projected values

self.d_model = d_model # Dimensionality of the model

self.W_q = Dense(d_k) # Learned projection matrix for the queries

self.W_k = Dense(d_k) # Learned projection matrix for the keys

self.W_v = Dense(d_v) # Learned projection matrix for the values

self.W_o = Dense(d_model) # Learned projection matrix for the multi-head output

def reshape_tensor(self, x, heads, flag):

if flag:

# Tensor shape after reshaping and transposing: (batch_size, heads, seq_length, -1)

x = reshape(x, shape=(shape(x)[0], shape(x)[1], heads, -1))

x = transpose(x, perm=(0, 2, 1, 3))

else:

# Reverting the reshaping and transposing operations: (batch_size, seq_length, d_k)

x = transpose(x, perm=(0, 2, 1, 3))

x = reshape(x, shape=(shape(x)[0], shape(x)[1], self.d_k))