the values for the position vector are fixed and come from

One approach is to encode the absolute positions

we have done this by adding the position vector \(p_i\) with the word embedding \(x_i\)

Why not relative positions?

Why not multiply? 

Why not learn it? 

-2

-1

0

1

Why not add PE with the attention score?

It is still an active area of research (alongside attention mechanisms)

Absolute 

Relative

Types

Fixed

(non-parametric)

Learnable

(parametric)

Learnable

(parametric)

BERT

GPT-x

(Vanilla Transformer)

We can broadly group existing methods under two categories

Fixed

(non-parametric)

LLaMA

GPT-J 

PALM

BLOOM

T5

OPT

BART

Transformer-XL

0

1

2

7

It is hypothesized that APE does learn to attend to tokens by relative position (Why is this important? We will come back to this later) 

Does it? Let's see

0

1

2

3

4

5

6

7

Suppose we take a pre-trained model (say, GPT-2) and use it for grammatical acceptance test (based on perplexity score).

Ideally, we expect the model prediction to be consistent for the given input sentence with (learned) absolute position encoding

What if we change the starting position of the input by shifting by \(k\) units to the right. (i.e., the embeddings of first token added to \((0+k)\))

100

101

102

103

104

105

107

If the model learned to attend to relative positions, then the prediction should not change. Therefore, the performance should not degrade.

This motivated the need to explicitly encode relative positions

The other limitation of APE is its inability to extrapolate to a longer sequence beyond the context length of a model!

Let's take an example to understand these better

closely related to the above is the "length generalization" ability to extrapolate from shorter training sequences to longer test ones (within the context length of a model)

Suppose a GPT (uses APE) model was pre-trained with a context length \(T\) equals to 512. 

Therefore, enabling the model to extrapolate to sequences longer than \(T\) is important

However, we want the model to generate a story of 1024 tokens.

Well, we do not have the learned embeddings for the positions 513 to 1024. This would require us to re-train the model with \(T=1024\)

However, re-training the same model with increased context length is computationally expensive.

Now, consider the  task of adding two integers using  LLMs with a context length of 128

Say, all training examples demonstrate how to add two 4 digit integers

Can the model generalize the algorithm to add, say, two 12 digit integers?

Yes, the type of PE is a major factor in enabling this ability, but APE is not a good choice! [Paper]

Moreover, what if the sequence itself has no absolute ordering (such as graphs)? 

All these reasons motivate us to use relative position encoding!

However, in the case of Relative Position Encoding (RPE), we have \(T\) relative positions for each token in a sequence of length \(T\) 

Therefore, RPE encodes pairwise relationships between tokens. (APE does not)

Let us understand this better

-1

-2

-3

1

2

3

4

We have \(T\) relative positions for each token in a sequence of length \(T\) and, in general, \(2T-1\) embeddings to encode all relative positions (whereas APE requires only \(T\) embeddings)

Therefore, the relative positions for the word "is" are (-3,-2,-1,0,1,2,3,4)

Let us illustrate by displaying all relative positions for the above sequence

We can view the pairwise relationship of tokens in the sequence as a labelled, directed, fully-connected graph

For example, a directed edge from the word "somewhere" to "something" has a label 1. Whereas, a directed edge from the word "something" to "somewhere" has a label -1.

For \(T=8\)

With a naive implementation, for each token embedding \(x_i\), we need a way to combine (say, addition) \(T\) relative position embeddings. 

In general, it requires \(T^2\) additions instead of \(T\) additions as in APE.

Can we do better?

recall that,

Therefore,

score for the given query and key without any position encoding

Adding a constant that comes from position encoding

 PE just adds a constant to the pre-attention score. Therefore, we can inject position information in the attention layer directly by adjusting the pre-attention score \(e_{ij}\)

Correlation beween a word and position 

Yes, by modifying the pre-attention score \(e_{ij}\) directly

There are multiple ways of modifying the attention score  !

Adding a constant to the attention score is the simplest approach!

\(p_{ij}^K\)  is also called position bias (a bias \(b\) as in \(xW+b\))

Note that in this formulation, the dimension of \(p_{ij}\) is  \(d\) (head dimension), whereas in APE, it is \(d_{model}\) 

No. Because the size of the position embedding matrix still depends on \(T\) !

Suppose the relative distance information is clipped beyond a certain distance \(k\) as follows

Does this approach help the model to generalize to sequence lengths not seen during training?

Then, the size of position embedding matrix to \((2k+1)\), making it independent of \(T\)

It is also empirically observed that clipping the distance does not hurt the performance.

Finally, add the relative position information to value vectors

The position embeddings \( p_{ij}^{\tiny{K}} , p_{ij}^{\tiny{V}} \) are shared across heads  (not across layers)

Replacing the APE in vanilla transformer architecture by RPE with \(k=16\) leads to better performance

Source: [paper]

We can compute the attention score as usual

We can expand it further

The transpose happens along the last dimension of \((P^K)^T \in \mathbb{R}^{T \times d \times T}\). 

The space complexity of \(P^K\) is \(O(T^2d)\)

How can we reduce the space complexity? 

The second term,\(x_iW_Q \in \mathbb{R}^{1 \times d}\) gets multiplied with all relative positions (\(j=0,\cdots,T-1\)) given by the  matrix \((p_{i:}^{K})^T \in \mathbb{R}^{d \times T}\) . 

