ALiBi › Worked Example: Tracing Through ALiBi Step by Step

Michael BrenndoerferUpdated May 30, 20253 min read
Language AI HandbookAll chapters
Part 14 of 67Positional Encoding
Chapter 101ALiBi › Worked Example: Tracing Through ALiBi Step by Step
Reading section 8 of 16 · Worked Example: Tracing Through ALiBi Step by Step · 3 minSection 8 of 16 · 3 min

Book contents

Section 8 of 16 · 3 min read

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 dk=4d_k = 4. We'll use simple numbers so you can verify each step by hand.

In[19]:
Code
# 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 heads
Out[20]:
Console
Step 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 Q=KQ = K with identity-like structure, each token scores highest against itself. In step 2, we scale by 4=2\sqrt{4} = 2, which shrinks all values without changing their relative ordering. In step 3, the ALiBi bias is purely a function of distance: entry (i,j)(i, j) is −m⋅∣i−j∣-m \cdot |i-j| where mm 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 −∞-\infty 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.