ALiBi › ALiBi in Attention: Complete Implementation
Language AI HandbookAll chapters ALiBi › ALiBi in Attention: Complete ImplementationOne section · Sections in this chapter7 / 16 ↓
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
Now we have all the pieces to implement ALiBi-augmented attention from scratch. The implementation closely mirrors standard attention. We compute and scale the score matrix, add the bias, apply the optional causal mask, take the softmax, and use the result to weight the values.
The only new element is the construction and addition of the ALiBi bias tensor. For multi-head attention, we need one bias matrix per head, which we stack into a three-dimensional tensor of shape (num_heads, seq_len, seq_len). Each slice along the first dimension is the bias matrix for one head, with that head's specific slope.
def softmax(x, axis=-1):
"""Numerically stable softmax."""
x_max = np.max(x, axis=axis, keepdims=True)
exp_x = np.exp(x - x_max)
return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
def alibi_attention(Q, K, V, slopes, causal=True):
"""
Compute attention with ALiBi position encoding.
Args:
Q: Query matrix of shape (num_heads, seq_len, d_k)
K: Key matrix of shape (num_heads, seq_len, d_k)
V: Value matrix of shape (num_heads, seq_len, d_v)
slopes: ALiBi slopes, one per head, shape (num_heads,)
causal: Whether to apply causal masking
Returns:
Output of shape (num_heads, seq_len, d_v)
Attention weights of shape (num_heads, seq_len, seq_len)
"""
num_heads, seq_len, d_k = Q.shape
# Compute raw attention scores: Q @ K^T
# Shape: (num_heads, seq_len, seq_len)
scores = np.matmul(Q, K.transpose(0, 2, 1))
# Scale by sqrt(d_k)
scores = scores / np.sqrt(d_k)
# Create ALiBi biases for each head
# Shape: (num_heads, seq_len, seq_len)
positions = np.arange(seq_len)
distances = np.abs(positions[:, None] - positions[None, :])
# Broadcast slopes: (num_heads, 1, 1) * (seq_len, seq_len)
alibi_bias = -slopes[:, None, None] * distances[None, :, :]
# Apply causal mask if needed
if causal:
causal_mask = np.triu(np.ones((seq_len, seq_len)), k=1)
alibi_bias = np.where(causal_mask == 1, -np.inf, alibi_bias)
# Add ALiBi bias to scores
scores = scores + alibi_bias
# Apply softmax to get attention weights
attention_weights = softmax(scores, axis=-1)
# Handle NaN from -inf (all masked positions)
attention_weights = np.nan_to_num(attention_weights, nan=0.0)
# Compute output: weighted sum of values
output = np.matmul(attention_weights, V)
return output, attention_weightsThe broadcasting step deserves attention. We have slopes with shape (num_heads,). Reshaping to (num_heads, 1, 1) and multiplying by distances[None, :, :] with shape (1, seq_len, seq_len) produces a tensor of shape (num_heads, seq_len, seq_len). Each head's slice is the distance matrix multiplied by that head's slope. This is a single vectorized operation that replaces what would otherwise be a loop over heads.
Let's test this implementation with a concrete example to verify the shapes and inspect the attention patterns:
# Create a simple example
np.random.seed(42)
num_heads = 4
seq_len = 8
d_k = 16
d_v = 16
# Random Q, K, V matrices
Q = np.random.randn(num_heads, seq_len, d_k) * 0.5
K = np.random.randn(num_heads, seq_len, d_k) * 0.5
V = np.random.randn(num_heads, seq_len, d_v) * 0.5
# Get ALiBi slopes
slopes = get_alibi_slopes(num_heads)
# Compute ALiBi attention
output, attention_weights = alibi_attention(Q, K, V, slopes, causal=True)Output shape: (4, 8, 16) Attention weights shape: (4, 8, 8) Attention weights for head 1 (steep slope = 0.250): [[1. 0. 0. 0. 0. 0. 0. 0. ] [0.45 0.55 0. 0. 0. 0. 0. 0. ] [0.233 0.37 0.397 0. 0. 0. 0. 0. ] [0.227 0.205 0.262 0.306 0. 0. 0. 0. ] [0.118 0.068 0.121 0.279 0.414 0. 0. 0. ] [0.083 0.086 0.13 0.201 0.176 0.324 0. 0. ] [0.065 0.089 0.092 0.127 0.137 0.272 0.218 0. ] [0.025 0.038 0.057 0.073 0.136 0.214 0.233 0.224]] Attention weights for head 4 (gentle slope = 0.003906): [[1. 0. 0. 0. 0. 0. 0. 0. ] [0.562 0.438 0. 0. 0. 0. 0. 0. ] [0.36 0.453 0.187 0. 0. 0. 0. 0. ] [0.344 0.23 0.245 0.181 0. 0. 0. 0. ] [0.184 0.232 0.181 0.169 0.233 0. 0. 0. ] [0.121 0.125 0.286 0.214 0.096 0.158 0. 0. ] [0.104 0.124 0.171 0.176 0.08 0.175 0.169 0. ] [0.109 0.137 0.063 0.124 0.158 0.147 0.163 0.099]]
Compare the attention patterns between heads. Head 1 with its steep slope concentrates attention heavily on recent positions, with weights dropping rapidly as distance increases. The rows show a clear diagonal structure: most of the weight is on the most recent 1 to 3 tokens, with essentially nothing beyond 5 positions. Head 4 with its gentle slope distributes attention more evenly across the available context. The rows are more uniform, with distant positions receiving non-negligible weight.
To understand the impact of ALiBi more directly, let's compare attention patterns with and without the position bias. We'll compute attention for the same queries and keys, once with standard attention (no position encoding) and once with ALiBi:
The difference is striking. Standard attention distributes weight based purely on content similarity, sometimes attending strongly to distant positions because the content happens to match well. ALiBi reshapes this pattern, pulling attention toward recent tokens while still allowing content to influence the final distribution. This locality bias emerges from a single matrix addition, requiring no learned position embeddings and no modifications to the token representations themselves.
