Common Positional Encodings and Their Implementations

Learnable Position Encoding

Learnable position encoding refers to directly adding position encoding as a learnable parameter to word embedding vectors. This method is simple to operate and easy to understand, with all positional information of words being learned entirely by the model itself.

Since the learnable position encoding matrix is fixed, the model does not have extrapolation capability during inference. When encountering long texts exceeding the training length, it will fail to infer.

Large Model Extrapolation (Length Extrapolation)

Extrapolation refers to the problem where the input length during training and inference is inconsistent, leading to a decline in the model's generalization capability.

For example, if a model is trained using only 512 tokens of text, then during inference, if the input exceeds 512 tokens, the model may not be able to process it correctly. This limits the effectiveness of large models in tasks such as processing long texts or multi-turn dialogues.

Sinusoidal Position Encoding

Trigonometric position encoding requires no training and provides absolute positional information and a certain degree of extrapolation capability.

The sinusoidal position encoding formula is as follows:

$$PE(pos, 2i) = \sin(\frac{pos}{10000^{2i /d_{model}}})\\ PE(pos, 2i + 1) = \cos(\frac{pos}{10000^{2i / d_{model}}})$$

At the same position pos, the sinusoidal values differ across dimensions. $\frac{2i}{d_{model}}$ is used for normalization. In lower dimensions (smaller i), the denominator of the formula is smaller, resulting in higher frequency, thus paying more attention to local positional changes. In higher dimensions (larger i), the frequency is lower, thus paying more attention to global positional changes.

When extrapolating, if longer positions are encountered, the position embedding vector for any position can be derived using the trigonometric transformation formula $sin(a + b) = sin(a) cos(b) + cos(a) sin(b)$.

Limitations of sinusoidal position encoding in extrapolation:

As long as the Attention vector computation result can represent the position difference $m - n$, it has relative position representation capability.

In sinusoidal position encoding, the QK computation between the word $word_i$ at position $i$ and the word $word_{i+k}$ at position $i + k$ is:

$$\begin{align} q^T_{word_1}k_{word_{2}} & = ((embed_{1} + PE_1) W_q )^T \cdot ((embed_2 + PE_2) W_k )\\ & = W_q^T (embed_1 + PE_1)^T (embed_2 + PE_2) W_k\\ & = W_q^T (embed_1^T embed_2 + embed_1^T PE_2 + embed_2^T PE_1 + PE_1^T PE_2) W_k \end{align}$$

Continuing the analysis of $PE_2$ and $PE_1$:

$$\begin{align} PE(pos + k, 2i ) & = sin ((pos + k) / (10000^{2i /dmodel})) \\ & = sin(pos / (10000^{2i /dmodel}) ) cos(k / (10000^{2i /dmodel})) + cos(pos / (10000^{2i /dmodel}) ) sin(k / (10000^{2i /dmodel})) \\ & = PE(pos, 2i) u + PE(pos, 2i + 1) v \end{align}$$

Similarly, $PE(pos + k , 2i + 1) = u PE(pos, 2i + 1) - v PE(pos, 2i)$

Then:

$$PE_1^T PE_2 = 2u = 2 cos(k / 10000^{2i / dmodel}) \text{ (Strictly speaking, this should be a summation)}$$

Therefore, within the formula of $q_{word_1}^T k_{word_2}$, there exists a component $cos(k / 10000^{2i / dmodel})$ that can represent $k$, thus providing a certain degree of relative position representation capability.

As $k$ increases, the cosine function gradually decreases, which means that the farther apart two tokens are, the weaker the correlation between them becomes. This remote attenuation is the reason for the insufficient relative position representation capability.

Rotary Position Encoding (RoPE)

RoPE provides a more natural relative position representation compared to sinusoidal position encoding. Relative position refers to considering the relative distance between the current position and the attended position when computing Attention.

In sinusoidal position encoding, the vector obtained by adding the word vector and position vector is fed into Attention for computation, which can introduce relative position information to a certain extent, though not explicitly.

In RoPE, a uniquely designed multiplication between word vectors and queries is equivalent to a relative position operation, thus explicitly introducing relative position.

Advantages:

• Can be extended to arbitrary sequence lengths
• As the relative distance increases, the dependency between tokens weakens
• Equips linear self-attention with the capability of relative position encoding

We assume that we have found a very elegant and well-behaved function $f(q, i)$ that can provide position vectors for different positions, where $q$ is the word vector and $i$ is the position of the word. Now, when computing Attention between this word and another word, given that the distance between the two words is $m - n$, we expect $Attention(word_m, word_n) = f(q_m, m) * f(q_n, n) = g(q_m, q_n, m - n)$. The question now is where to find functions $f$ and $g$ that satisfy these conditions.

In the complex number space, we know that

$$z \cdot e^{i \theta} = (a + b i) (cos(\theta) + i \cdot sin(\theta)) = (a \cdot cos(\theta) - b \cdot sin(\theta)) + i \cdot(a \cdot sin(\theta) + b \cdot cos(\theta)) = R(\theta) \cdot z \\ R(\theta) = \left[ \begin{matrix} cos(\theta) & - sin(\theta) \\ sin(\theta) & cos(\theta) \end{matrix} \right]$$

Then it can be derived that:

$$(R(\theta_1) z_1)^T \cdot (R(\theta_2) z_2) = z_1^T R(\theta_1)^T R(\theta_2) z_2 = z_1^T R(\theta_1 - \theta_2) z_2$$

We can exactly find the required $f$ and $g$, where $f = R(\theta) z$ and $g = (z_1, z_2, \theta_1 - \theta_2)$.

ROPE is based on the $f$ and $g$ found above. It treats the word vector as a complex vector and performs a rotation operation on each dimension, with the rotation angle $\theta = pos / 10000^{2i /d_{model}}$. Then, when performing Attention computation, relative position is automatically introduced because, according to the above formula, the Attention computation process is equivalent to the $g$ function, which explicitly introduces the relative position operation $m - n$.

ROPE Linear Interpolation

Although ROPE has relatively good relative position representation capability in its formulation, it still suffers from performance degradation during extrapolation. For example, if a model is trained with a context window of 2048 but encounters a context window of 5096 during inference, performance will degrade.

One solution is to perform linear interpolation, where $L$ is the maximum window length used in pre-training, and $L^{'}$ is the current sample length. Then $f^{`}(x, m) = f(x, \frac{m L}{L^{`}})$, as shown in the figure below:

Linear interpolation causes the rotation angle of ROPE to decrease. For example, if the original distance between two tokens is $k$, the difference in rotation angles between the two tokens is $\theta_1$. However, due to linear interpolation, the rotation angle decreases, so the angle difference between two tokens at distance $k$ also becomes smaller, which in turn reduces the correlation of local information. Summary:

• Position interpolation shrinks the rotation arc
• Reduces the rotation speed
• Causes the model to lose high-frequency information, thereby affecting model performance.

ROPE Nonlinear Interpolation Scheme (NTK-Aware scaled ROPE)

The idea behind nonlinear interpolation is to modify the Base value of ROPE, which is that 10000.

So why is nonlinear interpolation better? This is a placeholder for future work; I don't know either.

Alibi (Attention Linear Bias)

Alibi similarly does not directly add position embedding to word vectors. Instead, it modifies the Attention mechanism by adding relative position information.

The Attention computation is $q_{token_i}^T k_{token_j}$. The relative position information is directly added to the computation: $q_{token_i}^T k_{token_j} + (pos_{token_i} - pos_{token_j})$. Simple and clear.

In the multi-head implementation of Attention, different coefficients can also be assigned to different heads to achieve diversified relative position information representation.

Code Implementation

import math
import torch
import torch.nn as nn

class PositionEncoding(nn.Module):
    def __init__(self, embed_dim, max_len=5000, strategy="sinusoidal"):
        """
        Supports multiple position encoding strategies
        :param embed_dim: Embedding dimension
        :param max_len: Maximum sequence length
        :param strategy: Position encoding strategy ("sinusoidal", "learnable", "rope")
        """
        super().__init__()
        self.strategy = strategy
        self.embed_dim = embed_dim

        if strategy == "sinusoidal":
            self.position_encoding = self._create_sinusoidal_encoding(embed_dim, max_len)
        elif strategy == "learnable":
            self.position_encoding = nn.Parameter(torch.zeros(max_len, embed_dim))
            nn.init.normal_(self.position_encoding, mean=0, std=0.02)
        elif strategy == "rope":
            # ROPE does not store the encoding directly but dynamically computes rotations
            self.inv_freq = 1.0 / (10000 ** (torch.arange(0, embed_dim, 2).float() / embed_dim))
        else:
            raise ValueError(f"Unknown position encoding strategy: {strategy}")

    def _create_sinusoidal_encoding(self, embed_dim, max_len):
        position = torch.arange(max_len).unsqueeze(1) # [max_len, 1]
        div_term = torch.exp(torch.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim))
        pe = torch.zeros(max_len, embed_dim)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        return pe

    def forward(self, x):
        """
        Applies position encoding according to the strategy
        :param x: Input tensor, shape [batch_size, seq_len, embed_dim]
        """
        if self.strategy == "sinusoidal":
            return x + self.position_encoding[:x.size(1), :].to(x.device)
        elif self.strategy == "learnable":
            return x + self.position_encoding[:x.size(1), :].to(x.device)
        elif self.strategy == "rope":
            return self._apply_rope(x)
        else:
            raise ValueError(f"Unknown position encoding strategy: {self.strategy}")

    def _apply_rope(self, x):
        """
        Applies ROPE encoding
        """
        batch_size, seq_len, _ = x.size()
        pos_seq = torch.arange(seq_len, device=x.device).unsqueeze(1)
        sin, cos = torch.sin(pos_seq * self.inv_freq), torch.cos(pos_seq * self.inv_freq)
        sin_cos = torch.stack((sin, cos), dim=-1).reshape(seq_len, -1)  # [seq_len, embed_dim]
        x_even, x_odd = x[..., 0::2], x[..., 1::2]
        return torch.cat((x_even * cos - x_odd * sin, x_even * sin + x_odd * cos), dim=-1)