Attention is All You need

Hey, welcome back! Today, we’re going to dive into one of the most influential papers in the history of AI, Attention Is All You Need. We’ll look at the Transformer architecture introduced by this paper in depth and understand the reasoning behind its design choices.

Transformers are neural networks that take sequential data as input and transform it into another sequence as output. The architecture is divided into two main parts: the encoder and the decoder.

But people had already used encoder-decoder approaches before, so what unique problem did Transformers solve?

When dealing with sequential data using recurrent models, in order to compute the next output, information from previous steps needs to be carried through the sequence. This creates several disadvantages:

• Higher sequential computation cost
• Less parallelization and inefficient use of GPU resources
• Difficulty learning relationships between tokens that are far apart in long sequences

Transformers dropped recurrence and instead compute relationships between tokens primarily through attention. Because tokens within a layer can be processed in parallel, this makes much better use of modern hardware.

Transformers are composed of two main parts:

1. Encoder: Takes our input and outputs a matrix representation of that input.
2. Decoder: Takes the encoded representation and autoregressively generates an output sequence.

Both contain a stack of N layers (N = 6 in the original paper).

Encoder

The encoder is composed of two main subcomponents: Multi-Head Attention (MHA) and a feed-forward network. Each sublayer additionally contains a residual connection followed by layer normalization:

LayerNorm(x + Sublayer(x))

Now, let’s dive deeper and see what is happening inside each component.

Input Embeddings

These are vectors of size 512 in the original Transformer, representing the information encoded in each token.

For the encoder, these embeddings are generated from the input tokens. For the decoder, embeddings are generated from the previously produced output tokens.

Positional Encoding

By nature, Transformers don’t have an inherent understanding of the order of the input like RNNs do. Since self-attention processes tokens without recurrence, we need to encode the positional information of the tokens in some way.

There are many ways to do this, including learned positional embeddings, where we start from randomly initialized embeddings and learn them during training.

The original paper instead uses sinusoidal positional encodings, which are combinations of sine and cosine functions. The authors hypothesized that this could make it easier for the model to learn to attend by relative positions and potentially extrapolate to sequence lengths longer than those encountered during training.

1
import torch
2
import torch.nn as nn
3
import math
4

5
class LearnedPositionalEncoding(nn.Module):
6
    def __init__(self, max_len, d_model):
7
        super().__init__()
8
        self.emb = nn.Embedding(max_len, d_model)
9

10
    def forward(self, x: torch.Tensor):
11
        _, seq_len, _ = x.shape
12

13
        pos = torch.arange(seq_len, device=x.device)
14

15
        pos_emb = self.emb(pos)
16

17
        # [seq_len, d_model] -> [1, seq_len, d_model]
18
        return x + pos_emb.unsqueeze(0)
19

20
class SinusoidalPositionalEncoding(nn.Module):
21
    def __init__(self, max_len, d_model):
22
        super().__init__()
23

24
        pos_enc = torch.zeros(max_len, d_model)
25

26
        # [max_len, 1]
27
        pos = torch.arange(max_len, dtype=torch.float32).unsqueeze(1)
28

29
        # [ceil(d_model / 2)]
30
        div_term = torch.exp(
31
            torch.arange(0, d_model, 2, dtype=torch.float32)
32
            * (-math.log(10000.0) / d_model)
33
        )
34

35
        pos_enc[:, 0::2] = torch.sin(pos * div_term)
36

37
        # Handle odd d_model safely
38
        pos_enc[:, 1::2] = torch.cos(pos * div_term[: pos_enc[:, 1::2].shape[1]])
39

40
        self.register_buffer("pos_enc", pos_enc)
41

42
    def forward(self, x: torch.Tensor):
43
        _, seq_len, _ = x.shape
44

45
        return x + self.pos_enc[:seq_len].unsqueeze(0)

Multi-Head Self-Attention Mechanism

Self-attention takes the positionally encoded representations and projects them into three different vectors: the Query (Q), Key (K), and Value (V) vectors.

We then perform scaled dot-product attention to compute how each token relates to the other tokens.

One advantage of self-attention is that it reduces the path length between tokens. Instead of information having to travel through many sequential steps, tokens can directly attend to one another, making long-range dependencies easier to learn.

There are multiple ways to compute attention, including additive attention and dot-product attention. While they have similar theoretical complexity, dot-product attention is generally faster and more space-efficient in practice because it can take advantage of highly optimized matrix multiplication.

The problem with using dot-product attention without scaling is that when the key dimension dₖ becomes large, the dot products can grow large in magnitude. This pushes the softmax into regions where its gradients become extremely small.

To address this, we scale the dot products by the square root of the key dimension:

We then apply softmax to convert the scaled attention scores into probability-like attention weights. These weights are multiplied by the Value (V) vectors to produce the final output of the attention head.

Instead of performing one attention operation across the entire embedding dimension, we split the representation across multiple attention heads. Each head performs attention in parallel on different learned projections of the input.

1
import math
2
import torch
3
import torch.nn as nn
4
import torch.nn.functional as F
5

6
from causal_masking import causal_masking
7

8
class SingleHeadAttention(nn.Module):
9
    def __init__(self, d_model: int, d_k: int, dropout: float = 0.0):
10
        super().__init__()
11

12
        self.q = nn.Linear(d_model, d_k, bias=False)
13
        self.k = nn.Linear(d_model, d_k, bias=False)
14
        self.v = nn.Linear(d_model, d_k, bias=False)
15

16
        self.dropout = nn.Dropout(dropout)
17

18
    def forward(self, x: torch.Tensor):
19
        _, seq_len, _ = x.shape
20

21
        # (batch, seq_len, d_k)
22
        Q = self.q(x)
23
        K = self.k(x)
