30. ALiBi Relative Positional Bias in Attention Scores

Problem

Standard attention has no built-in notion of distance — a query attends to a key 1 token back and a key 500 tokens back with equal ease unless positional information is injected some other way. ALiBi (Attention with Linear Biases) is one of the simplest ways to inject that notion directly into the attention logits: instead of adding position information to the input embeddings, subtract a distance-proportional penalty from each attention score before the softmax, with a different penalty slope per head so different heads can specialize in different effective ranges.

Implement causal self-attention with an added ALiBi bias term: for query position i and key position j (j <= i), add -slope_h * (i - j) to the raw attention logit for head h, where slope_h is that head's fixed (not learned) scalar slope. Then apply the usual causal mask and softmax.

Source: src/30_relative_position_bias.py

def alibi_attention(
    q: torch.Tensor,       # (B, H, T, D)
    k: torch.Tensor,       # (B, H, T, D)
    v: torch.Tensor,       # (B, H, T, D)
    slopes: torch.Tensor,  # (H,) one ALiBi slope per head
) -> torch.Tensor:  # (B, H, T, D)
    ...

>>> q = k = v = torch.randn(1, 2, 5, 4)
>>> slopes = torch.tensor([1.0, 0.5])
>>> out = alibi_attention(q, k, v, slopes)
>>> out.shape
torch.Size([1, 2, 5, 4])

Step-by-Step Approach

1. Compute the ordinary scaled dot-product scores, (Q @ K^T) / sqrt(D), shape (B, H, T, T).
2. Build the (T, T) relative-distance matrix i - j from two broadcasted torch.arange(T) grids, clamped at 0 since only j <= i distances are ever used.
3. Reshape slopes from (H,) to (H, 1, 1) so it broadcasts against the (T, T) distance matrix, producing a per-head bias of shape (H, T, T): bias = -slopes.view(H,1,1) * distance.
4. Add that bias to the raw scores — note it needs an extra leading batch dimension ((1, H, T, T)) to broadcast correctly against the (B, H, T, T) scores tensor.
5. Apply the ordinary causal mask on top (j <= i to -inf) — the ALiBi bias and the causal mask are two independent, additive corrections to the same logits, not a replacement for one another.
6. Softmax and matmul with V as usual. Since the bias is always <= 0 for valid positions, it never needs a numerical-stability fix of its own — softmax's max-subtraction trick still handles the combined logits fine.

The key insight is that ALiBi needs no learned parameters and no change to Q/K/V at all — it's a pure post-hoc additive correction to the logits, which is why per-head broadcasting shape discipline ((H,1,1) against (T,T), then an extra batch dim against (B,H,T,T)) is the entire implementation difficulty.

Reference solution

import math

import torch


def alibi_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    slopes: torch.Tensor,
) -> torch.Tensor:
    B, H, T, D = q.shape
    scale = 1.0 / math.sqrt(D)

    scores = torch.matmul(q, k.transpose(-1, -2)) * scale  # (B, H, T, T)

    qi = torch.arange(T, device=q.device).unsqueeze(1)  # (T, 1)
    kj = torch.arange(T, device=q.device).unsqueeze(0)  # (1, T)
    causal = kj <= qi                                    # (T, T) bool

    # relative distance i-j (only meaningful where causal is True; clamp so
    # non-causal entries don't produce negative "distances" before masking)
    rel = (qi - kj).clamp(min=0).float()                 # (T, T)

    # per-head linear penalty: slopes is (H,) -> (H, 1, 1) broadcasts against (T, T)
    bias = -slopes.view(H, 1, 1) * rel.unsqueeze(0)       # (H, T, T)

    scores = scores + bias.unsqueeze(0)                   # (B, H, T, T), broadcast over batch
    scores = scores.masked_fill(~causal.view(1, 1, T, T), float("-inf"))

    attn = torch.softmax(scores, dim=-1)
    return torch.matmul(attn, v)


TEST_CASES = [
    {"name": "two heads, distinct slopes", "B": 1, "H": 2, "T": 6, "D": 4, "slopes": [1.0, 0.5]},
    {"name": "single head", "B": 2, "H": 1, "T": 4, "D": 8, "slopes": [0.25]},
    {"name": "four heads, geometric slopes", "B": 1, "H": 4, "T": 7, "D": 4, "slopes": [1.0, 0.5, 0.25, 0.125]},
]


def main():
    import torch.nn.functional as F

    torch.manual_seed(0)

    for i, case in enumerate(TEST_CASES):
        B, H, T, D = case["B"], case["H"], case["T"], case["D"]
        slopes = torch.tensor(case["slopes"])
        print(f"Test {i}: {case['name']} (B={B}, H={H}, T={T}, D={D}, slopes={case['slopes']})")
        q = torch.randn(B, H, T, D)
        k = torch.randn(B, H, T, D)
        v = torch.randn(B, H, T, D)

        out = alibi_attention(q, k, v, slopes)
        print(f"  out.shape={tuple(out.shape)}")
        assert out.shape == (B, H, T, D)

        qi = torch.arange(T).unsqueeze(1)
        kj = torch.arange(T).unsqueeze(0)
        causal = kj <= qi
        rel = (qi - kj).clamp(min=0).float()
        bias = -slopes.view(H, 1, 1) * rel.unsqueeze(0)
        additive_mask = bias.masked_fill(~causal.unsqueeze(0), float("-inf"))
        expected = F.scaled_dot_product_attention(q, k, v, attn_mask=additive_mask)

        torch.testing.assert_close(out, expected, atol=1e-4, rtol=1e-4)
        print("PASSED")


if __name__ == "__main__":
    main()

Key Functions & Tricks

• qi - kj with broadcasting — column-vector minus row-vector torch.arange grids produce the full (T, T) pairwise relative-distance matrix in one op.
• tensor.clamp(min=0) — zeroes out distances for non-causal (future) positions before they're multiplied by a slope, avoiding a spurious sign flip that masking alone wouldn't catch.
• slopes.view(H, 1, 1) — reshapes a per-head 1-D slope vector so it broadcasts elementwise against a (T, T) or (H, T, T) tensor without an explicit loop over heads.
• bias.unsqueeze(0) — adds the batch dimension back so a (H, T, T) bias broadcasts correctly against (B, H, T, T) scores.
• masked_fill(~causal, float('-inf')) — applied after the bias is added, so the causal mask always wins regardless of what the (always <=0) bias contributed.
• F.scaled_dot_product_attention(..., attn_mask=bias_plus_mask) — used as the test oracle, taking one combined additive mask (bias + -inf causal terms) that mirrors what the student's own logits should look like.

How to Recognize This Pattern