By reshaping the dimension of \(XW_Q \in \mathbb{R}^{T \times d }\) to \( \mathbb{R}^{T \times 1 \times d }\)

Computing the first term can be parallelized. However, the second term requires a small tweak

resulting dim of product\(XW_Q (P^K)^T \)  in \(T \times 1 \times T\)

Writing \(e_{ij}\)  in vectorized form,

 Note that we construct \(P\) from the embedding matrix which leads to additional memory requirement.

Is there a way to transform the embedding matrix directly instead of computing a separate relative position matrix \(P\)?. 

Yes. You can refer to [here]

Transformer-XL (recursivevly extend context)

Where, 

 \({\color{red}{u,v}}\) are learnable vectors.

\({\color{blue}R_{i-j}^T}\) is the relative distance between the position \(i\) and \(j\) computed from sinusoidal function (non-learnable)

T5-bias

Where, 

 \({\color{blue}r_{i-j}} \in \mathbb{R}\) are learnable scalars, shared across layers. Maximum relative distance is clipped to 128

It also greatly reduces number of learnable parameters when the model size is scaled

RPE injects the position information in the pre-attention score

What about the attention mechanism that does not directly compute or modify \(e_{ij}\) as in the case of kernel attention and low-rank approximation of attention

Can we inject the relative position into the token embeddings (without combining \(T\) PEs per token)?

In kernel and low-rank approximation of attention, the PEs are assumed to be injected into the token embeddings

Encode the relative position in angles and multiply it with the tokens embedding

We can generalize this by letting,

However, we want the inner product between tokens embeddings to encode the relative position information. That is we want it to be

Adding absolute position embedding changes the angle between the representation of words "the" and "book" based on where they appear in a sentence despite their relative distance being the same in both sentences

Can we preserve the angle between word representations so that the relative information can be encoded in angles?

Let us assume the following functional form for the left-hand side of the above equation (keeping rotation in mind)

Recall that the inner product between two complex vectors \(x,y \in \mathbb{C}^d\) is \(x \cdot y^H\), where \(H\) is a Hermitian transpose 

is a non-zero constant (why?). 

Here, we multiply a (complex) constant to the affine-transformed word embedding

then

(replaced \(i,j\) by \(m,n\) to avoid confusion with the imaginary number)

we can write 

Can you recognize this matrix?

Therefore, for \(x_m \in \mathbb{R}^2\)

as a matrix as follows

in

yeah, a rotation matrix \(\mathbf{R}\)

Essentially, it rotates affine-transformed \(x_m\) by \(m\theta\) 

Essentially, the tranformation rotates the embeddings \(x_m,x_n\) by \(m\theta, n\theta\), respectively 

Then the inner product results in

We can generalize this to \(x \in \mathbb{R}^{d}, d>2\)

where \(\mathbf{R}_{\Theta,m }^d\) is given by  

Then, the pre-attention score  is computed by encoding the relative position as follows

It is similar to the one we used in sinusoidal embedding

APE of \(m,n\)

RPE of \(m-n\)

Unlike sinusoidal APE, RoPE injects the positional information at every layer of the model and the position embedding is not injected into the value vectors.

How do they help in length generalization?

Let's define length generalization as the ability of a model to extrapolate to longer sequences than the one  it has seen during training

From the transformer paper

However, no experiments were conducted to validate the same

Which approach (APE or RPE) is good at extrapolation? 

Suppose we train a causal language model with \(T=512\) tokens

During inference, we validate the model by allowing the model to predict next tokens up to \(T_{valid}\) tokens

Perplexity score is used to measure the quality of predictions.

\(T_{valid}\) is varied from 512 to 16000 tokens

We repeat the experiment only by changing the position encoding method and keeping the other aspects as is. 

The figure below shows the performance of using sinusoidal,T5, RoPE and ALiBi  position encoding methods

Does increasing the training context window \(T=1024\)  help?

Yes, it helps T5 Bias RPE more than rotary and sinusoidal APE. For example, for \(T_{valid}=8000\), the model trained with \(T=512\) has a higher perplexity score than the model trained with \(T=1024\)

This suggests that modifying T5-bias might improve the performance. 

 \({\color{blue}r_{j-i}} \in \mathbb{R}\) are learnable scalars

where

Why not hard-code \({\color{blue}r_{j-i}}\)?

as it helps increase the inference speed

How do we do that?

Just add a constant (negative bias) to the attention score based on its position as show in the figure

The attention score is computed as follows

\(m\) is a head specific scalar that follows geometric progression.

For \(n=8\) heads, \(m_i=\frac{1}{2^i},i=1,\cdots,n\)

Since it is RPE, position information is added at every layer of the model.

It is called linear because the bias value increase linearly with respect to the distance

We can  observe that the constant value added for nearby tokens is larger (0,-1,-2..) than the one added for distant tokens say (-128,-129..)

This effectively acts as a local windowed attention that enables the ALiBi PE to predict the next token effectively from the recent past \(k\) tokens (a natural structure in languages)

Why does this work remarkably well even for \(T_{valid}=16K\)?

Subsequent studies generalized ALiBi such as  KERPLE (Kernalized Relative Position embedding for Length Extrapolation) and improved RoPE such as xPOS. 

Validation perplexity

Despite all these studies, length generalization remains a challenge for transformer based models

This raises a question

Is there an inherent limitation in Transformers’ design preventing effective length generalization? 

a recent study (Feb 2024) suggests so