Positional Embedding

Self-attention itself doesn’t inherently encode the order of tokens. So, for input sequences like “the dog sat on cat” and “the cat sat on dog” self-attention intrinsically can’t distinguish whether cat comes before on or whether dog comes before on, both of them drastically change the meaning of the sentence . We want something like this:

$attention(m,n)=f(\text{content of }m,\text{content of }n,m-n)$

so that the interaction naturally contains the relative distance between the two tokens.

Absolute PE

Assign a vector to each absolute position in the given sequence: $p_1,p_2,p_3,\dots ,p_n$. So if token embeddings are $e_1,e_2,e_3,\dots ,e_n$, we will perform: $t_i=e_i+p_i$ and $t_i$ is what we pass to the Transformer block.

The representation receives information about position $i$ itself, rather than directly encoding only relative distance $i-j$

Drawbacks:

• Model just remembers the position vectors
• No information about relative positions (though a sufficiently expressive model can learn that relationship)
• Struggle to predict unseen (during training) positions

Sinusoidal PE

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

Then, $t_{pos}=e_{pos}+PE(pos)$

Instead of using a lookup table we use a deterministic function to calculate the PE.

Implementation:

class SinusoidalPositionalEncoding(nn.Module):
    
    def __init__(self,d_model:int, seq_len:int):
        super().__init__()
        self.d_model = d_model
        self.seq_length = seq_len
        
        pos_encode = torch.zeros(seq_len,d_model) # (seq_len,d_model)
        
        position = torch.arange(0,seq_len,dtype=torch.float).unsqueeze(1)  # (seq_len,1)
        div_term = torch.exp(
            torch.arange(0,d_model,2).float()* # float tensor because we will perform exponentials and multplication 
            (-math.log(10_000.0)/d_model) # # we need 1 scaling factor for a pair of sin-cos
            )
        
        pos_encode[:,0::2] = torch.sin(position*div_term) #sin at even terms
        pos_encode[:,1::2] = torch.cos(position*div_term) #cos at odd terms
        
        pos_encode = pos_encode.unsqueeze(0)  # (1,seq_length,d_model) for batch 
        
        self.register_buffer('pos_encode',pos_encode) 
        
    def forward(self, x):
        if x.shape[1] > self.seq_length:
            raise ValueError(
                f"Input sequence length {x.shape[1]} exceeds configured maximum "
                f"{self.seq_length}"
            )
        pos_encode = self.get_buffer("pos_encode")
        x = x + pos_encode[:, :x.shape[1], :]
        return x

Gains:

• Word position
• Similar Performance as learned position embeddings
• May allow the model to extrapolate to longer sequences than the ones encountered during training

Drawbacks:

• Limited generalization to sequence lengths much longer than those seen during training.
• Fixed/pre-defined position; no learnable parameters.

Rotatry PE (RoPE)

Instead of adding positional embedding to the input embedding vectors we introduce positional information by rotating query and key vectors according to their token positions. The idea is:

$q_m^{\prime }=R_m{q}_m,k_n^{\prime }=R_n{k}_n$

where $R_m$ and $R_n$ are position-dependent rotation matrices. We don’t add any positional vectors instead we rotate $Q$ and $K$.

Easiest way to understand is through an 2-dimensional vector example:

$x=\left(\begin{array}{c}{x}_1\\ x_2\end{array}\right)$

We will use a rotation matrix with angle $\varphi$:

$$R(\varphi )=\left(\begin{array}{c}cos(\varphi )-\sin (\varphi )\\ sin(\varphi )+\cos (\varphi )\end{array}\right)$$

Multiplying $x$ with $R(\varphi )$ :

$$R(\varphi )x=\left(\begin{array}{c}{x}_{1}cos(\varphi )-x_2\sin (\varphi )\\ x_{1}sin(\varphi )+x_2\cos (\varphi )\end{array}\right)$$

