When processing sequential data such as text, order information is fundamentally critical. To incorporate sequence order rather than treating the input as an unordered set, encoding positional information becomes essential. Positional encoding achieves this by assigning unique embedding vectors to each position and adding them to corresponding token representations. While absolute and relative position encodings remain the most common approaches, this article explores more advanced techniques:
- Rotary Positional Encoding and its variants
- Contextual Positional Encoding (CoPE)
Rotary Positional Encoding Fundamentals
Rotary Positional Encoding (RoPE) represents an innovative technique for handling sequential data in natural language processing. Its primary objective is to incorporate position information through geometric rotation, enabling the model to better understand relationships between sequence elements. The core concept involves encoding relative positions by multiplying contextual representations with rotation matrices. RoPE exhibits natural decay as relative distance increases.
The fundamental principle of RoPE involves applying position-specific rotation matrices to token embeddings, effectively embedding positional information directly into the token representations. This encoding mechanism enhances the model's ability to leverage position information during sequence processing, improving both semantic understanding and generation quality.
Let's implement a basic version of RoPE:
def rotate_embeddings(keys, queries, rotation_complex):
# Concept: For vector v = [x1, y1, x2, y2, ...] with shape = dim
# Convert to complex numbers: [x1, y1, x2, y2, ...] -> x1 + iy1, x2 + iy2
# Complex multiplication equals vector rotation
# (x + iy) * (cos + isin) -> x' + iy'
# Restack: x' + iy' -> [x', y', x2', y2', ...]
# This rotates the vector in pairs
_, sequence_length, _, _ = queries.shape
cos_freq, sin_freq = rotation_complex
cos_freq, sin_freq = cos_freq[:sequence_length], sin_freq[:sequence_length]
# Reshape from (..., n) to (..., n//2, 2)
queries_rot = queries.float().reshape(
queries.shape[:-1] + (-1, 2)
) # (B, T, n_head, C) -> (B, T, n_head, C_pair, 2)
keys_rot = keys.float().reshape(keys.shape[:-1] + (-1, 2))
# Split into real and imaginary components
q_real, q_imag = queries_rot.unbind(-1)
k_real, k_imag = keys_rot.unbind(-1)
cos_freq = broadcast_frequency(cos_freq, q_real)
sin_freq = broadcast_frequency(sin_freq, q_real)
# Apply rotation formula
q_rot_real = q_real * cos_freq - q_imag * sin_freq
q_rot_imag = q_real * sin_freq + q_imag * cos_freq
k_rot_real = k_real * cos_freq - k_imag * sin_freq
k_rot_imag = k_real * sin_freq + k_imag * cos_freq
# Combine real and imaginary parts
queries_out = torch.stack([q_rot_real, q_rot_imag], dim=-1)
keys_out = torch.stack([k_rot_real, k_rot_imag], dim=-1)
# Flatten last two dimensions
queries_out = queries_out.flatten(3)
keys_out = keys_out.flatten(3)
return queries_out.type_as(queries), keys_out.type_as(keys)
This implementation serves as the foundation for the variants we'll explore next. We'll now examine several extensions based on different approaches to rotation matrices, angles, and frequency precomputation. These techniques are particularly valuable for extending context length beyond pre-training limits.
Linear Interpolation for RoPE
Linear RoPE introduces method-dependent functions g(m) and h(θ_d) to modify the standard RoPE equation:
Where s represents the scaling factor (ratio of extended to original context length), and θ_d is defined as follows with base b (typically 10000):
The wavelength (inversely proportional to frequency) represents the token length required for a complete rotation (2π) in dimension d.
Here's the implementation:
def compute_linear_frequencies(embed_dim: int, max_seq_len: int, base_freq: float = 10000.0):
freq_exponents = torch.arange(0, embed_dim, 2)[: (embed_dim // 2)].float() / embed_dim
frequencies = 1.0 / (base_freq ** freq_exponents)
positions = torch.arange(max_seq_len, device=frequencies.device)
angular_freqs = torch.outer(positions, frequencies).float()
cos_components = torch.cos(angular_freqs)
sin_components = torch.sin(angular_freqs)
return cos_components, sin_components
Neural Tangent Kernel Extension
The Neural Tangent Kernel (NTK) provides a framework for analyzing neural network training dynamics in the infinite-width limit. In this regime, network behavior can be described through a fixed kernel function, connecting traditional neural network training with kernel methods.
NTK-aware interpolation addresses the loss of high-frequency information when interpolating RoPE embeddings. Instead of uniformly scaling all RoPE dimensions by a factor s, it applies selective scaling - reducing high-frequency scaling while increasing low-frequency scaling. This is achieved by modifying the θ parameter:
def compute_ntk_frequencies(embed_dim: int, max_seq_len: int, base_freq: float = 10000.0, scale_alpha: int = 16):
adjusted_base = base_freq * scale_alpha ** (embed_dim / (embed_dim - 2))
freq_exponents = torch.arange(0, embed_dim, 2)[: (embed_dim // 2)].float() / embed_dim
frequencies = 1.0 / (adjusted_base ** freq_exponents)
positions = torch.arange(max_seq_len, device=frequencies.device)
angular_freqs = torch.outer(positions, frequencies).float()
cos_components = torch.cos(angular_freqs)
sin_components = torch.sin(angular_freqs)
return cos_components, sin_components
YaRN Implementation
YaRN (Yet another RoPE extensioN) efficiently extends model context windows, requiring 10× fewer tokens and 2.5× fewer training steps than previous methods. It introduces a ramp function integrated into the method-dependent function:
def compute_yarn_frequencies(embed_dim: int, original_ctx_len: int, base_freq: float = 10000.0,
scale_factor: int = 16, beta_fast: int = 32, beta_slow: int = 1,
mscale: float = 0.707, max_ctx_len: int = 2048):
pos_freqs = base_freq ** (torch.arange(0, embed_dim, 2)[: (embed_dim // 2)].float() / embed_dim)
inv_freq_extrap = 1.0 / pos_freqs
inv_freq_interp = 1.0 / (scale_factor * pos_freqs)
low_bound = max(math.floor(embed_dim * math.log(original_ctx_len / (beta_fast * 2 * math.pi)) /
(2 * math.log(base_freq))), 0)
high_bound = min(math.ceil(embed_dim * math.log(original_ctx_len / (beta_slow * 2 * math.pi)) /
(2 * math.log(base_freq))), embed_dim - 1)
linear_weights = (torch.arange(embed_dim // 2, dtype=torch.float32) - low_bound) / (high_bound - low_bound)
ramp_weights = torch.clamp(linear_weights, 0, 1).float().to(device=pos_freqs.device)
mask_weights = 1 - ramp_weights
inv_freqs = inv_freq_interp * (1 - mask_weights) + inv_freq_extrap * mask_weights
final_mscale = float((0.1 * math.log(scale_factor) + 1.0) * mscale)
positions = torch.arange(max_ctx_len, device=inv_freqs.device, dtype=inv_freqs.dtype)
angular_freqs = torch.outer(positions, inv_freqs)
cos_components = angular_freqs.cos() * final_mscale
sin_components = angular_freqs.sin() * final_mscale
return cos_components, sin_components
Contextual Positional Encoding (CoPE)
Contextual Positional Encoding (CoPE) represants a novel approach for sequence processing tasks, particularly valuable in natural language processing and time series analysis. Unlike tradisional static position encodings (such as sinusoidal encodings in Transformers), CoPE dynamically adjusts positional encodings based on actual sequence content, better capturing contextual relationships between elements.
Gating Mechanism
The gating mechanism determines which tokens contribute to position count calculations for each query-key pair. It computes a gate value:
A value of 1 indicates the token is considered in position counting, while 0 means it's ignored.
Position Embedding Computation
To compute position embeddings, we accumulate gate values from the current token and all preceding tokens. Each position represents a token/word/sentence number within a given sequence.
For finite position computation (especially with sparse gate activation, common in sentence processing), fewer positions can cover the entire context of sequence length T. Each position is clamped within the maximum possible position range.
Since sigmoid outputs range [0,1], each position value i results in a float within [0,i]. Consequently, positions are non-learnable and cannot be computed by an embedding layer.
Position Embedding Interpolation
To overcome the limitation of non-learnible floating position values, we assign learnable position embeddings e[p] to each integer position in the sequence. The position embedding for element ij is computed through simple interpolation between the two nearest integer embeddings, weighted by the fractional position value calculated above.
Finally, attention is computed by adding position embeddings to key vectors.
CoPE Implementation
For memory and computational efficiency, the q.e[p] matrix is precomputed to enable interpolation and context addition. The interpolation is computed as:
class ContextualPositionalEncoding(nn.Module):
def __init__(self, max_positions, head_dimension):
super().__init__()
self.max_positions = max_positions
self.position_embeddings = nn.Parameter(torch.zeros(1, head_dimension, max_positions))
def forward(self, query_vectors, attention_scores):
# Compute gated positions
gate_values = torch.sigmoid(attention_scores)
position_counts = gate_values.flip(-1).cumsum(dim=-1).flip(-1)
position_counts = position_counts.clamp(max=self.max_positions - 1)
# Interpolate between integer positions
pos_ceil = position_counts.ceil().long()
pos_floor = position_counts.floor().long()
# Compute integer position logits
integer_logits = torch.matmul(query_vectors, self.position_embeddings)
ceil_logits = integer_logits.gather(-1, pos_ceil)
floor_logits = integer_logits.gather(-1, pos_floor)
# Linear interpolation weights
interpolation_weights = position_counts - pos_floor
return ceil_logits * interpolation_weights + floor_logits * (1 - interpolation_weights)
Given query matrices and query-key products, the CoPE forward pass returns interpolated position embeddings. These are then integrated into the attention context:
class AttentionModule(nn.Module):
def __init__(self, config):
super().__init__()
self.embed_dim = config.d_model
self.num_heads = config.num_heads
self.head_dim = config.d_model // config.num_heads
self.num_kv_heads = config.num_heads if config.num_kv_heads == 0 else config.num_kv_heads
assert self.num_heads % self.num_kv_heads == 0
self.queries_per_kv = self.num_heads // self.num_kv_heads
self.cope = ContextualPositionalEncoding(config.seq_len, self.head_dim)
self.key_proj = nn.Linear(self.embed_dim, self.head_dim * self.num_kv_heads)
self.query_proj = nn.Linear(self.embed_dim, self.head_dim * self.num_heads)
self.value_proj = nn.Linear(self.embed_dim, self.head_dim * self.num_kv_heads)
self.output_proj = nn.Linear(self.embed_dim, self.embed_dim, config.bias)
self.attn_dropout = nn.Dropout(config.dropout)
self.residual_dropout = nn.Dropout(config.dropout)
self.use_flash_attention = hasattr(torch.nn.functional, "scaled_dot_product_attention")
def forward(self, inputs: torch.Tensor, mask: torch.Tensor, freq_components) -> torch.Tensor:
batch_size, seq_len, _ = inputs.shape
# Project to queries, keys, values
keys = self.key_proj(inputs)
queries = self.query_proj(inputs)
values = self.value_proj(inputs)
# Reshape for multi-head attention
keys = keys.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
queries = queries.view(batch_size, seq_len, self.num_heads, self.head_dim)
values = values.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
# Apply RoPE
queries, keys = rotate_embeddings(keys, queries, freq_components)
# Handle grouped query attention
if self.num_kv_heads != self.num_heads:
keys = torch.repeat_interleave(keys, self.queries_per_kv, dim=2)
values = torch.repeat_interleave(values, self.queries_per_kv, dim=2)
# Prepare for attention computation
keys = keys.transpose(1, 2)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)
# Compute attention scores
attention_scores = torch.matmul(queries, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
attention_scores = attention_scores + mask[:, :, :seq_len, :seq_len]
# Add CoPE embeddings
attention_scores = attention_scores + self.cope(queries, attention_scores)
# Apply softmax and dropout
attention_weights = F.softmax(attention_scores.float(), dim=-1).type_as(keys)
attention_weights = self.attn_dropout(attention_weights)
# Compute weighted sum
context = torch.matmul(attention_weights, values)
# Reshape and project output
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
output = self.output_proj(context)
output = self.residual_dropout(output)
return output
The line attention_scores = attention_scores + self.cope(queries, attention_scores) is where CoPE embeddings are integrated into the context.
CoPE enhances the model's ability to accurately understand and process long-range dependencies in language by introducing dynamic, content-aware position information. This is particularly valuable for complex sentences or documents where correct interpretation depends on understanding distant relationships between words. For domain-specific or diverse data, CoPE improves model flexibility and generalization by adapting to different text features and structures. In highly context-sensitive tasks like machine translation, text summarization, or dialogue systems, CoPE can significantly boost performance.
References
For further reading on these methods, consider the following research papers:
https://avoid.overfit.cn/post/91fd4283a7944bebabb6017f5ee285e9
Author: Zain ul Abideen