Attention is all you need

Transformers: Attention is all you need

*RNN seq2seq models

enjoyed

the

movie

transformers

h_1

h_2

h_3

h_4

h_5

h_0

s_0=h_5

Naan

transfarmar

padaththai

rasiththen

Naan

transfarmar

padaththai

s_1

s_2

s_3

Transition to Transformers

Encoder

Decoder

h_1

h_2

h_3

h_4

h_5

Naan

transfarmar

padaththai

rasiththen

Naan

transfarmar

padaththai

s_1

s_2

s_3

Transition to Transformers

Encoder

Decoder

With Attention

\begin{bmatrix} \alpha_{11}\\ \alpha_{12}\\ \alpha_{13}\\ \alpha_{14}\\ \alpha_{15} \end{bmatrix}

enjoyed

the

movie

transformers

Naan

transfarmar

padaththai

rasiththen

Transition to Transformers

Encoder

Decoder

Feed Forward Networks

Self-Attention

Encoder-Decoder Attention

Feed Forward Networks

We will see each of these components (and a few more) one at a time in detail and connect them together to synthesize the final architecture

Word Embedding

h_1

h_2

h_3

h_4

h_5

s_1

s_2

s_3

s_4

s_5

z_1

z_2

z_3

z_4

z_5

enjoyed

the

movie

transformers

Self-Attention

Encoder

Self-Attention

Note that the inputs are vectors (word embeddings) and the outputs are also vectors

1	0	0

0	1	0

1	1	0

1	0	1

1	1	1

You can think of these word embeddings as \(h_i\) in RNN encoder (if you wish to compare)

It is important to establish a strong connection between the word "it" to the word "trophy" or "suitcase" based on the context.

It is called self-attention as it identify the attention within the context. It is different from cross-attention which we will see later.

The goal

Given a word in a sentence, we want to compute the relational score between the word and the rest of the words in the sentence.

	The	animal	didn't	cross	the	street	bacause	it
The
animal
didn't
cross
the
street
because
it

The goal

Given a word in a sentence, we want to compute the relational score between the word and the rest of the words in the sentence....

	The	animal	didn't	cross	the	street	bacause	it
The	0.6	0.1	0.05	0.05	0.02	0.02	0.02	0.1
animal	0.02	0.5	0.06	0.15	0.02	0.05	0.01	0.12
didn't	0.01	0.35	0.45	0.1	0.01	0.02	0.01	0.03
cross	.
the	.
street	.
because	.
it	0.01	0.6	0.02	0.1	0.01	0.2	0.01	0.01

...such that the score is higher if they are related contextually

s_{i}

h_{j}

The goal

	The	animal	didn't	cross	the	street	bacause	it
The	0.6	0.1	0.05	0.05	0.02	0.02	0.02	0.1
animal	0.02	0.5	0.06	0.15	0.02	0.05	0.01	0.12
didn't	0.01	0.35	0.45	0.1	0.01	0.02	0.01	0.03
cross	.
the	.
street	.
because	.
it	0.01	0.6	0.02	0.1	0.01	0.2	0.01	0.01

Now both the vectors \((s_i)\) and \((h_j)\), for all \(i,j\) are available for all the time (whereas in the seq2seq model, \(h_j\) for all \(j\) were available and \(s_i\) was obtained one at a time)

s_{i}

h_{j}

enjoyed

the

movie

transformers

Self-Attention

1	0	0

0	1	0

1	1	0

1	0	1

1	1	1

Let's focus on first calculating the first output from self-attention layer

transformers

Self-Attention

1	0	0

1	1	1

0.3	0.2

\cdots

0.1	0.5

-0.1	0.25

0.11	0.89

0	0.4

0.2	0.7

k_1

v_1

q_1

k_5

v_5

q_5

W_K

W_V

W_Q

W_k

W_V

W_Q

z_1

score(s_{t-1},h_j)

Fixed

variable

score(q_1, \ \ k_j)

Score func: dot product

e_{1j}=[q_1 \cdot k_1, \quad q_1 \cdot k_2, \quad \cdots, \quad \quad \cdots \quad \quad q_1 \cdot k_5]

\alpha_{1j} = softmax(e_{1j})

z_{1} = \sum \limits_{j=1}^5 \alpha_{1j}v_j

What about the \(z_2\)?

Let's focus on first calculating the first output from self-attention layer

transformers

Self-Attention

1	0	0

1	1	1

Let's focus on first calculating the first output from self-attention layer

0.3	0.2

\cdots

0.1	0.5

-0.1	0.25

0.11	0.89

0	0.4

0.2	0.7

k_1

v_1

q_1

k_5

v_5

q_5

W_K

W_V

W_Q