For a token at position m, rotate it query/key representation by angle proportional to $m$ using a predefined angular frequency $\theta$: $x_m^{\prime }=R(m\theta )x_m$

So,
position $0\to \theta \ast 0=0$
position $1\to \theta \ast 1=\theta$
position $2\to \theta \ast 2=2\theta$
position $3\to \theta \ast 3=3\theta$
$\dots$
position $m\to \theta \ast m=m\theta$

Applying this to $Q,K$ :

$$q_m=W_Q{x}_m^{\prime }\to q_m^{\prime }=W_Q{x}_m{R}^T(m\theta )=q_m{R}_m^T$$ $$k_n=W_K{x}_n^{\prime }\to k_n^{\prime }=W_K{x}_n{R}^T(n\theta )=k_n{R}_n^T$$

Calculating attention scores:

$$q_m^{\prime }{k}_n^{\prime T}=(q_m{R}_m^T)(k_n{R}_n^T{)}^T=q_m{R}_m^T{R}_n{k}_n^T$$

and it can be shown that $R_m^T=R(-m\theta )$, we get:

$$q_m^{\prime }{k}_n^{\prime T}=q_{m}R((n-m)\theta )k_n^T$$

So, attention interaction depends on $(n-m)\to$ relative position.
Although we inject absolute position the relative position emerges in $Q{K}^T$

For real usage, where $d>2$:

Question

Why do we need multiple frequencies ${\theta }_i?$
If we only used a single frequency for all dimensions, then resulting positional representation would be periodic. After the rotation accumulates $2\pi$ the representation returns to the same orientation which will make different positions indistinguishable. To reduce this ambiguity, RoPE applies different rotation frequencies to different dimensions. Some dimensions pairs rotate quickly while others rotate much slowly. The former capture the information from small-positional differences whereas the later capture information over longer distances.

Implementation:

class RotaryPositionalEncoding(nn.Module):
    def __init__(self, head_dim: int, seq_len: int):
        super().__init__()
        self.h_dim = head_dim
        self.seq_len = seq_len
        
        assert self.h_dim % 2 == 0, f"Given d_model {self.h_dim} is not even. We need an even d_model for RoPE!"
        
        div_term = torch.exp(
            torch.arange(0,head_dim,2).float()*(-math.log(10_000.0)/head_dim) # same div term as sinusoidal 
        )
        
        positions = torch.arange(0,seq_len,dtype=torch.float).unsqueeze(1) # (S,1)
        
        angles = positions * div_term  # (S,1) * (1,d_model/2) = (S,d_model/2)
        
        cos = torch.cos(angles)
        sin = torch.sin(angles)
        
        self.register_buffer("cos", cos)
        self.register_buffer("sin", sin)
        
    def forward(self,x):
        
        # x: (B,h,S,d_model)
        assert x.shape[2] <= self.seq_len, f"input of size {x.shape[2]} exceeds seq_length of {self.seq_len}"
        
        curr_seq_len = x.shape[2]
        
        sin = self.get_buffer('sin')[:curr_seq_len]
        cos = self.get_buffer('cos')[:curr_seq_len]
        
        x_even = x[:,:,:,0::2]
        x_odd = x[:,:,:,1::2]
        
        x_even_rot = x_even*cos - x_odd*sin
        x_odd_rot = x_even*sin + x_odd*cos
        
        x_rotated = torch.empty_like(x)
        x_rotated[:,:,:,0::2] = x_even_rot
        x_rotated[:,:,:,1::2] = x_odd_rot
        
        return x_rotated

Gains:

• Flexibility of Sequence length
• Decaying inter-token dependency with increasing relative distance
• Capability of equipping with linear self-attention with relative position encoding
• Gives a geometric interpretation of how positional information influences attention mechanism

Drawbacks:

• Implementation Complexity ;*)

Resources:

• How Rotary Position Embedding Supercharges Modern LLMs [RoPE]