Implementation of Handwritten Multi-Head Attention (MHA)

import torch


import torch.nn as nn


import torch.nn.functional as F





class MultiHeadAttention(nn.Module):


    def __init__(self, embed_size, num_heads):


        super(MultiHeadAttention, self).__init__()


        self.embed_size = embed_size


        self.num_heads = num_heads


        self.head_dim = embed_size // num_heads





        assert(


            self.head_dim * num_heads == embed_size


        ), "Embedding size needs to be divisible by num_heads"





        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)


        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)


        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)


        self.fc_out = nn.Linear(num_heads * self.head_dim, embed_size)





    def forward(self, values, keys, query, mask):


        N = query.shape[0]


        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]





        # Divide values, keys, queries into multiple headers


        values = values.reshape(N, value_len, self.num_heads, self.head_dim)


        keys = keys.reshape(N, key_len, self.num_heads, self.head_dim)


        queries = query.reshape(N, query_len, self.num_heads, self.head_dim)


	   # Perform linear transformation


        values = self.values(values)


        keys = self.keys(keys)


        queries = self.queries(queries)





        # Scaled dot-product attention (using Einstein summation)


        #Here directly find the kq matrix


        attention = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])


        if mask is not None:


            attention = attention.masked_fill(mask == 0, float("-1e20"))


		# Use softmax to normalize into attention scores. The denominator is to prevent the attention scores from being too different.


        attention = torch.softmax(attention / (self.embed_size ** (1 / 2)), dim=3)


		# Calculate the weighted sum of the attention score and the corresponding value


        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(


            N, query_len, self.num_heads * self.head_dim


        )





        out = self.fc_out(out)


        return out








embed_size = 256


num_heads = 8


values = torch.randn(64, 10, embed_size)


keys = torch.randn(64, 10, embed_size)


query = torch.randn(64, 10, embed_size)


mask = None # optional mask





multihead_attention = MultiHeadAttention(embed_size, num_heads)


output = multihead_attention(values, keys, query, mask)


print(output.shape) #Expected output: torch.Size([64, 10, 256])

Experimental results:

Einsum notation is an elegant way to perform complex operations on tensors, essentially using a domain-specific language. Once understood and mastered, einsum can help us write more concise and efficient code faster.

Basic content: When two variables have the same index, then traverse and sum. In this case, the sum sign can be omitted.

When implementing some algorithms, the mathematical expressions have already been calculated and need to be converted into code implementation. It is simpler to use einsum, but the readability is very poor.