W_k

W_V

W_Q

z_1

score(s_{t-1},h_j)

Fixed

variable

score(q_2, \ \ k_j)

e_{2j}=[q_2 \cdot k_1, \quad q_2 \cdot k_2, \quad \cdots, \quad \quad \cdots \quad \quad q_2 \cdot k_5]

\alpha_{2j} = softmax(e_{2j})

z_{2} = \sum \limits_{j=1}^5 \alpha_{2j}v_j

z_1

z_2

Repeat the procedure for all other \(z\).

Score func: dot product

Wait, can we vectorize all these computations and compute the outputs (\(z_1,z_2,\cdots,z_T\)) in one go?

W_Q

h_1

q_T

h_2

h_T

q_2

q_1

\cdots

\Bigg[

\Bigg]

\Bigg[

\Bigg]

Q \in \mathbb{R}^?

Q \in \mathbb{R}^{64 \times T}

W_K

h_1

k_T

h_2

h_T

k_2

k_1

\cdots

\Bigg[

\Bigg]

\Bigg[

\Bigg]

K \in \mathbb{R}^?

K \in \mathbb{R}^{64 \times T}

Wait, can we vectorize all these computations and compute the outputs (\(z_1,z_2,\cdots,z_T\)) in one go?

W_V

h_1

v_1

h_2

h_T

v_2

v_1

\cdots

\Bigg[

\Bigg]

\Bigg[

\Bigg]

V \in \mathbb{R}^?

V \in \mathbb{R}^{64 \times T}

Wait, can we vectorize all these computations and compute the outputs (\(z_1,z_2,\cdots,z_T\)) in one go?

\begin{aligned} Z &= [z_1,z_2,\cdots,z_T] \\ &=softmax\big(\frac{Q^TK}{\sqrt{d_k}}\big)V^T \end{aligned}

Where \(d_k\) is the dimension of key vector.

Since \(d_k\) scales the values of \(Q^TK\), it is called a scaled-dot product.

dim(Q^TK): T \times 64 \times64 \times T = T \times T

dim(Z): T \times T \times T \times 64 = T \times 64

Vectorized Output

\text{MatMul:} \ Q^TK

\text{Scale}:\frac{1}{\sqrt{d_k}}

\text{Softmax}

\text{MatMul}

\(^*\)Note: The original paper follows \(softmax(\frac{QK^T}{\sqrt{d_k}})V\) as it represents embedding \(h_j\) as a row vector

\text{MatMul:} \ Q^TK

\text{Scale}:\frac{1}{\sqrt{d_k}}

\text{Softmax}

\text{MatMul}

Scaled Dot Product

Head

W_Q

W_K

W_V

H=\{h_1,h_2,\cdots,h_T\}

Two-head Attention

Scaled Dot Product

Head-1

W_Q

W_K

W_V

H=\{h_1,h_2,\cdots,h_T\}

Scaled Dot Product

Head-2

W_Q

W_K

W_V

H=\{h_1,h_2,\cdots,h_T\}

h_j

W_Q^1

W_K^1

W_V^1

W_Q^2

W_K^2

W_V^2

Scaled Dot Product

Attention

Scaled Dot Product

Attention

Two head Attention

Concatenate

Linear

How do we extend this to multi-head attention?

Recall How \(Q,K,V\) are obtained from \(H\).

h_j

\cdots

W_Q^1

W_K^1

W_V^1

W_Q^2

W_K^2

W_V^2

W_Q^8

W_K^8

W_V^8

Scaled Dot Product

Attention

Scaled Dot Product

Attention

Scaled Dot Product

Attention

Concatenate (:\(T \times 512\))

Linear

Multi-Head Attention

Concatenate (:\(T \times 512\))

Linear

W_Q^8

W_K^8

W_V^8

Scaled Dot Product

Attention

W_Q^2

W_K^2

W_V^2

Scaled Dot Product

Attention

Scaled Dot Product

Attention

h_j

W_Q^1

W_K^1

W_V^1

Multi-Head Attention

Concatenate (:\(T \times 512\))

Linear

Multi-Head Attention

W_Q^8

W_K^8

W_V^8

Scaled Dot Product

Attention

W_Q^2

W_K^2

W_V^2

Scaled Dot Product

Attention

Scaled Dot Product

Attention

h_j

W_Q^1

W_K^1

W_V^1

\small \text{MultiHead}(Q,K,V) =\small \text{Concatenate}(head_1,\cdots,head_8)

head_i=Attention

W_O

(Q^i,K^i,V^i)

The input is projected into \(h=8\) different representation subspaces.

So, the multi-head attention allows the model to jointly

attend to information from different representation

subspaces at different positions.

