Position Encoding Comparison › Implementation Complexity

# Sinusoidal: Standalone function, applied once to input
def apply_sinusoidal(embeddings, pos_encodings):
    """Add pre-computed sinusoidal encodings to token embeddings."""
    seq_len = embeddings.shape[0]
    return embeddings + pos_encodings[:seq_len]


# Learned: Simple lookup and add
def apply_learned(embeddings, pos_embeddings):
    """Add learned position embeddings to token embeddings."""
    seq_len = embeddings.shape[0]
    return embeddings + pos_embeddings[:seq_len]


# RoPE: Requires modifying Q and K within attention
def apply_rope(q, k, cos, sin):
    """
    Apply rotary position embedding to queries and keys.
    Rotates pairs of dimensions by position-dependent angles.
    """
    # Split into pairs for rotation
    q_even, q_odd = q[..., 0::2], q[..., 1::2]
    k_even, k_odd = k[..., 0::2], k[..., 1::2]

    # Apply 2D rotation to each pair
    q_rotated = np.concatenate(
        [q_even * cos - q_odd * sin, q_even * sin + q_odd * cos], axis=-1
    )

    k_rotated = np.concatenate(
        [k_even * cos - k_odd * sin, k_even * sin + k_odd * cos], axis=-1
    )

    return q_rotated, k_rotated


# ALiBi: Requires modifying attention scores
def apply_alibi(attention_scores, alibi_bias):
    """Add ALiBi bias to attention scores before softmax."""
    return attention_scores + alibi_bias

def attention_with_sinusoidal(X, W_q, W_k, W_v, pos_encodings):
    """Standard attention with sinusoidal position encoding."""
    # Position info added at input
    X_pos = X + pos_encodings[: X.shape[0]]

    Q = X_pos @ W_q
    K = X_pos @ W_k
    V = X_pos @ W_v

    scores = Q @ K.T / np.sqrt(K.shape[-1])
    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights = weights / weights.sum(axis=-1, keepdims=True)

    return weights @ V


def attention_with_rope(X, W_q, W_k, W_v, cos, sin):
    """Attention with RoPE applied to Q and K."""
    Q = X @ W_q
    K = X @ W_k
    V = X @ W_v

    # RoPE modifies Q and K before score computation
    Q_rot, K_rot = apply_rope(Q, K, cos, sin)

    scores = Q_rot @ K_rot.T / np.sqrt(K.shape[-1])
    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights = weights / weights.sum(axis=-1, keepdims=True)

    return weights @ V


def attention_with_alibi(X, W_q, W_k, W_v, alibi_bias):
    """Attention with ALiBi bias on scores."""
    Q = X @ W_q
    K = X @ W_k
    V = X @ W_v

    scores = Q @ K.T / np.sqrt(K.shape[-1])
    # ALiBi modifies scores before softmax
    scores = scores + alibi_bias

    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights = weights / weights.sum(axis=-1, keepdims=True)

    return weights @ V