ALiBi — Attention with Linear Biases

Press, Smith, & Lewis, "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation", ICLR 2022

Learning Objectives

By the end of this chapter, you will be able to:

1. Explain why learned and sinusoidal positional embeddings fail to extrapolate beyond training sequence lengths, and why this is a critical problem for deployment.
2. Describe the core insight of ALiBi: replace position embeddings with a simple linear distance penalty added directly to attention scores — requiring zero additional parameters.
3. Derive the ALiBi formula and explain the geometric slope schedule that allows different heads to attend at different distance scales simultaneously.
4. Compute ALiBi attention weights and outputs by hand for “The cat sat on mat” and compare them to the standard attention results from Chapter 1.
5. Implement a complete ALiBi class in both NumPy and PyTorch that you can run on any input.
6. Connect ALiBi to RoPE (Chapter 8), Flash Attention (Chapter 13), and modern inference systems, understanding the tradeoffs between rotation-based and bias-based position encoding.

Where this appears: ALiBi was introduced at Facebook AI Research and adopted by the BLOOM model (176B parameters, BigScience), MPT-7B/30B (MosaicML), and other production LLMs. It offers the simplest possible approach to positional encoding — no learned parameters, no trigonometric functions, just subtraction. Understanding ALiBi illuminates a deep question: how much positional information does attention truly need?

The Real Problem

Chapters 7 and 8 established two approaches to injecting position information into attention: relative position biases (Chapter 7) and rotary position embeddings (Chapter 8, RoPE). Both solve the problem of making attention aware of token order. But there is a deeper problem that neither the original sinusoidal embeddings nor learned position embeddings solve reliably: what happens when the model encounters sequences longer than those seen during training?

Position Embedding Failure Modes

Consider a language model trained on sequences of up to 2048 tokens. During training, the model learns (or is given) position representations for positions 0 through 2047. At inference time, a user provides a 4096-token document. What happens?

1. Learned position embeddings (GPT-2, BERT): The model has an embedding lookup table with exactly 2048 rows. Position 2049 simply does not exist. The model either crashes or truncates the input. There is no extrapolation at all.
2. Sinusoidal position embeddings (original Transformer): The sine/cosine functions can produce values at any position. But the model has never seen the specific activation patterns produced by positions beyond 2048. The learned attention weights are calibrated for a specific range of positional signals, and positions outside that range produce unpredictable behavior. Perplexity collapses.
3. RoPE (Chapter 8): Rotation-based encoding can generate rotations for any position. However, without techniques like NTK-aware scaling or YaRN, the effective frequency distribution at unseen positions differs from training. RoPE degrades more gracefully than learned embeddings but still requires extension methods for reliable extrapolation.

The Extrapolation Challenge

The root cause is architectural: all three methods above encode position into the token representations themselves. The model must then learn to interpret these positional signals during attention computation. When the positional signals at inference time differ from those seen during training, the learned interpretation breaks down.

Press, Smith, and Lewis (2022) asked a radical question: what if we bypass token representations entirely and inject position information directly into the attention scores? If the position signal is a simple, predictable function of distance — one that the model never needs to “learn to interpret” — then extrapolation becomes trivial. A token 3 positions away receives the same penalty regardless of whether position 3 appeared during training at index 100 or index 10,000.

From Intuition to Mathematics

The Distance Penalty Insight

Imagine you are sitting in a library, and someone asks you a question (your query). You can look at every book on the shelf (the keys and values), but you have a natural preference: the closer a book is to you, the more likely it is to be relevant. A book one shelf away is easy to reach; a book across the room requires more effort. You can still access distant books, but you need a stronger reason (a higher content-based score) to bother.

ALiBi formalizes this intuition. Given the standard scaled dot-product score between query $i$ and key $j$, ALiBi subtracts a penalty proportional to their distance:

The penalty $m_h \times |i - j|$ grows linearly with distance. Token $j$ that is 1 position away from $i$ receives a penalty of $m_h$; 2 positions away receives $2 m_h$; and so on. After this subtraction, the scores enter softmax normally. The effect is that distant tokens must have proportionally higher content scores to achieve the same attention weight as nearby tokens.

