Mistral 7B, annotated

The shape of the model, and GQA · original

Everyday picture

A model's shape is a short list of numbers, like a building's floor plan: how many floors, how wide each floor, how many rooms. From that list alone you can work out how much memory the model needs to remember each token of a conversation.

Tiny example

Every token leaves behind one key and one value per key/value head in every layer. With 8 key/value heads of 128 numbers, 32 layers, and 16-bit numbers (2 bytes, an assumption: the paper does not state the precision), that is 2 × 32 × 8 × 128 × 2 = 131,072 bytes, 128 KiB per token. If each of the 32 query heads kept its own keys and values, it would be 512 KiB.

In words: “a key and a value, per key/value head, per layer, at so many bytes per number; and each key/value head is shared by g query heads.”

With the numbers: 2 × 32 × 8 × 128 × 2 = 131,072 bytes; with H_kv = H = 32 it would be 524,288. The group size is g = 32 / 8 = 4, and the cache is 4 times smaller.

In Python:

L, H, H_kv, d_h, b = 32, 32, 8, 128, 2
per_token = 2 * L * H_kv * d_h * b
per_token  # → 131072
# in KiB
per_token / 1024  # → 128.0
# if every query head had its own keys and values
2 * L * H * d_h * b  # → 524288
# query heads per key/value head
g = H // H_kv
g  # → 4

Why it matters

Grouped-query attention sits between the two designs of the multi-query attention paper: one key/value head per query head (multi-head) and one for all of them (multi-query). The paper cites Ainslie et al. (2023) for it and uses it, in its words, to accelerate inference and allow higher batch sizes. The primer's MultiHeadAttention is grouped-query attention whenever n_kv_heads is smaller than n_heads, and kv_cache_bytes_per_token is the formula above.

Sliding window attention · original

“SWA exploits the stacked layers of a transformer to attend information beyond the window size W.”Jiang et al. (2023), §2

Everyday picture

A rumour passes down a line of people, but each person only talks to the few people nearest them. Nobody speaks to the far end of the line directly, yet the rumour gets there, one short hop at a time. In sliding-window attention each token, in each layer, attends only to a window of nearby earlier tokens. Because each layer reads the previous layer's results, which already mixed in their own neighbours, information hops further back with every layer.

Tiny example

Twelve tokens: “It was the best of times it was the worst of times”, numbered 0 to 11. With a window of 3, as in the paper's Figure 1, token 11 (“times”) reads tokens 9, 10 and 11 in each layer. After one layer it has heard from token 9; token 9 had itself heard from token 7 in the layer below; so after two layers token 11 reaches token 7, and after three layers token 5. Each layer adds 2 positions of reach.

Reading it: the upper grid is one layer's mask, like the middle panel of the paper's figure: a row for each token doing the looking, a column for each token being looked at, filled where a score is computed; row 11, the token followed below, is drawn darker. It is a band of width W along the diagonal, never above it (no token sees the future). The lower picture stacks the layers, with the input at the bottom and the last layer on top. The filled cells in each row are the tokens that can still influence token 11 at the top: at every layer down, the band spreads W − 1 positions further back, so the bottom row shows everything token 11 can hear. Slide the window up and each layer's band widens; add layers and the reach grows in steps. With W = 3 and three layers, token 11 hears from tokens 5 to 11, and tokens 0 to 4 are out of reach, yet no single layer ever looked more than 2 positions back.

In words: “the state at position i in layer k reads the previous layer's states from W positions back up to itself; stacking k such layers lets information travel up to W times k positions.”

With the numbers: Mistral 7B has W = 4,096 and 32 layers, so the span is 4,096 × 32 = 131,072 tokens, the paper's “approximately 131K”. The widget, the paper's figures and the primer's lesson count the window as W positions including the token itself, which reaches W − 1 back per layer: 32 × 4,095 = 131,040. The paper's text counts from i − W to i, which is W + 1 positions. Either way the answer is about 131 thousand tokens from a window of 4 thousand.

In Python:

W, k = 4096, 32
# the paper's span: W × k
W * k  # → 131072
# counting the window as W positions including the token itself
k * (W - 1)  # → 131040
# the tiny example: W = 3 (itself and 2 back), token 11, three layers
w, layers, i = 3, 3, 11
i - layers * (w - 1)  # → 5

What it saves. The paper reports that at a sequence length of 16K with W = 4,096, changes it made to FlashAttention and xFormers give a 2x speed-up over a vanilla attention baseline. Counting pairs gives a sense of why (this page's own arithmetic, not a measurement): a full causal mask over 16,384 tokens scores 134 million query-key pairs per head, a window of 4,096 scores 58.7 million, 2.3 times fewer.

In Python:

n, W = 16384, 4096
# full causal attention: token i reads i + 1 tokens
full = sum(i + 1 for i in range(n))
# sliding window: token i reads at most W tokens
window = sum(min(i + 1, W) for i in range(n))
full, window  # → (134225920, 58722304)
round(full / window, 2)  # → 2.29

Why it matters

The window turns attention's cost from growing with the square of the length to growing in proportion to it, while stacked layers keep a long, if indirect, reach. The price is that a fact from far back must survive several hops to arrive. The lesson builds the band with sliding_window_mask, counts it with pairs_computed, and checks the reach by applying the mask layer after layer in receptive_field, the same computation as the widget. The idea itself comes from earlier work the paper cites, Sparse Transformers and Longformer.

Rolling buffer cache · original

“On a sequence length of 32k tokens, this reduces the cache memory usage by 8x, without impacting the model quality.”Jiang et al. (2023), §2

Everyday picture

A whiteboard with room for exactly W notes. When it is full, the next note is written over the oldest one. Since no layer ever looks more than W tokens back, the notes that get erased are exactly the ones nobody will read again.

Tiny example

With W = 4, token i goes into slot i mod 4, the remainder after dividing by 4. Tokens 0 to 3 fill slots 0 to 3. Token 4 (“of”) goes into slot 0, over token 0 (“It”); token 9 (“worst”) goes into slot 1, over token 5 (“times”). The board never grows.

Reading it: press Next token to generate one more token. The four boxes are the whole cache. The slot just written is filled solid, as the latest token is coloured in the paper's figure, and the row of words below shows the text so far, with the tokens still in the cache underlined. From token 4 on, each new token lands in slot i mod 4 and erases the token exactly four positions earlier, which has just left the window. The cache never holds more than four tokens, however long the text grows.

In words: “token i's keys and values go into slot i mod W, so the cache holds the last W tokens once there are more than W of them, and all of them before that.”

With the numbers: W = 4: tokens 4 and 9 land in slots 0 and 1. For Mistral 7B at n = 32,768 tokens, the cache holds 4,096 of them instead of 32,768: 8 times less, the paper's figure. At 128 KiB per token (the GQA example above) that is 0.54 GB instead of 4.3 GB per sequence.

In Python:

W = 4
# slot(i) = i mod W
[i % W for i in range(12)]  # → [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
# Mistral 7B at 32k tokens
W, n = 4096, 32768
n / min(n, W)  # → 8.0
per_token = 131072
# gigabytes cached, with and without the buffer
round(min(n, W) * per_token / 1e9, 2), round(n * per_token / 1e9, 2)  # → (0.54, 4.29)

Why it matters

This is what makes the window pay off in memory, not only in compute: the cache has a fixed size, so a server knows in advance how much memory a conversation can ever need. The lesson's sliding_window_decode generates token by token with exactly this buffer and checks that its outputs match the masked attention, and cache_bytes_by_method plots the flat line it produces against every other design. One detail on the page of the paper: its text says Figure 2 shows W = 3, while the figure and its caption use W = 4.

Pre-fill and chunking · original

Everyday picture

The model writes its answer one token at a time, but the prompt is known in full before it starts, so its keys and values can be computed in advance: the prefill. A very long prompt would need a very large block of attention scores all at once, so the paper cuts it into pieces and pre-fills them one after another, like reading a long letter a page at a time while keeping the last page in mind.

Tiny example

The 12-token sentence is split into three chunks of 4 (the paper chooses the chunk size equal to the window, here 4): “It was the best”, “of times it was”, “the worst of times”. When the third chunk is processed, its tokens attend to themselves with an ordinary causal mask, to the second chunk (now sitting in the cache) through the sliding window, and not at all to the first chunk, which is out of reach.

Reading it: each row is one token of the chunk being pre-filled, each column one token of the whole prompt, and a filled cell (marked 1) is a score that gets computed. The dashed lines split the columns into the past (chunks already dropped from the cache), the cache (the previous chunk) and the current chunk. For chunk 3 the current block is a lower triangle (causal), the cache block is an upper triangle (the window reaches into the previous chunk less and less as you go down), and the past block is empty. Together each row has exactly 4 filled cells: the same window as always. Chunk 1 has no cache to look at, and chunk 2 has no past.

In words: “each of the C tokens in a chunk is scored against at most the W cached tokens plus the C tokens of the chunk.” (This bound is this page's accounting; the paper draws it rather than writing it.)

With the numbers: in the tiny example, C = W = 4: at most 4 × 8 = 32 cells, of which the figure fills 16 (4 per row). For a 16,384-token prompt with C = W = 4,096, the largest block of scores is 4,096 × 8,192 = 33.5 million per head, against 16,384² = 268 million for the whole prompt at once: 8 times smaller.

In Python:

def block(C, W):
    # C queries against W cached keys plus the C keys of the chunk
    return C * (W + C)
block(4, 4)  # → 32
n, C, W = 16384, 4096, 4096
block(C, W), n * n  # → (33554432, 268435456)
n * n // block(C, W)  # → 8

Why it matters

Chunking caps the size of every attention computation during prefill, and with the rolling buffer it caps the cache too, so a prompt of any length fits in fixed memory. With the chunk size equal to the window, the rolling buffer holds exactly the previous chunk when the next one arrives, which is all the window can reach. The inference lesson explains why prefill and decode behave so differently in the first place.

The benchmarks · original

Everyday picture

To compare two students fairly, you give them the same exams, marked the same way. The authors re-ran every model, their own and Llama's, through their own evaluation pipeline, rather than copying numbers from other papers. The exams are grouped into categories:

• Commonsense reasoning (0-shot): HellaSwag, Winogrande, PIQA, SIQA, OpenbookQA, ARC-Easy, ARC-Challenge, CommonsenseQA.
• World knowledge (5-shot): NaturalQuestions, TriviaQA.
• Reading comprehension (0-shot): BoolQ, QuAC.
• Maths: GSM8K (8-shot, maj@8) and MATH (4-shot, maj@4).
• Code: HumanEval (0-shot) and MBPP (3-shot). The HumanEval companion explains how code is scored.
• Aggregated: MMLU (5-shot), BBH (3-shot), and AGI Eval (3 to 5 shots, English multiple-choice only).

“5-shot” means five worked examples are placed in the prompt before each question; “maj@8” means eight answers are sampled and the most common one counts.

All numbers are accuracies in percent. Pick a benchmark to compare the four models:

Reading it: each bar is one model's accuracy on the chosen benchmark, on an axis from 0 to 100%, with its name on the left and its score on the right; Mistral 7B's bar is solid and the others are striped. On most benchmarks Mistral's bar is the longest, and on MMLU it leads Llama 2 13B by 4.5 points. Choose HumanEval or MBPP and Code Llama 7B, a model fine-tuned for code, comes out ahead, with Mistral close behind; Code Llama trails well behind on everything that is not code. Choose NQ (NaturalQuestions, a knowledge test) and Mistral's 28.8 sits just below Llama 2 13B's 29.0, the one column where the caption's “outperforms Llama 2 13B on all metrics” does not hold in its own table.

Why it matters

The paper lists two places where its evaluation differs from Llama 2's own: MBPP uses the hand-verified subset, and TriviaQA is asked without Wikipedia passages. Scores from different pipelines are not directly comparable, which is why re-running every model matters. The benchmarks lesson covers what such numbers can and cannot tell you.

Size and efficiency · original

“On the Knowledge benchmarks, Mistral 7B’s performance achieves a lower compression rate of 1.9x, which is likely due to its limited parameter count that restricts the amount of knowledge it can store.”Jiang et al. (2023), §3

Everyday picture

Line up the Llama 2 family by size (7B, 13B, 70B) and draw the curve of their scores. Put Mistral 7B's score on that curve and read off where a Llama 2 model would have to sit to score the same. That imaginary Llama's size is Mistral's equivalent model size.

Tiny example

The paper's Figure 5 does this for four categories. On MMLU, Mistral's 60.1 matches a Llama 2 of about 23B, 3.3 times its size. On reasoning it matches about 38B (5.4 times), on comprehension about 21B (3 times), and on knowledge only 13B (1.9 times).

In words: “how many times bigger a Llama 2 model would have to be to match Mistral 7B.”

With the numbers: 23 / 7 = 3.3; 38 / 7 = 5.4; 21 / 7 = 3.0; 13 / 7 = 1.9, the ratios printed on the paper's figure.

In Python:

N_mistral = 7
equivalent = {"MMLU": 23, "reasoning": 38, "comprehension": 21, "knowledge": 13}
{k: round(v / N_mistral, 1) for k, v in equivalent.items()}  # → {'MMLU': 3.3, 'reasoning': 5.4, 'comprehension': 3.0, 'knowledge': 1.9}

Why it matters

The pattern is the interesting part. Skills (reasoning, reading, maths) compress well into a small model; facts do not, because every fact has to be stored somewhere in the weights. This is also why a small model paired with retrieval, which looks facts up instead of memorising them, is a common design (the RAG lesson builds one).

5.1 A system prompt to enforce guardrails · original

Everyday picture

Before a new employee meets customers, a manager gives them a few sentences of standing instructions. A system prompt is that briefing for a model, and a guardrail is any rule that keeps its output in bounds. The paper's recommended briefing is four sentences:

“Always assist with care, respect, and truth. Respond with utmost utility yet securely. Avoid harmful, unethical, prejudiced, or negative content. Ensure replies promote fairness and positivity.”Jiang et al. (2023), §5.1

Tiny example

With its system prompt, the model declined all of a set of 175 unsafe prompts, at a cost of 0.26 MT-Bench points (about 4%) in helpfulness; Llama 2's system prompt, given to the same model, cost 0.46. The paper's example of the difference: asked how to kill a Linux process, Mistral 7B Instruct with its prompt explains the kill command, while Llama 2 13B Chat with its own prompt refuses. With no system prompt, both models answer correctly.

Why it matters

Guardrails move a model along a trade-off between usefulness and caution; refusing a harmless question about operating systems is the cost of leaning too far. The guardrails lesson builds input and output checks around a model rather than relying on the prompt alone.

5.2 Content moderation with self-reflection · original

Everyday picture

The same model can be asked to act as a moderator: given a prompt, or an answer, say whether it is acceptable or falls into one of three categories (illegal activities, hateful or violent content, unqualified legal, medical or financial advice). The paper calls this self-reflection.

Tiny example

On the authors' balanced set of adversarial and ordinary prompts, counting acceptable prompts as positives, it reached a precision of 99.4% and a recall of 95.6%. The paper does not give counts, so take 1,000 acceptable prompts as an illustration: recall 95.6% means 956 of them are let through (44 wrongly flagged), and precision 99.4% means that among everything let through, about 6 were not acceptable.

In words: “precision is the share of prompts let through that deserved it; recall is the share of acceptable prompts that were let through.”

With the numbers (illustrative counts): TP = 956, FN = 44, FP = 6: precision = 956 / 962 = 99.4%, recall = 956 / 1,000 = 95.6%.

In Python:

# illustrative counts matching the paper's rates
TP, FN, FP = 956, 44, 6
round(100 * TP / (TP + FP), 1)  # → 99.4
round(100 * TP / (TP + FN), 1)  # → 95.6

Why it matters

Because acceptable prompts are the positives, high precision means very little harmful content slips through as “acceptable”, and the 95.6% recall means about one acceptable prompt in 23 is wrongly flagged. The paper notes that users can choose which categories to filter for their own application. The metrics lesson explains precision and recall and the trade-off between them.