24
        V = self.v(x)
25

26
        # 1 / sqrt(d_k)
27
        scale = 1 / math.sqrt(Q.size(-1))
28

29
        # (B, S, d_k) @ (B, d_k, S)
30
        # -> (B, S, S)
31
        attn = torch.matmul(Q, K.transpose(-2, -1)) * scale
32

33
        # causal_masking returns (1, 1, S, S)
34
        mask = causal_masking(seq_len, device=x.device)
35

36
        # Remove the extra head dimension:
37
        # (1, 1, S, S) -> (1, S, S)
38
        mask = mask.squeeze(1)
39

40
        attn = attn.masked_fill(mask, float("-inf"))
41

42
        # Attention probabilities
43
        weight = F.softmax(attn, dim=-1)
44

45
        weight = self.dropout(weight)
46

47
        # (B, S, S) @ (B, S, d_k)
48
        # -> (B, S, d_k)
49
        outp = torch.matmul(weight, V)
50

51
        return outp, weight

Finally, the outputs of all the heads are concatenated and passed through another linear projection.

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as fun
4
from causal_masking import causal_masking
5
import math
6

7
class MultiHeadAttention(nn.Module):
8
    def __init__(self, d_model: int, n_head: int, dropout: float = 0.0):
9
        super().__init__()
10
        if d_model % n_head == 0:
11
            self.n_head = n_head
12
            self.d_head = d_model // n_head
13
            self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
14
            self.proj = nn.Linear(d_model, d_model, bias=False)
15
            self.dropout = nn.Dropout(dropout)
16
        else:
17
            print("d_model and n_head are incompatable")
18

19
    def forward(self, x: torch.Tensor):
20
        batch_s, seq_len, emb_dim = x.shape
21
        qkv = self.qkv(x)
22
        qkv = qkv.view(batch_s, seq_len, 3, self.n_head, self.d_head)
23
        q, k, v = qkv.unbind(dim=2)
24
        q = q.transpose(1, 2)
25
        k = k.transpose(1, 2)
26
        v = v.transpose(1, 2)
27
        scale = 1 / math.sqrt(self.d_head)
28
        attn = torch.matmul(q, k.transpose(-2, -1)) * scale
29
        mask = causal_masking(seq_len, device=x.device)
30
        attn = attn.masked_fill(mask, float("-inf"))
31
        weight = fun.softmax(attn, dim=-1)
32
        weight = self.dropout(weight)
33
        ctx = torch.matmul(weight, v)
34
        out = ctx.transpose(1, 2).contiguous().view(batch_s, seq_len, emb_dim)
35
        out = self.proj(out)
36
        return out, weight

Normalization and Residual Connections

To help stabilize training and improve gradient flow, the original Transformer applies a residual connection followed by layer normalization around each sublayer.

A residual connection takes the input to a sublayer, lets the sublayer transform it, and then adds the original input back:

LayerNorm(x + Sublayer(x))

This allows information and gradients to flow more easily through the deep network.

Feed-Forward Network

Each encoder layer also contains a position-wise feed-forward network composed of two linear transformations with a ReLU activation between them:

FFN(x) = max(0, xW₁ + b₁)W₂ + b₂

The same feed-forward network is applied independently to every token position.

The output of the encoder will be a rich contextual representation of the input sequence, which will then be used by the decoder.

1
import torch.nn as nn
2
import torch
3

4
class FFN(nn.Module):
5
    def __init__(self, d_model, mult: int, dropout: float = 0.0):
6
        super().__init__()
7
        self.network = nn.Sequential(
8
            nn.Linear(d_model, d_model * mult),
9
            nn.GELU(),
10
            nn.Linear(d_model * mult, d_model),
11
            nn.Dropout(dropout),
12
        )
13

14
    def forward(self, x: torch.Tensor):
15
        return self.network(x)

Decoder

The decoder transforms the rich contextual representations produced by the encoder into an output sequence.

It contains similar Multi-Head Attention and Feed-Forward Network components, but it also introduces masked multi-head self-attention.

The decoder additionally contains encoder-decoder attention (cross-attention), where the decoder attends to the output representations produced by the encoder.

Masked Self-Attention Mechanism

In order to preserve the autoregressive nature of generation and make sure we predict the next token using only the tokens that came before it, we mask attention to future positions.

This prevents a token from attending to tokens that have not yet been generated.

1
import torch
2

3
def causal_masking(seq_len: int, device=None):
4
    mask = torch.triu(
5
        torch.ones((seq_len, seq_len), dtype=torch.bool, device=device), diagonal=1
6
    )
7
    return mask.view(1, 1, seq_len, seq_len)

Linear Classifier and Softmax

After the representation passes through the decoder stack, we use a final linear layer followed by softmax to produce a probability distribution over the vocabulary for the next token.

During generation, the next token can then be selected according to a decoding strategy such as greedy decoding, sampling, or top-k sampling.

1
import torch
2
import torch.nn as nn
3
from ffn import FFN
4
from multi_head import MultiHeadAttention
5

6
class TransformerBlock(nn.Module):
7
    def __init__(self, d_model: int, n_head: int, dropout: float = 0.0):
8
        super().__init__()
9
        self.first = nn.LayerNorm(d_model)
10
        self.attn = MultiHeadAttention(d_model, n_head, dropout)
11
        self.second = nn.LayerNorm(d_model)
12
        self.ffn = FFN(d_model, mult=4, dropout=dropout)
13

14
    def forward(self, x: torch.Tensor):
15
        x = x + self.attn(self.first(x))[0]
16
        x = x + self.ffn(self.second(x))
17
        return x

References

Original paper: Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS), 30.

image source: DataCamp How Transformers Work

Code implementation: https://github.com/yeabwang/transformer