Head-Specific Slopes

The slope $m_h$ is different for each attention head, creating a multi-scale view of the sequence:

• Steep slopes (large $m_h$ like 0.5): The penalty grows quickly with distance. These heads become “local experts” that focus strongly on nearby tokens — useful for syntax, morphology, and local coherence.
• Gentle slopes (small $m_h$ like 0.0625): The penalty grows slowly. These heads can attend broadly across the sequence — useful for coreference, topic tracking, and long-range dependencies.

Together, the set of heads covers multiple distance scales simultaneously, much like wavelet analysis decomposes a signal into components at different frequencies. The model does not learn these slopes; they are fixed by a geometric schedule, eliminating any risk of position overfitting.

The Mathematical Definition

Symbol-by-Symbol Breakdown

For head $h$ with slope $m_h$, the ALiBi attention score between query position $i$ and key position $j$ is:

The attention weights are then obtained by applying softmax row-wise:

And the output is the weighted sum of values:

The Geometric Slope Schedule

Press et al. define the slopes using a geometric sequence:

For $H = 8$ heads, this gives:

Head 1 penalizes a token 10 positions away by $-5.0$ — essentially killing that attention score. Head 8 penalizes the same distance by only $-0.039$ — barely noticeable. Together, the heads form a multi-resolution view of the sequence, from hyper-local to near-global.

Interactive: ALiBi Bias Explorer

Use the controls below to explore how ALiBi reshapes attention. Adjust the slope to see how stronger penalties create more local attention patterns. Switch between heads to compare their different distance sensitivities. Click any row to see its attention distribution as a bar chart.

Interactive: Extrapolation Comparison

The visualization below compares how different positional encoding methods degrade when the input sequence exceeds the training length. ALiBi degrades most gracefully because its bias is a deterministic function of distance, not a learned representation that must generalize.

Step-by-Step Calculation

We use the same shared example as every chapter: tokens = ["The", "cat", "sat", "on", "mat"], with $d_{\text{model}} = 4$, 2 heads, and $d_k = 2$. Head 1 uses slope $m_1 = 0.5$ and Head 2 uses slope $m_2 = 0.25$.

Step-by-Step for "The" (Row 0), Head 1 ($m = 0.5$)

Step 1: Extract head slices. Head 1 uses columns 0–1:

Step 2: Compute raw dot products.

• $Q_{h1}[\text{The}] \cdot K_{h1}[\text{The}] = 1 \times 0 + 0 \times 1 = 0.0$
• $Q_{h1}[\text{The}] \cdot K_{h1}[\text{cat}] = 1 \times 1 + 0 \times 0 = 1.0$
• $Q_{h1}[\text{The}] \cdot K_{h1}[\text{sat}] = 1 \times 1 + 0 \times 1 = 1.0$
• $Q_{h1}[\text{The}] \cdot K_{h1}[\text{on}] = 1 \times 0 + 0 \times 0 = 0.0$
• $Q_{h1}[\text{The}] \cdot K_{h1}[\text{mat}] = 1 \times 1 + 0 \times 0 = 1.0$

Step 3: Scale by $\sqrt{d_k} = \sqrt{2} \approx 1.4142$.

Step 4: Compute ALiBi penalties. For slope $m = 0.5$ and query position $i = 0$:

Step 5: Add bias to scaled scores.

Notice: "cat" (score 0.707, penalty -0.500 = 0.207) beats "mat" (score 0.707, penalty -2.000 = -1.293) despite having identical content scores. Distance matters.

Step 6: Softmax.

• Shift by max (0.207): $[-0.207,\; 0.000,\; -0.500,\; -1.707,\; -1.500]$
• Exponentiate: $[0.813,\; 1.000,\; 0.607,\; 0.181,\; 0.223]$
• Sum: $2.824$
• Normalize: $A_{h1}[\text{The}] = [0.2879,\; 0.3541,\; 0.2148,\; 0.0642,\; 0.0790]$

Compare to standard attention (no ALiBi): $[0.1237,\; 0.2509,\; 0.2509,\; 0.1237,\; 0.2509]$. ALiBi dramatically shifts weight toward nearby tokens: "The" goes from 12.4% to 28.8%, while distant "on" drops from 12.4% to 6.4%.

