#

Multi-Headed Attention (MHA)

This is a tutorial/implementation of multi-headed attention from paper Attention Is All You Need in PyTorch. The implementation is inspired from Annotated Transformer.

Here is the training code that uses a basic transformer with MHA for NLP auto-regression.

Here is an experiment implementation that trains a simple transformer.

25import math
26from typing import Optional, List
27
28import torch
29from torch import nn as nn
30
31from labml import tracker
32from labml_helpers.module import Module

#

Prepare for multi-head attention

This module does a linear transformation and splits the vector into given number of heads for multi-head attention. This is used to transform key, query, and value vectors.

35class PrepareForMultiHeadAttention(Module):

#

46    def __init__(self, d_model: int, heads: int, d_k: int, bias: bool):
47        super().__init__()

#

Linear layer for linear transform

49        self.linear = nn.Linear(d_model, heads * d_k, bias=bias)

#

Number of heads

51        self.heads = heads

#

Number of dimensions in vectors in each head

53        self.d_k = d_k

#

55    def forward(self, x: torch.Tensor):

#

Input has shape [seq_len, batch_size, d_model] or [batch_size, d_model] . We apply the linear transformation to the last dimension and split that into the heads.

59        head_shape = x.shape[:-1]

#

Linear transform

62        x = self.linear(x)

#

Split last dimension into heads

65        x = x.view(*head_shape, self.heads, self.d_k)

#

Output has shape [seq_len, batch_size, heads, d_k] or [batch_size, d_model]

68        return x

#

Multi-Head Attention Module

This computes scaled multi-headed attention for given query , key and value vectors.

$A tt e n t i o n (Q, K, V) = se q so f t ma x (\frac{Q K ^{⊤}}{d _{k}}) V$

In simple terms, it finds keys that matches the query, and gets the values of those keys.

It uses dot-product of query and key as the indicator of how matching they are. Before taking the $so f t ma x$ the dot-products are scaled by $\frac{1}{d _{k}}$ . This is done to avoid large dot-product values causing softmax to give very small gradients when $d_{k}$ is large.

Softmax is calculated along the axis of of the sequence (or time).

71class MultiHeadAttention(Module):

#

heads is the number of heads.
d_model is the number of features in the query , key and value vectors.

92    def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1, bias: bool = True):

#

98        super().__init__()

#

Number of features per head

101        self.d_k = d_model // heads

#

Number of heads

103        self.heads = heads

#

These transform the query , key and value vectors for multi-headed attention.

106        self.query = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=bias)
107        self.key = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=bias)
108        self.value = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=True)

#

Softmax for attention along the time dimension of key

111        self.softmax = nn.Softmax(dim=1)

#

Output layer

114        self.output = nn.Linear(d_model, d_model)

#

Dropout

116        self.dropout = nn.Dropout(dropout_prob)

#

Scaling factor before the softmax

118        self.scale = 1 / math.sqrt(self.d_k)

#

We store attentions so that it can be used for logging, or other computations if needed

121        self.attn = None

#

Calculate scores between queries and keys

This method can be overridden for other variations like relative attention.

123    def get_scores(self, query: torch.Tensor, key: torch.Tensor):

#

Calculate $Q K^{⊤}$ or $S_{ijbh} = \sum_{d} Q_{ibh d} K_{jbh d}$

131        return torch.einsum('ibhd,jbhd->ijbh', query, key)

#

mask has shape [seq_len_q, seq_len_k, batch_size] , where first dimension is the query dimension. If the query dimension is equal to $1$ it will be broadcasted.

133    def prepare_mask(self, mask: torch.Tensor, query_shape: List[int], key_shape: List[int]):

#

139        assert mask.shape[0] == 1 or mask.shape[0] == query_shape[0]
140        assert mask.shape[1] == key_shape[0]
141        assert mask.shape[2] == 1 or mask.shape[2] == query_shape[1]

#

Same mask applied to all heads.

144        mask = mask.unsqueeze(-1)

#

resulting mask has shape [seq_len_q, seq_len_k, batch_size, heads]

147        return mask

#

query , key and value are the tensors that store collection of query, key and value vectors. They have shape [seq_len, batch_size, d_model] .

mask has shape [seq_len, seq_len, batch_size] and mask[i, j, b] indicates whether for batch b , query at position i has access to key-value at position j .

149    def forward(self, *,
150                query: torch.Tensor,
151                key: torch.Tensor,
152                value: torch.Tensor,
153                mask: Optional[torch.Tensor] = None):

#

query , key and value have shape [seq_len, batch_size, d_model]

165        seq_len, batch_size, _ = query.shape
166
167        if mask is not None:
168            mask = self.prepare_mask(mask, query.shape, key.shape)

#

Prepare query , key and value for attention computation. These will then have shape [seq_len, batch_size, heads, d_k] .

172        query = self.query(query)
173        key = self.key(key)
174        value = self.value(value)

#

Compute attention scores $Q K^{⊤}$ . This gives a tensor of shape [seq_len, seq_len, batch_size, heads] .

178        scores = self.get_scores(query, key)

#

Scale scores $\frac{Q K ^{⊤}}{d _{k}}$

181        scores *= self.scale

#

Apply mask

184        if mask is not None:
185            scores = scores.masked_fill(mask == 0, float('-inf'))

#

$so f t ma x$ attention along the key sequence dimension $se q so f t ma x (\frac{Q K ^{⊤}}{d _{k}})$

189        attn = self.softmax(scores)

#

Save attentions if debugging

192        tracker.debug('attn', attn)

#

Apply dropout

195        attn = self.dropout(attn)

#

Multiply by values $se q so f t ma x (\frac{Q K ^{⊤}}{d _{k}}) V$

199        x = torch.einsum("ijbh,jbhd->ibhd", attn, value)

#

Save attentions for any other calculations

202        self.attn = attn.detach()

#

Concatenate multiple heads

205        x = x.reshape(seq_len, batch_size, -1)

#

Output layer

208        return self.output(x)