ALiBi › The Mathematical Formulation

Now that we understand ALiBi's core insight, let's develop the complete mathematical framework. We will build from standard attention to ALiBi-augmented attention, showing exactly where and how the position bias enters the computation.

Standard attention and ALiBi attention are nearly identical. The only difference is one additional term added to the score matrix before the softmax. This surgical simplicity is part of what makes ALiBi attractive: it requires minimal changes to an existing attention implementation and adds almost no computational overhead.

Recall that standard self-attention operates on three matrices derived from the input sequence. For a sequence of $n$ tokens, the attention mechanism computes:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

where:

• $Q \in \mathbb{R}^{n \times d_k}$: the query matrix containing query vectors for all $n$ positions
• $K \in \mathbb{R}^{n \times d_k}$: the key matrix containing key vectors for all positions
• $V \in \mathbb{R}^{n \times d_v}$: the value matrix containing value vectors for all positions
• $d_k$: the dimension of queries and keys
• $QK^T \in \mathbb{R}^{n \times n}$: the raw attention score matrix where entry $(i, j)$ is the dot product between query $i$ and key $j$
• $\sqrt{d_k}$: scaling factor that prevents dot products from growing too large in high dimensions, which would push the softmax into regions of very small gradients

The matrix $QK^T$ captures content-based similarity: how much each query "wants" to attend to each key based purely on their learned representations. This matrix is completely blind to position. Token 1 attending to token 2 produces the same content-based score whether they appear next to each other or 500 positions apart. This is by design for the core attention mechanism, which wants to capture global relationships. But it means we need another mechanism to introduce positional awareness. That mechanism is the ALiBi bias.

Why does ALiBi add bias at this stage rather than modifying the queries and keys? Adding to the pre-softmax scores is the most direct possible intervention. The scores are the raw votes about which tokens matter. By adjusting the scores before normalization, we influence the attention distribution in a clear and simple way. Modifying queries and keys (as RoPE does) achieves a similar effect but through a more complex indirect path involving embedding geometry.

ALiBi's modification is surgical. Rather than changing how queries, keys, or values are computed, it adds a single term to the attention scores. The complete ALiBi-augmented attention reads:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + B\right)V$$

The bias matrix $B \in \mathbb{R}^{n \times n}$ is added directly to the scaled scores before softmax. This is the only change to standard attention. The matrix $B$ encodes position information through a remarkably simple rule: each entry is determined entirely by the distance between the corresponding positions. Specifically,

$$B_{ij} = -m \cdot |i - j|$$

where:

• $B_{ij}$: the bias added to the attention score between query position $i$ and key position $j$
• $m$: the slope parameter for this attention head, controlling the strength of the distance penalty
• $|i - j|$: the absolute distance between positions $i$ and $j$ in the token sequence

Let's unpack why this formula works. The absolute distance $|i - j|$ measures how far apart two positions are in the sequence. Adjacent tokens have distance 1, tokens separated by 10 positions have distance 10. Multiplying by a positive slope $m$ and negating creates a penalty that grows linearly with distance. A token that is close gets a small penalty; a token that is far gets a large penalty.

Consider what happens for a query at position 5 attending to various keys:

• Attending to position 5 (itself): $B_{55} = -m \cdot 0 = 0$ (no penalty, full score preserved)
• Attending to position 4 (adjacent): $B_{54} = -m \cdot 1 = -m$ (small penalty)
• Attending to position 3 (two steps away): $B_{53} = -m \cdot 2 = -2m$ (moderate penalty)
• Attending to position 0 (distant): $B_{50} = -m \cdot 5 = -5m$ (large penalty)

The negative bias reduces the attention score, and larger distances produce more negative biases. After softmax normalization, this translates to lower attention weights for distant tokens. The content-based similarity in $QK^T$ still matters completely: a very semantically relevant distant token can overcome its distance penalty if its dot product score is large enough. But now content has to compete against a structural penalty that favors proximity.

Why does this formula make sense? Notice that the formula simply implements a universal assumption that language models hold implicitly: nearby context tends to be more relevant than distant context. Rather than forcing the model to learn this bias from data (which is expensive and may lead to overfitting on long-range spurious correlations), ALiBi builds it in as an inductive bias. The model is free to override it with strong content signals, but the default behavior is local attention. This is similar to how convolutional neural networks build in a translation-invariance inductive bias, rather than forcing the network to learn that from scratch.

For decoder-style models that generate text left-to-right, we need causal masking: position $i$ can only attend to positions $j \le i$. The combined bias matrix integrates both the ALiBi distance penalties and the causal mask:

$$B = \begin{bmatrix} 0 & -\infty & -\infty & -\infty & \cdots \\ -m & 0 & -\infty & -\infty & \cdots \\ -2m & -m & 0 & -\infty & \cdots \\ -3m & -2m & -m & 0 & \cdots \\ \vdots & \vdots & \vdots & \vdots & \ddots \end{bmatrix}$$

The structure reveals two distinct components working together. The lower triangle, including the diagonal, contains the ALiBi distance penalties. Entry $(i, j)$ with $j \le i$ holds the value $-m \cdot (i - j)$, which is zero on the diagonal and grows more negative as you move left along any row. The upper triangle holds $-\infty$ values from the causal mask.

The $-\infty$ entries ensure that after applying $\exp(-\infty) = 0$ inside the softmax, future positions contribute nothing to the weighted sum. The linear penalties in the lower triangle then shape how attention flows among the allowed past positions. These two components operate in entirely different regimes: the causal mask is a hard prohibition (zero weight, always), while the ALiBi penalty is a soft preference (lower weight, but nonzero).

Notice that the causal mask interacts naturally with ALiBi. For a model processing position $i$, the only tokens it can attend to are positions 0 through $i$. These are exactly the tokens for which $j \le i$, so the distance is always $i - j$, always positive, and always measured in the backward direction. The ALiBi penalty is therefore always a backward-looking decay, consistently favoring the most recent tokens in the context.

Let's translate this mathematics into code, building up the bias matrix piece by piece. The core insight behind the implementation is that we can compute all pairwise distances simultaneously using broadcasting, rather than iterating over all pairs with nested loops.