Step-by-Step for "The" (Row 0), Head 2 ($m = 0.25$)

Head 2 uses columns 2–3: $Q_{h2}[\text{The}] = [1.0,\; 0.0]$. The scaled scores are $[0.000,\; 0.707,\; 0.000,\; 0.707,\; 0.354]$.

With the gentler slope of 0.25, the penalties are halved compared to Head 1:

Result: $A_{h2}[\text{The}] = [0.2142,\; 0.3384,\; 0.1299,\; 0.2052,\; 0.1122]$. Compare to Head 1: "on" receives 20.5% in Head 2 vs 6.4% in Head 1. The gentler slope lets Head 2 maintain broader context while Head 1 focuses locally.

Step-by-Step for "cat" (Row 1), Head 1

$Q_{h1}[\text{cat}] = [0.0,\; 2.0]$. The raw dot products:

• $[0,2] \cdot [0,1] = 2.0, \quad [0,2] \cdot [1,0] = 0.0, \quad [0,2] \cdot [1,1] = 2.0, \quad [0,2] \cdot [0,0] = 0.0, \quad [0,2] \cdot [1,0] = 0.0$

Scaled: $[1.414,\; 0.000,\; 1.414,\; 0.000,\; 0.000]$. ALiBi bias (distances from position 1): $[-0.5,\; 0,\; -0.5,\; -1.0,\; -1.5]$.

Weights: $[0.3791,\; 0.1520,\; 0.3791,\; 0.0559,\; 0.0339]$. "The" and "sat" (both distance 1) get identical penalties, so they tie at 37.9%. Distant "mat" (distance 3) drops to just 3.4%.

Full Attention Weights and Output

ALiBi vs Standard: Side-by-Side

Head 1 Weights ($m = 0.5$):

Head 2 Weights ($m = 0.25$):

Averaged Weights (Head 1 + Head 2) / 2:

Final Output ($5 \times 4$):

Interpreting the Differences

The averaged weight matrix reveals ALiBi's defining characteristic: diagonal dominance. Every token gives itself or its immediate neighbor the highest weight. The diagonal values (self-attention) are: The=0.251, cat=0.171, sat=0.365, on=0.416, mat=0.403. Contrast with standard attention where "on" distributes weight uniformly (all 0.200 in Head 1) because its query vector is zero.

Even when a token has zero content-based preference (like "on" in Head 1), ALiBi still creates meaningful attention by position alone: self-attention (distance 0) receives the highest weight, and the weight decays with distance. Position becomes the tiebreaker when content scores are equal.

Applications Across Domains

Long-Document Language Modeling

The original motivation for ALiBi was enabling language models to process documents longer than their training context. BLOOM (176B, BigScience) and MPT-7B/30B (MosaicML) adopted ALiBi for this reason. A model trained on 2048-token sequences can be deployed on 4096+ token documents without fine-tuning. This is critical for summarization, document QA, and legal/medical text processing where documents routinely exceed training lengths.

Code Generation

Source code files are often much longer than natural language training samples. A model trained on 2048-token code snippets may need to process 10,000+ token files at inference. ALiBi's extrapolation property allows the model to maintain coherent attention across entire files, keeping track of function definitions, variable scope, and import relationships that span thousands of tokens.

Scientific Sequence Modeling

Protein sequences (up to 30,000+ amino acids), genomic sequences, and time series data all benefit from length extrapolation. A model trained on 1024-residue protein fragments can process full-length proteins without retraining. The linear distance penalty maps naturally to the biological intuition that nearby residues in a sequence interact more strongly.

Connection to Modern Systems

RoPE vs ALiBi

In practice, RoPE has become more widely adopted (LLaMA 2/3, Mistral, Gemma) partly because extension methods like NTK-aware scaling and YaRN have largely solved its extrapolation weakness. ALiBi remains attractive for its simplicity and when extrapolation without any fine-tuning is paramount.

Flash Attention Compatibility

