ALiBi › Worked Example: Tracing Through ALiBi Step by Step
Language AI HandbookAll chapters ALiBi › Worked Example: Tracing Through ALiBi Step by StepOne section · Sections in this chapter8 / 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
Let's work through a concrete numerical example to solidify the mechanics. We'll use a very small sequence so every number is legible.
Consider a sequence of 4 tokens with 2 attention heads and . We'll use simple numbers so you can verify each step by hand.
# Small worked example with traceable numbers
np.random.seed(0)
n = 4 # sequence length
h = 2 # number of heads
dk = 4 # key/query dimension
# Simple Q and K with small values
Q_ex = np.array(
[
[
[1.0, 0.0, 0.5, 0.0],
[0.0, 1.0, 0.0, 0.5],
[0.5, 0.0, 1.0, 0.0],
[0.0, 0.5, 0.0, 1.0],
],
[
[1.0, 0.0, 0.5, 0.0],
[0.0, 1.0, 0.0, 0.5],
[0.5, 0.0, 1.0, 0.0],
[0.0, 0.5, 0.0, 1.0],
],
]
) # shape (2, 4, 4)
K_ex = Q_ex.copy() # Use K = Q for simplicity
V_ex = np.random.randn(h, n, dk) * 0.1
slopes_ex = get_alibi_slopes(h) # slopes for 2 headsStep 1: Raw dot-product scores QK^T (head 1): [[1.25 0. 1. 0. ] [0. 1.25 0. 1. ] [1. 0. 1.25 0. ] [0. 1. 0. 1.25]] Step 2: Scaled scores (divided by sqrt(4) = 2.000) (head 1): [[0.625 0. 0.5 0. ] [0. 0.625 0. 0.5 ] [0.5 0. 0.625 0. ] [0. 0.5 0. 0.625]] Step 3: ALiBi bias matrix (head 1, slope = 0.062): [[-0. -0.062 -0.125 -0.188] [-0.062 -0. -0.062 -0.125] [-0.125 -0.062 -0. -0.062] [-0.188 -0.125 -0.062 -0. ]] Step 4: ALiBi bias with causal mask applied (head 1): [[-0. -inf -inf -inf] [-0.062 -0. -inf -inf] [-0.125 -0.062 -0. -inf] [-0.188 -0.125 -0.062 -0. ]] Step 5: Final scores (scaled + bias) (head 1): [[ 0.625 -inf -inf -inf] [-0.062 0.625 -inf -inf] [ 0.375 -0.062 0.625 -inf] [-0.188 0.375 -0.062 0.625]] Step 6: Attention weights after softmax (head 1): [[1. 0. 0. 0. ] [0.335 0.665 0. 0. ] [0.341 0.22 0.438 0. ] [0.163 0.286 0.184 0.367]]
Let's read through this trace carefully. In step 1, the raw dot-product scores capture content similarity. Since we used with identity-like structure, each token scores highest against itself. In step 2, we scale by , which shrinks all values without changing their relative ordering. In step 3, the ALiBi bias is purely a function of distance: entry is where is the head's slope. For head 1 with the steeper slope, the penalties at distance 1, 2, and 3 are meaningful numbers. In step 4, the upper triangle is set to to enforce causality. In step 5, we add the bias to the scaled scores. Notice how the off-diagonal entries drop: the content similarity gets penalized by distance, pulling the distribution toward the diagonal. In step 6, softmax turns these into probabilities. The diagonal entries (attending to self) dominate strongly for head 1 because the ALiBi penalty has suppressed the off-diagonal scores.
This step-by-step trace illustrates the key property of ALiBi: it does not change what the model "wants" to attend to (the content signal), but it imposes a systematic cost for looking farther away. When content signals are strong, they can dominate. When content signals are weak or comparable, the distance penalty tips the balance toward nearby tokens.