Back to Basic Block

enjoyed

the

movie

transformers

Encoder

Feed Forward Network

Multi-Head Attention

The encoder is composed of a stack of \(N=6\) layers

Each layer in the encoder is composed of two sublayers, namely, multi-head attention and feed forward neural networks

enjoyed

the

movie

transformers

Multi-Head Attention

FFN

z_1=\mathbb{R}^{512}

l={2048}

o_1=\mathbb{R}^{512}

o_1

enjoyed

the

movie

transformers

Multi-Head Attention

FFN

z_2=\mathbb{R}^{512}

l={2048}

o_2=\mathbb{R}^{512}

o_2

enjoyed

the

movie

transformers

Multi-Head Attention

FFN

z_5=\mathbb{R}^{512}

l={2048}

o_5=\mathbb{R}^{512}

o_5

\text{FFN}(z)=max(0,W_1z+b_1)W_2+b_2

Identical network for each position \(z_i\)

Let's calculate the number of learnable parameters in the encoder layer

T=40

d_k,d_q,d_v=64

W_Q :64 \times 512=32,178

W_K :64 \times 512=32,178

W_V :64 \times 512=32,178

8*64 \times 512 = 2,62144

h=8

8*64 \times 512 = 2,62144

8 * 64 \times 512 = 2,62144

\begin{aligned} FFN &=(512*2048)+2048+(2048*512)+512 \\ &=2 \times 10^6 \end{aligned}

Feed Forward Network

Multi-Head Attention

enjoyed

the

movie

transformers

W_O :512 \times 512 =262,144

Therefore, about 3 million parameters per layer

Considering 8 heads

\vdots

Encoder Stack

The encoder is composed of \(N\) identical layers and each layer is composed of 2 sub-layers.

Number of Parameters: \(N \times 3 \times 10^6\)

Number of Parameters: \(6 \times \) 3 million

about 18 million parameters

The computation is parallelized in the horizontal direction (i.e., within a training sample) of the encoder stack, not along the vertical direction.

enjoyed

the

movie

transformers

Layer-1

Layer-2

Layer-6

T \times 512

Let us denote the output sequence of vectors from the encoder as \(e_j\), for \(j=1,\cdots,?\)

Decoder

enjoyed

the

movie

transformers

T \times 512

e_j

Encoder

T \times 512

1 \times 37000

Decoder

Naan

transfarmar

padaththai

rasithen

h_j

e_j

What will be the output dimension?

Decoder Stack

Naan

transfarmar

padaththai

rasithen

\vdots

Layer-1

Layer-2

Layer-6

The decoder is a stack of \(N=6\) Layers. However, each layer is composed of three sublayers.

Feed Forward Network

Masked Multi-Head (Self)Attention

Multi-Head (Cross) Attention

Masked (Self) Attention

\text{MatMul:} \ Q^TK

\text{Scale}:\frac{1}{\sqrt{d_k}}

\text{Softmax}

\text{MatMul}

With one important difference:

Recall that in self-attention we computed the query, key, and value vectors \(q,k\&v\) by multiplying the word embeddings \(h_1,h_2,\cdots,h_T\) with the transformation matrices \(W_Q,W_K,W_V\),respectively.

The same is repeated in the decoder layers. This time the \(h_1,h_2,\cdots,h_T\) are the word embeddings of target sentence. But,

\text{MatMul:} \ Q^TK

\text{Scale}:\frac{1}{\sqrt{d_k}}

\text{Softmax}

\text{MatMul}

Mask

With one important difference: Masking to implement the teacher-forcing approach during training.

Note: Encoder block also uses masking in attention sublayer in practice to mask the padded tokens in sequences having length < T

Of course, we can't use teacher forcing during inference. Instead, the decoder act as a auto-regressor.

Recall that in self-attention we computed the query, key, and value vectors \(q,k\&v\) by multiplying the word embeddings \(h_1,h_2,\cdots,h_T\) with the transformation matrices \(W_Q,W_K,W_V\),respectively.

The same is repeated in the decoder layers. This time the \(h_1,h_2,\cdots,h_T\) are the word embeddings of target sentence. But,

Feed Forward Network

Masked Multi-Head (Self)Attention

Multi-Head (Cross) Attention

First decoder layer

Naan

transfarmar

~~padaththai~~

~~rasiththen~~

<Go>

Feed Forward Network

e_{1:T}

Naan

transfarmar

~~padaththai~~

~~rasiththen~~

<Go>

f\big(W_{Q_1},W_{K_1},W_{V_1},h_{1:T})

Q_2

f(Q_1,K_1,V_1;Mask;W_{O_1})

s_{1:T}

W_{Q_2}

W_{K_2}

W_{V_2}

V_2

K_2

f(Q_2,K_2,V_2;W_{O_2})

\{Q_1,K_1,V_1\}

Number of Parameters:

FFN = 2\times(512 \times2048)+2048+512 \\

=2 \times 10^6 \\

About 2 million parameters from FFN layer

About 1 million parameters from Masked-Multi Head Attention layer

About 1 million parameters from Multi Head Cross Attention layer

About 4 million parameter per decoder layer

Decoder output

Feed Forward Network

e_{1:T}

f\big(W_{Q_1},W_{K_1},W_{V_1},h_{1:T})

Q_2

f(Q_1,K_1,V_1;Mask;W_{O_1})

s_{1:T}

W_{Q_2}

W_{K_2}

W_{V_2}

V_2

K_2

f(Q_2,K_2,V_2;W_{O_2})

\{Q_1,K_1,V_1\}

Linear \(W_D\)

Softmax

The output from the top most decoder layer is linearly transformed by the matrix \(W_D\) of size \(512 \times |V|\) where \(|V|\) is the size of the vocbulary.

The probability distribution for the predicted words is obtained by applying softmax function.

This alone contributes about 19 million parameters of the total 65 million parameters of the architecture.

Positional Encoding

The position of words in a sentence was encoded in the hidden states of RNN based sequence to sequence models.

However, in transformers, no such information is available to either encoder or decoder. Moreover, the output from self-attention is permutation-invariant.

How do we embed positional information in the word embedding \(h_j\) (of size 512)?

h_0^\prime \in \mathbb{R}^{512}

i=0

i=1

i=2

i=511

i=510

i=509

\vdots

\text{position,} \ j \rightarrow \quad

\text{dimension} \ i \rightarrow \quad

Enjoyed

h_1^\prime \in \mathbb{R}^{512}

So, it is necessary to encode the positional information.

Positional Encoding

i=0

i=1

i=2

i=511

i=510

i=509

\vdots

h_0\prime \in \mathbb{R}^{512}

p_0

h_0 \in \mathbb{R}^{512}

How do we fill the elements of the positional vector \(p_0\)?

Could it be a constant vector (i.e, all elements are of constant (position) value \(j\), for \(p_j\))?

Can we use one hot encoding for the position \(j\), \(j=0,1,\cdots,T\)?

or learn embedding for all possible positions?

Not suitable if the sentence length is dynamic.

PE_{(j,i)} = \begin{cases} \sin\left(\frac{j}{10000^{2i/d_{\text{model}}}}\right) & i=0,1,\cdots,255 \\ \cos\left(\frac{j}{10000^{(2i)/d_{\text{model}}}}\right) & i=0,1,\cdots,255\\ \end{cases}

Sinusoidal encoding function

Hypothesis: Embed a unique pattern of features for each position \(j\) and the model will learn to attend by the relative position.

How do we generate the features?

\text{where,}d_{model} = 512

For the fixed position \(j\), the value for \(i\) is sampled from \(sin()\) if \(i\) is even or \(cos()\) if \(i\) is odd

Let's evaluate the function \(PE_{(j,i)}\) for \(j={0,1,\cdots,8}\) and \(i={0,1,\cdots,63}\)

Then, we can visualize this matrix as a heat map.

p_0=\begin{bmatrix} 0\\1\\0\\1\\ 0\\\vdots\\0\\1 \end{bmatrix}

PE_{(j,i)} = \begin{cases} \sin\left(\frac{j=0}{10000^{2i/d_{\text{model}}}}\right) & i=0,1,\cdots,255 \\ \cos\left(\frac{j=0}{10000^{(2i)/d_{\text{model}}}}\right) & i=0,1,\cdots,255\\ \end{cases}

This alternating 0's and 1's will be added to the first word(embedding) of all sentences.

Multi-Head Attention

Feed forward NN

Add&Norm

Multi-Head cross Attention

Feed forward NN

Add&Norm

\sim

h_{source}

Multi-Head Maksed Attention

Add&Norm

\sim

h_{target}

Residual Connections

Adds the original input back to the sub-layer output.

Prevents vanishing gradients: gradients flow directly through the shortcut.

Layer Normalization (Norm)

Normalizes each word vector to have mean = 0, variance = 1

Keeps activation in a stable range across all layers helps much faster convergence

\text{Output} = \text{LayerNorm}(x + \text{SubLayer}(x))

\text{Output} = \text{SubLayer}(x) + x

References

The original "Attention is all you need" paper

Blog: The illustrated transformer

LayerNormalization paper

A book chapter on transformers by Dan Jurafsky (Draft Jan 2023 version)

For NLP-Loyed-Course

By Suman Kundu