ALiBi is naturally compatible with Flash Attention (Chapter 13). The bias matrix can be computed on-the-fly within each tile of the tiled attention computation, without materializing the full $N \times N$ bias matrix in GPU HBM. Since the bias at position $(i,j)$ is just $-m_h \times |i-j|$, it can be computed from the tile indices alone. Flash Attention 2 natively supports ALiBi-style additive biases.

KV-Cache and Inference

ALiBi does not modify Q, K, or V vectors, so it has zero impact on KV-cache size. Compare with RoPE, where rotated K vectors must be stored in the cache (though this has the same size as unrotated vectors). ALiBi's bias is applied after the Q·K dot product, meaning the cache stores only the original unmodified K and V vectors. The bias is recomputed at each decoding step from the position indices — a negligible cost.

ALiBi composes cleanly with MQA (Chapter 5) and GQA (Chapter 6). The bias is applied per-head to the attention scores regardless of whether keys and values are shared.

Complexity Analysis

The computational overhead of ALiBi is negligible: adding a scalar to each attention score costs $O(N^2)$ per head, which is dominated by the $O(N^2 \cdot d_k)$ cost of the dot products themselves. In practice, the bias can be fused into the attention kernel with zero additional memory passes.

Python Implementation

The class below implements ALiBi from scratch. The critical method is apply_alibi_bias(), which subtracts the linear distance penalty from scaled scores. Compare this with the standard attention class from Chapter 1 — only one additional step is needed.

     0  1  2  3  4
0 [  0  1  2  3  4 ]
1 [  1  0  1  2  3 ]
2 [  2  1  0  1  2 ]
3 [  3  2  1  0  1 ]
4 [  4  3  2  1  0 ]

        The    cat    sat     on    mat
The   0.00  -0.50  -1.00  -1.50  -2.00
cat  -0.50   0.00  -0.50  -1.00  -1.50
sat  -1.00  -0.50   0.00  -0.50  -1.00
on   -1.50  -1.00  -0.50   0.00  -0.50
mat  -2.00  -1.50  -1.00  -0.50   0.00

        The    cat    sat     on    mat
The   0.00  -0.25  -0.50  -0.75  -1.00
cat  -0.25   0.00  -0.25  -0.50  -0.75
sat  -0.50  -0.25   0.00  -0.25  -0.50
on   -0.75  -0.50  -0.25   0.00  -0.25
mat  -1.00  -0.75  -0.50  -0.25   0.00

     The  cat  sat   on  mat
The  0.0  1.0  1.0  0.0  1.0
cat  2.0  0.0  2.0  0.0  0.0
sat  1.0  1.0  2.0  0.0  1.0
on   0.0  0.0  0.0  0.0  0.0
mat  0.0  1.0  1.0  0.0  1.0

       The     cat     sat      on     mat
The  0.0000  0.7071  0.7071  0.0000  0.7071
cat  1.4142  0.0000  1.4142  0.0000  0.0000
sat  0.7071  0.7071  1.4142  0.0000  0.7071
on   0.0000  0.0000  0.0000  0.0000  0.0000
mat  0.0000  0.7071  0.7071  0.0000  0.7071

        The    cat    sat     on    mat
The   0.00  -0.50  -1.00  -1.50  -2.00
cat  -0.50   0.00  -0.50  -1.00  -1.50
sat  -1.00  -0.50   0.00  -0.50  -1.00
on   -1.50  -1.00  -0.50   0.00  -0.50
mat  -2.00  -1.50  -1.00  -0.50   0.00

        The     cat     sat      on     mat
The   0.000   0.207  -0.293  -1.500  -1.293
cat   0.914   0.000   0.914  -1.000  -1.500
sat  -0.293   0.207   1.414  -0.500  -0.293
on   -1.500  -1.000  -0.500   0.000  -0.500
mat  -2.000  -0.793  -0.293  -0.500   0.707

        The      cat      sat       on      mat
The  0.2879   0.3541   0.2148   0.0642   0.0790
cat  0.3791   0.1520   0.3791   0.0559   0.0339
sat  0.1003   0.1653   0.5527   0.0815   0.1003
on   0.0796   0.1312   0.2163   0.3566   0.2163
mat  0.0341   0.1140   0.1880   0.1528   0.5110

       d0       d1
The  0.3274   0.3936
cat  0.3961   0.1689
sat  0.1504   0.2154
on   0.1877   0.2393
mat  0.2896   0.3695

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

        d0       d1       d2       d3
The  0.3274   0.3936   0.1861   0.2613
cat  0.3961   0.1689   0.2120   0.2977
sat  0.1504   0.2154   0.2544   0.3573
on   0.1877   0.2393   0.1811   0.5662
mat  0.2896   0.3695   0.2731   0.4746

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

        d0       d1       d2       d3
The  0.3274   0.3936   0.1861   0.2613
cat  0.3961   0.1689   0.2120   0.2977
sat  0.1504   0.2154   0.2544   0.3573
on   0.1877   0.2393   0.1811   0.5662
mat  0.2896   0.3695   0.2731   0.4746

1import numpy as np
2import math
3
4class ALiBiAttention:
5    """
6    ALiBi — Attention with Linear Biases (Press et al., ICLR 2022)
7
8    Eliminates positional embeddings entirely.
9    Instead, subtracts a head-specific linear distance penalty
10    from attention scores before softmax.
11
12    slope_h = 1 / 2^h  (geometric schedule, no learned params)
13    """
14
15    def __init__(self, d_model: int, n_heads: int):
16        self.d_model = d_model
17        self.n_heads = n_heads
18        self.d_k = d_model // n_heads
19        self.scale = math.sqrt(self.d_k)
20        self.slopes = self._compute_slopes(n_heads)
21
22    def _compute_slopes(self, n_heads: int) -> list:
23        """Geometric slope schedule: m_h = 1/2^h."""
24        return [1.0 / (2 ** (h + 1)) for h in range(n_heads)]
25
26    def _build_alibi_bias(self, N: int, slope: float) -> np.ndarray:
27        """Build the N×N distance penalty matrix."""
28        i_idx = np.arange(N)[:, None]
29        j_idx = np.arange(N)[None, :]
30        return -slope * np.abs(i_idx - j_idx)
31
32    def _softmax(self, x: np.ndarray) -> np.ndarray:
33        """Numerically stable softmax along last axis."""
34        x_shifted = x - np.max(x, axis=-1, keepdims=True)
35        exp_x = np.exp(x_shifted)
36        return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
37
38    def compute_scores(self, Qh: np.ndarray, Kh: np.ndarray):
39        """Raw dot-product scores: Qh @ Kh^T."""
40        return Qh @ Kh.T
41
42    def scale_scores(self, scores: np.ndarray) -> np.ndarray:
43        """Divide by sqrt(d_k) to control variance."""
44        return scores / self.scale
45
46    def apply_alibi_bias(self, scaled: np.ndarray, slope: float):
47        """Add the linear distance penalty to scores."""
48        N = scaled.shape[0]
49        bias = self._build_alibi_bias(N, slope)
50        return scaled + bias
51
52    def compute_weights(self, alibi_scores: np.ndarray):
53        """Apply softmax to get attention weights."""
54        return self._softmax(alibi_scores)
55
56    def compute_output(self, weights: np.ndarray, Vh: np.ndarray):
57        """Weighted sum of value vectors."""
58        return weights @ Vh
59
60    def forward(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray):
61        """
62        Full forward pass.
63
64        Args:
65            Q: Query matrix  (N, d_model)
66            K: Key matrix    (N, d_model)
67            V: Value matrix  (N, d_model)
68
69        Returns:
70            output:      Concatenated head outputs  (N, d_model)
71            all_weights: List of weight matrices per head
72        """
73        head_outputs = []
74        all_weights = []
75
76        for h in range(self.n_heads):
77            s = h * self.d_k
78            e = s + self.d_k
79            Qh, Kh, Vh = Q[:, s:e], K[:, s:e], V[:, s:e]
80
81            raw_scores = self.compute_scores(Qh, Kh)
82            scaled = self.scale_scores(raw_scores)
83            alibi_scores = self.apply_alibi_bias(scaled, self.slopes[h])
84            weights = self.compute_weights(alibi_scores)
85            output = self.compute_output(weights, Vh)
86
87            head_outputs.append(output)
88            all_weights.append(weights)
89
90        return np.hstack(head_outputs), all_weights
91
92    def explain(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray,
93                tokens: list, query_idx: int = 0):
94        """Print a detailed trace for a specific query token."""
95        token = tokens[query_idx]
96        print(f"\n=== ALiBi trace for '{token}' (row {query_idx}) ===")
97
98        for h in range(self.n_heads):
99            s = h * self.d_k
100            e = s + self.d_k
101            Qh, Kh, Vh = Q[:, s:e], K[:, s:e], V[:, s:e]
102
103            raw = self.compute_scores(Qh, Kh)
104            scaled = self.scale_scores(raw)
105            alibi = self.apply_alibi_bias(scaled, self.slopes[h])
106            w = self.compute_weights(alibi)
107            out = self.compute_output(w, Vh)
108
109            print(f"\n--- Head {h+1} (slope={self.slopes[h]}) ---")
110            print(f"Q_h{h+1}[{token}] = {Qh[query_idx]}")
111            for j, t in enumerate(tokens):
112                dist = abs(query_idx - j)
113                penalty = self.slopes[h] * dist
114                print(f"  score[{t}] = {scaled[query_idx,j]:.4f}"
115                      f"  penalty = -{penalty:.4f}"
116                      f"  alibi = {alibi[query_idx,j]:.4f}")
117            print(f"  softmax = {np.round(w[query_idx], 4)}")
118            print(f"  output  = {np.round(out[query_idx], 4)}")
119
120
121# ── Shared Example (used in every chapter) ──
122tokens = ["The", "cat", "sat", "on", "mat"]
123
124Q = np.array([
125    [1.0, 0.0, 1.0, 0.0],   # The
126    [0.0, 2.0, 0.0, 1.0],   # cat
127    [1.0, 1.0, 1.0, 0.0],   # sat
128    [0.0, 0.0, 1.0, 1.0],   # on
129    [1.0, 0.0, 0.0, 1.0],   # mat
130])
131
132K = np.array([
133    [0.0, 1.0, 0.0, 1.0],   # The
134    [1.0, 0.0, 1.0, 0.0],   # cat
135    [1.0, 1.0, 0.0, 0.0],   # sat
136    [0.0, 0.0, 1.0, 1.0],   # on
137    [1.0, 0.0, 0.5, 0.5],   # mat
138])
139
140V = np.array([
141    [1.0, 0.0, 0.0, 0.0],   # The
142    [0.0, 1.0, 0.0, 0.0],   # cat
143    [0.0, 0.0, 1.0, 0.0],   # sat
144    [0.0, 0.0, 0.0, 1.0],   # on
145    [0.5, 0.5, 0.5, 0.5],   # mat
146])
147
148# ── Run ──
149alibi = ALiBiAttention(d_model=4, n_heads=2)
150output, all_weights = alibi.forward(Q, K, V)
151
152print("ALiBi Output (5x4):")
153print(np.round(output, 4))
154
155print("\nHead 1 Weights (5x5):")
156print(np.round(all_weights[0], 4))
157
158print("\nHead 2 Weights (5x5):")
159print(np.round(all_weights[1], 4))
160
161# Detailed trace for "The" and "cat"
162alibi.explain(Q, K, V, tokens, query_idx=0)
163alibi.explain(Q, K, V, tokens, query_idx=1)

PyTorch Implementation

The PyTorch version supports GPU acceleration and automatic differentiation. The key difference from the NumPy version: slopes are registered as a buffer (not a parameter), ensuring they move to GPU with the model but are never updated by the optimizer. The bias tensor is built for all heads at once using broadcasting: slopes $(H,) \rightarrow (H,1,1)$ times distances $(N,N) \rightarrow (H,N,N)$.

     0  1  2  3  4
0 [  0  1  2  3  4 ]
1 [  1  0  1  2  3 ]
2 [  2  1  0  1  2 ]
3 [  3  2  1  0  1 ]
4 [  4  3  2  1  0 ]

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import math
5
6class ALiBiAttention(nn.Module):
7    """
8    ALiBi — Attention with Linear Biases (Press et al., ICLR 2022)
9    PyTorch implementation with GPU support.
10
11    Subtracts head-specific linear distance penalty from scores.
12    Zero additional learned parameters.
13    """
14
15    def __init__(self, d_model: int, n_heads: int):
16        super().__init__()
17        self.d_model = d_model
18        self.n_heads = n_heads
19        self.d_k = d_model // n_heads
20        self.scale = math.sqrt(self.d_k)
21
22        # Register slopes as a non-learnable buffer
23        slopes = torch.tensor(
24            [1.0 / (2 ** (h + 1)) for h in range(n_heads)]
25        )
26        self.register_buffer("slopes", slopes)
27
28    def _build_alibi_bias(
29        self, N: int, device: torch.device
30    ) -> torch.Tensor:
31        """Build (n_heads, N, N) ALiBi bias tensor."""
32        positions = torch.arange(N, device=device)
33        dist = torch.abs(
34            positions.unsqueeze(1) - positions.unsqueeze(0)
35        ).float()
36        # slopes: (n_heads,) -> (n_heads, 1, 1) for broadcast
37        bias = -self.slopes.unsqueeze(1).unsqueeze(2) * dist
38        return bias
39
40    def forward(
41        self,
42        Q: torch.Tensor,
43        K: torch.Tensor,
44        V: torch.Tensor,
45        mask: torch.Tensor | None = None,
46    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
47        """
48        Args:
49            Q: (N, d_model)  — query matrix
50            K: (N, d_model)  — key matrix
51            V: (N, d_model)  — value matrix
52            mask: Optional (N, N) boolean mask
53
54        Returns:
55            output:      (N, d_model) — concatenated heads
56            all_weights: list of (N, N) per-head weights
57        """
58        N = Q.size(0)
59        device = Q.device
60
61        # Build ALiBi bias for all heads at once
62        alibi_bias = self._build_alibi_bias(N, device)
63
64        head_outputs = []
65        all_weights = []
66
67        for h in range(self.n_heads):
68            s = h * self.d_k
69            e = s + self.d_k
70            Qh = Q[:, s:e]
71            Kh = K[:, s:e]
72            Vh = V[:, s:e]
73
74            # Scaled dot-product
75            scores = torch.matmul(Qh, Kh.transpose(-2, -1))
76            scores = scores / self.scale
77
78            # Add ALiBi bias (the only difference from standard)
79            scores = scores + alibi_bias[h]
80
81            if mask is not None:
82                scores = scores.masked_fill(mask, float("-inf"))
83
84            weights = F.softmax(scores, dim=-1)
85            output = torch.matmul(weights, Vh)
86
87            head_outputs.append(output)
88            all_weights.append(weights)
89
90        return torch.cat(head_outputs, dim=-1), all_weights
91
92
93# ── Shared Example ──
94tokens = ["The", "cat", "sat", "on", "mat"]
95
96Q = torch.tensor([
97    [1.0, 0.0, 1.0, 0.0],
98    [0.0, 2.0, 0.0, 1.0],
99    [1.0, 1.0, 1.0, 0.0],
100    [0.0, 0.0, 1.0, 1.0],
101    [1.0, 0.0, 0.0, 1.0],
102])
103
104K = torch.tensor([
105    [0.0, 1.0, 0.0, 1.0],
106    [1.0, 0.0, 1.0, 0.0],
107    [1.0, 1.0, 0.0, 0.0],
108    [0.0, 0.0, 1.0, 1.0],
109    [1.0, 0.0, 0.5, 0.5],
110])
111
112V = torch.tensor([
113    [1.0, 0.0, 0.0, 0.0],
114    [0.0, 1.0, 0.0, 0.0],
115    [0.0, 0.0, 1.0, 0.0],
116    [0.0, 0.0, 0.0, 1.0],
117    [0.5, 0.5, 0.5, 0.5],
118])
119
120# ── Run ──
121alibi = ALiBiAttention(d_model=4, n_heads=2)
122output, all_weights = alibi(Q, K, V)
123
124print("ALiBi Output (5x4):")
125print(output.round(decimals=4))
126
127print("\nHead 1 Weights:")
128print(all_weights[0].round(decimals=4))
129
130print("\nHead 2 Weights:")
131print(all_weights[1].round(decimals=4))
132
133# ── GPU: just move tensors ──
134if torch.cuda.is_available():
135    alibi_gpu = alibi.cuda()
136    Q_gpu, K_gpu, V_gpu = Q.cuda(), K.cuda(), V.cuda()
137    out_gpu, _ = alibi_gpu(Q_gpu, K_gpu, V_gpu)
138    print("GPU matches CPU?",
139          torch.allclose(output, out_gpu.cpu(), atol=1e-4))

Key Takeaways

1. ALiBi replaces positional embeddings with a linear distance penalty subtracted from attention scores before softmax. This is the simplest possible position encoding: zero parameters, zero trigonometric functions, just $-m_h \times |i - j|$.
2. Different heads use different slopes via a geometric schedule $m_h = 1/2^h$. This creates a multi-resolution view: steep slopes for local syntax, gentle slopes for long-range dependencies.
3. Extrapolation is the primary advantage. Because the bias is a deterministic function of distance (not a learned representation), models trained at length $L$ generalize to $2L$ or beyond without fine-tuning.
4. Diagonal dominance emerges naturally. Even when content scores are uniform or zero, ALiBi creates meaningful attention patterns where self-attention and nearby tokens dominate.
5. ALiBi composes with all other optimizations. It works with MQA/GQA (KV-sharing), Flash Attention (tiled computation), and causal masking. The bias adds negligible overhead.

Exercises

Exercise 1: Compute for "sat"

Using the shared Q, K, V matrices, compute the ALiBi attention weights for “sat” (row 2) in both heads. Verify that Head 1 gives $[0.1003,\; 0.1653,\; 0.5527,\; 0.0815,\; 0.1003]$. Why does "sat" strongly self-attend (55.3%) in Head 1? Consider both the content score and the distance penalty.

Exercise 2: Slope = 0 Reduces to Standard Attention

Prove mathematically that when $m_h = 0$ for all heads, ALiBi produces identical results to standard scaled dot-product attention (Chapter 1). Then verify numerically using the interactive visualizer above (set slope to 0).

Exercise 3: ALiBi + Causal Mask

In autoregressive models (Chapter 3), causal masking sets future positions to $-\infty$. How does this interact with ALiBi? Write the combined score formula and explain whether the ALiBi penalty is redundant for future positions. Hint: $-\infty + (-m_h \times d) = -\infty$.

Exercise 4: Comparing Slope Schedules

Press et al. used $m_h = 1/2^h$. Propose two alternative schedules: (a) linear: $m_h = (H - h + 1) / H$ and (b) constant: $m_h = 0.25$ for all heads. For each, compute the attention weights for "The" across all 5 tokens using Head 1. What is lost when all heads have the same slope?

Exercise 5: Extrapolation Experiment

Extend the shared example to 10 tokens: ["The", "cat", "sat", "on", "mat", "and", "the", "dog", "ran", "home"]. Invent reasonable Q, K, V vectors for the new tokens. Compute ALiBi attention for the original 5 tokens and verify that their weights are nearly unchanged despite the longer sequence. This demonstrates ALiBi's extrapolation: adding more tokens barely affects existing attention patterns.

References

1. Press, O., Smith, N. A., & Lewis, M. (2022). “Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation.” International Conference on Learning Representations (ICLR).
2. Scao, T. L., et al. (2022). “BLOOM: A 176B-Parameter Open-Access Multilingual Language Model.” arXiv:2211.05100.
3. MosaicML (2023). “Introducing MPT-7B: A New Standard for Open-Source, Commercially Usable LLMs.” MosaicML Blog.
4. Su, J., et al. (2021). “RoFormer: Enhanced Transformer with Rotary Position Embedding.” arXiv:2104.09864.
5. Dao, T., et al. (2022). “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” arXiv:2205.14135.
6. Vaswani, A., et al. (2017). “Attention Is All You Need.” Advances in Neural Information Processing Systems (NeurIPS), 30.