Phi Models › Code Implementation: Phi-Style Attention
Language AI HandbookAll chapters Phi Models › Code Implementation: Phi-Style AttentionOne section · Sections in this chapter7 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
To truly understand how Phi models work, we need to examine their core computational mechanism: multi-head attention with rotary positional embeddings (RoPE). While the architecture is standard, implementing it from scratch reveals the elegant mathematics that enables transformers to process sequences.
Our implementation journey proceeds through three stages:
- Positional encoding: How RoPE embeds position information through rotation
- Attention computation: How QKV tensors interact to produce contextual representations
- Integration: How these components combine in a complete attention layer
Understanding Rotary Position Embeddings
Before diving into code, let's build intuition for why positional encoding matters. A transformer's attention mechanism computes similarity between all pairs of tokens, but this computation is inherently position-agnostic. The same query-key pair produces the same similarity score regardless of where the tokens appear in the sequence. Yet position clearly matters: "The cat sat on the mat" has different meaning than "mat the on sat cat The."
Traditional approaches add positional embeddings to token embeddings before attention. RoPE takes a different approach: it encodes position by rotating the query and key vectors in embedding space. The key insight is that rotation preserves vector magnitudes while changing the angle between vectors. When we compute the dot product (attention score) between a rotated query and key, the result depends on their relative positions, not just their content.
Think of it geometrically: if we rotate a query vector by angle and a key vector by angle , their dot product depends on the difference . By making these rotation angles position-dependent, we encode relative position directly into the attention computation.
import numpy as np
def apply_rotary_embeddings(x, cos, sin):
"""
Apply rotary positional embeddings to input tensor.
RoPE rotates pairs of dimensions in the embedding space by
position-dependent angles, letting the model to encode
relative positions through the attention dot product.
Args:
x: Input tensor of shape (batch, heads, seq_len, head_dim)
cos: Cosine of rotation angles (seq_len, head_dim)
sin: Sine of rotation angles (seq_len, head_dim)
Returns:
Rotated tensor of same shape as input
"""
# Split into pairs for rotation
x1 = x[..., ::2] # Even indices
x2 = x[..., 1::2] # Odd indices
# Reshape cos/sin for broadcasting
cos = cos[np.newaxis, np.newaxis, :, :]
sin = sin[np.newaxis, np.newaxis, :, :]
# Apply rotation
rotated_x1 = x1 * cos - x2 * sin
rotated_x2 = x1 * sin + x2 * cos
# Interleave back together
result = np.empty_like(x)
result[..., ::2] = rotated_x1
result[..., 1::2] = rotated_x2
return result
def compute_rope_frequencies(head_dim, max_seq_len, base=10000):
"""
Compute RoPE frequency components.
These frequencies determine how quickly each dimension
rotates as position increases, with lower-index dimensions
rotating more slowly to capture longer-range patterns.
"""
# Frequency for each dimension pair
dim_indices = np.arange(0, head_dim, 2)
freqs = 1.0 / (base ** (dim_indices / head_dim))
# Position indices
positions = np.arange(max_seq_len)
# Compute angles: (seq_len, head_dim/2)
angles = np.outer(positions, freqs)
# Expand to full head_dim by repeating
cos = np.cos(angles)
sin = np.sin(angles)
return cos, sinThe Mathematics of Rotation
RoPE applies a 2D rotation to each pair of consecutive dimensions in the embedding. A 2D rotation by angle transforms a vector to:
This is exactly what our apply_rotary_embeddings function computes: for each pair of dimensions, it applies the rotation using the precomputed cosine and sine values.
The rotation angle for position and dimension pair is computed as:
where:
- : the position index in the sequence (0, 1, 2, ...)
- : the dimension pair index (0, 1, 2, ..., )
- : the head dimension (e.g., 64)
- : a hyperparameter controlling the frequency scale (typically 10,000)
Lower-indexed dimensions rotate more slowly (smaller ), capturing longer-range positional relationships, while higher-indexed dimensions rotate faster, encoding fine-grained local position information.
This frequency hierarchy is important: it allows the model to represent both local patterns (through high-frequency dimensions) and long-range dependencies (through low-frequency dimensions) simultaneously. Let's see this in practice:
RoPE Frequency Computation ================================================== Head dimension: 64 Sequence length: 8 Cosine shape: (8, 32) Sine shape: (8, 32) First 4 positions, first 4 frequency components: Cosine values: [[ 1. 1. 1. 1. ] [ 0.54 0.732 0.846 0.912] [-0.416 0.071 0.431 0.665] [-0.99 -0.628 -0.116 0.301]] Sine values: [[0. 0. 0. 0. ] [0.841 0.682 0.533 0.409] [0.909 0.997 0.902 0.747] [0.141 0.778 0.993 0.954]]
The output shows how RoPE frequencies vary across positions and dimensions. Notice that the first column (lowest frequency) changes slowly across positions, while higher-frequency components (rightmost columns) oscillate more rapidly. At position 0, all sine values are 0 and cosine values are 1, meaning no rotation occurs. As position increases, the rotation angles grow, with higher-indexed dimensions accumulating rotation faster.
From Position Encoding to Attention
With position encoding understood, we can now build the complete attention mechanism. The core idea of attention is simple: each position in the sequence should be able to gather information from other positions, weighted by relevance. A query asks "what information do I need?", keys answer "here's what I have", and values provide the actual information to aggregate.
The mathematical formulation captures this intuition precisely. The core attention computation follows the standard scaled dot-product formula:
where:
- : the query matrix, where each row represents a position asking "what should I attend to?"
- : the key matrix, where each row represents a position answering "here's what I contain"
- : the value matrix, containing the actual information to be aggregated
- : the key/query dimension (typically 64 in Phi models)
- : the sequence length
- : scaling factor that prevents dot products from becoming too large
The softmax normalizes attention scores so each query position's weights sum to 1, creating a valid probability distribution over key positions.
Why divide by ? As the dimension grows, dot products tend to grow in magnitude (they're sums of terms). Large dot products push softmax into regions where gradients vanish, making training unstable. The scaling factor keeps the variance of dot products roughly constant regardless of .
Putting It All Together: The PhiAttention Class
Now we can combine RoPE with scaled dot-product attention in a complete implementation. The PhiAttention class orchestrates the full computation:
- Project input to QKV tensors using learned weight matrices
- Reshape to separate attention heads (each head attends independently)
- Apply RoPE to queries and keys (encoding position)
- Compute scaled dot-product attention
- Apply causal mask for autoregressive generation
- Project output back to model dimension
class PhiAttention:
"""
Multi-head attention as used in Phi models.
Combines standard scaled dot-product attention with:
- Rotary positional embeddings (RoPE)
- Pre-normalization
- Flash Attention compatibility (not implemented here)
"""
def __init__(self, hidden_size=2048, num_heads=32, head_dim=64, seed=42):
"""
Initialize attention weights.
Args:
hidden_size: Model hidden dimension
num_heads: Number of attention heads
head_dim: Dimension per attention head
seed: Random seed for initialization
"""
self.hidden_size = hidden_size
self.num_heads = num_heads
self.head_dim = head_dim
np.random.seed(seed)
# Phi uses small initialization scale
scale = 0.02
# Query, Key, Value projections
self.W_q = np.random.randn(hidden_size, num_heads * head_dim) * scale
self.W_k = np.random.randn(hidden_size, num_heads * head_dim) * scale
self.W_v = np.random.randn(hidden_size, num_heads * head_dim) * scale
# Output projection
self.W_o = np.random.randn(num_heads * head_dim, hidden_size) * scale
# Pre-compute RoPE frequencies for max sequence length
self.max_seq_len = 2048
self.cos, self.sin = compute_rope_frequencies(
head_dim // 2 * 2, # Ensure even
self.max_seq_len,
)
def forward(self, x, mask=None):
"""
Compute attention output.
Args:
x: Input tensor (batch, seq_len, hidden_size)
mask: Optional causal mask (seq_len, seq_len)
Returns:
output: Attention output (batch, seq_len, hidden_size)
weights: Attention weights (batch, heads, seq_len, seq_len)
"""
batch, seq_len, _ = x.shape
# Project to Q, K, V
Q = x @ self.W_q # (batch, seq_len, num_heads * head_dim)
K = x @ self.W_k
V = x @ self.W_v
# Reshape to (batch, num_heads, seq_len, head_dim)
Q = Q.reshape(batch, seq_len, self.num_heads, self.head_dim)
Q = Q.transpose(0, 2, 1, 3)
K = K.reshape(batch, seq_len, self.num_heads, self.head_dim)
K = K.transpose(0, 2, 1, 3)
V = V.reshape(batch, seq_len, self.num_heads, self.head_dim)
V = V.transpose(0, 2, 1, 3)
# Apply RoPE to Q and K
cos_slice = self.cos[:seq_len]
sin_slice = self.sin[:seq_len]
Q = apply_rotary_embeddings(Q, cos_slice, sin_slice)
K = apply_rotary_embeddings(K, cos_slice, sin_slice)
# Compute attention scores using scaled dot-product attention
# The scaling factor sqrt(d_k) prevents dot products from growing
# too large as the head dimension increases
scale = np.sqrt(self.head_dim)
scores = (Q @ K.transpose(0, 1, 3, 2)) / scale
# Apply causal mask if provided (set masked positions to -inf)
if mask is not None:
scores = np.where(mask, scores, -1e9)
# Softmax converts scores to attention weights that sum to 1
exp_scores = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights = exp_scores / exp_scores.sum(axis=-1, keepdims=True)
# Weighted sum of values produces the attention output
output = weights @ V # (batch, heads, seq_len, head_dim)
# Reshape back
output = output.transpose(
0, 2, 1, 3
) # (batch, seq_len, heads, head_dim)
output = output.reshape(batch, seq_len, self.num_heads * self.head_dim)
# Output projection
output = output @ self.W_o
return output, weightsThe implementation follows the mathematical formulation closely. Notice how the forward method mirrors our six-step process: projection and reshaping, RoPE application, attention computation, masking, then output projection.
Key implementation details:
- Small initialization scale (0.02): Phi uses smaller weight initialization than some models, which can help training stability
- Pre-computed RoPE frequencies: We compute cosine and sine values once and slice them per sequence length, avoiding redundant computation
- Causal masking: The mask ensures position can only attend to positions , letting autoregressive generation
Testing the Implementation
Let's verify our attention layer works correctly by processing a small sequence and examining the outputs:
# Test the attention implementation
batch_size = 1
seq_len = 10
hidden_size = 256
num_heads = 4
head_dim = 64
# Create deterministic test input (seeded once in the shared setup cell)
x = np.random.randn(batch_size, seq_len, hidden_size)
# Create causal mask
causal_mask = np.tril(np.ones((seq_len, seq_len), dtype=bool))
# Initialize and run attention
attention = PhiAttention(
hidden_size=hidden_size, num_heads=num_heads, head_dim=head_dim, seed=17
)
output, weights = attention.forward(x, mask=causal_mask)Phi-Style Attention Test ================================================== Input shape: (1, 10, 256) Output shape: (1, 10, 256) Weights shape: (1, 4, 10, 10) Attention weights sum per query (should be 1.0): [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
The output confirms our implementation is working correctly. The input and output shapes match (preserving dimensions through the attention layer), and each row of attention weights sums to exactly 1.0, showing proper softmax normalization. The weights tensor has shape (1, 4, 6, 6), representing 4 attention heads each creating a 6x6 attention matrix for the 6-position sequence.
Visualizing Attention Patterns
The causal mask creates a distinctive lower-triangular pattern in the attention weights. This pattern is basic to autoregressive language modeling: when predicting position 5, the model can only see positions 0-5, not future positions 6 and beyond. Let's visualize this:
