The Transformer Block

in progressThis lesson is still being built - content, demos, and checkpoints may change before release.

Transformer block

Loading interactive demo…

Six tokens through one block. Each row of the residual stream starts as an embedding plus a position code, is nudged by what every head attends to, normalized, nudged again by a feed-forward network, and read off as a distribution over the vocabulary. The weights are seeded, not trained: the flow is the lesson. Take the tour, or press and explore. Legend:

The panel beside the stage is the real TransformerBlock class from the published library executing in your browser, with the current line highlighted.

The residual stream on a diverging ramp: positive one way, negative the other, faint until written.
The token and stop being computed.
The tapped token; the sentence strip shades every token by its attention weight from that one.
The vocabulary strip, shaded by the tapped token's guess.

predicts is the tapped token's top word, residual norm the length of its row, attention entropy how spread its attention is, stage the stop in progress. Tap a token to inspect it.

The code panel shows bare executable code by default. The bar above it names the file (hover the name for the full published path); options holds the toggles: comments, the tracer instrumentation that produces this visualization, and verbose names for the terse identifiers (display only: the default is the byte-exact published file). The view switch adds plain-English narration: both interleaves it with the code, pseudo shows only the pseudocode. hide code collapses the panel; the rail it leaves behind brings it back.

Keys: Space play/pause · ←/→ step · Home/End jump · R reset · F fullscreen · G guided tour

The block. With $x$ the residual stream, one row per token,

h = \mathrm{LN}\big(x + \mathrm{MultiHead}(x)\big), \qquad y = \mathrm{LN}\big(h + \mathrm{FFN}(h)\big),

and each row of $y$ is read off as $p = \mathrm{softmax}(y\, W_{\text{out}} / T)$ .

Embedding and position. Row $i$ starts as $E[t_i] + \mathrm{PE}(i)$ , where the sinusoidal code is $\mathrm{PE}(i)_{2k} = \sin(i / 10000^{2k/d})$ and $\mathrm{PE}(i)_{2k+1} = \cos(i / 10000^{2k/d})$ . Attention treats tokens as a set, so position must be written into the numbers.

Multi-head attention. Each head runs the self-attention of the previous lesson over its own slice of the row, $d / \text{heads}$ wide; the head outputs are concatenated and mixed by $W_o$ . Two heads can look at different tokens for different reasons.

Residual and norm. The add stops compute $x + f(x)$ , so a sub-layer nudges a row rather than replacing it, and gradients have a straight path back through the block. Layer norm then maps a row to mean 0 and variance 1 (with $\gamma = 1$ , $\beta = 0$ here), so the residual norm returns to $\sqrt{d} = 2$ after every norm stop.

Feed-forward. $\mathrm{FFN}(h) = \max(0, h W_1)\, W_2$ , one row at a time, with a hidden width of 8: the only stop where tokens do not interact.

Invariants. Every probability row sums to 1; every normalized row has mean 0 and variance 1; with the position code off, identical tokens produce identical rows and identical guesses, since an unmasked block is equivariant under permuting its tokens. Edges: one head of width 4 versus two of width 2; the norm's epsilon keeps a constant row finite; a low temperature turns an untrained near-uniform guess into a confident wrong one.

Choices worth reading closely in the implementation (the code panel's comments and tracer checkboxes reveal the full published file):

Heads reuse SelfAttention untraced, one instance per head over a slice of the row; the block emits one attend per head and token with that row's weights, so the trace stays at 56 events.
residual and normalize are methods called twice, once after attention and once after the feed-forward, so each tag has one call site and the payload's stage says which pass it is.
No learned gain or shift in the norm and no biases in the feed-forward: the seeded weights are not trained, so the extra parameters would only be noise.
The position code is a zero vector when off, so the embed payload has the same shape either way and the sum is honest.
Every matrix comes from one seeded stream in a fixed order, so a seed names the whole block.

The 2017 paper stacked self-attention with three older ideas and removed the recurrence: the residual connection, which adds a sub-layer's result onto its input instead of replacing it; layer normalization, which keeps every row at a known scale; and a position code, since attention alone cannot tell the first token from the fifth. The block that resulted is the same in the encoder-only models that read text and the decoder-only models that write it; only the mask and the training differ. Every large language model since is this block repeated, made wider and deeper.

2015Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun - The residual network wins that year's ImageNet challenge: a skip connection lets a sub-layer learn a nudge, x plus f of x, which this block's two add stops copy.source
2016Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton - 'Layer Normalization' recentres and rescales across the features of one sample, the norm stops of this block.source
2017Ashish Vaswani and colleagues - 'Attention Is All You Need' publishes the encoder-decoder transformer, with sinusoidal position codes, on June 12.source
2018Alec Radford and colleagues at OpenAI - GPT-1, in 'Improving Language Understanding by Generative Pre-Training' on June 11: a decoder-only stack of these blocks.source
2018Jacob Devlin and colleagues at Google - BERT, in October: an encoder-only stack, bidirectional like this stage's unmasked block.source

Used today

Language models - Chat models are dozens of these blocks with thousands of numbers per row, a causal mask, and trained weights; the stops are the same six.
Vision and speech - Vision transformers embed image patches instead of words and run the identical block.
Protein and code - Amino-acid sequences and source code are token streams too, and the same block reads them.

heads

position

temperature1.00

sentence

predicts: ran 0.39
residual norm: 2.00
attention entropy: 2.12 bits
stage: predict

transformer-block.ts

import { Mulberry32 } from "../random";import type { Tracer } from "../trace";import { Relu } from "./activation";import { Kernels } from "./kernels";import { SelfAttention } from "./self-attention";export interface BlockResult {  /** The residual stream after the block, one row per token. */  stream: number[][];  /** attention[head][query][key]. */  attention: number[][][];  /** One probability row over the vocabulary per token. */  probabilities: number[][];}/** * One transformer block with a vocabulary head at every position: embed * each token and add its position code, let every head attend over its * slice of the stream, add the result back and normalize, push each row * through a small relu network, add and normalize again, and read a * distribution over the vocabulary off every row. Unmasked, so with the * position code off two identical tokens produce identical rows. The * weights are seeded, not trained: the lesson is the data flow. * * Time: O(n² · d + n · d · hidden) for n tokens. Space: O(n · d). */export class TransformerBlock {  readonly embeddings: number[][];  private readonly heads_: SelfAttention[];  private readonly wo: number[][];  private readonly w1: number[][];  private readonly w2: number[][];  private readonly wout: number[][];  private readonly relu = new Relu();  readonly headDim: number;  constructor(    readonly dim: number,    readonly heads: number,    readonly hidden: number,    readonly vocab: readonly string[],    seed = 1,    private readonly tracer?: Tracer,  ) {    if (!Number.isInteger(heads) || heads < 1 || dim % heads !== 0) throw new RangeError("the dimension must split evenly across the heads");    if (!Number.isInteger(hidden) || hidden < 1) throw new RangeError("the hidden width must be a whole number of at least one");    if (vocab.length < 2) throw new RangeError("a vocabulary needs at least two words");    this.headDim = dim / heads;    const random = new Mulberry32(seed);    const matrix = (rows: number, cols: number) => Array.from({ length: rows }, () => Array.from({ length: cols }, () => random.between(-1, 1)));    this.embeddings = matrix(vocab.length, dim);    this.heads_ = Array.from({ length: heads }, (_, head) => new SelfAttention(this.headDim, "seeded", seed * 101 + head + 1));    this.wo = matrix(dim, dim);    this.w1 = matrix(hidden, dim);    this.w2 = matrix(dim, hidden);    this.wout = matrix(vocab.length, dim);  }  private apply(m: readonly (readonly number[])[], v: readonly number[]): number[] {    return m.map((row) => row.reduce((s, weight, i) => s + weight * (v[i] ?? 0), 0));  }  /** The residual connection: the sub-layer's delta added back onto the row. */  private residual(stage: string, index: number, before: readonly number[], delta: readonly number[]): number[] {    const after = before.map((value, i) => value + delta[i]!);    // ▸ the sub-layer's result is added onto the row, not written over it    this.tracer?.emit("add", { stage, index, before: [...before], delta: [...delta], after: [...after] });    return after;  }  /** Layer norm: the row recentred and rescaled, with no learned gain or shift. */  private normalize(stage: string, index: number, row: readonly number[]): number[] {    const { out, mean, variance } = Kernels.layerNorm(row);    // ▸ the row is recentred to mean zero and rescaled to unit variance    this.tracer?.emit("norm", { stage, index, mean, variance, after: [...out] });    return out;  }  forward(tokens: readonly number[], positional: boolean, temperature = 1): BlockResult {    if (tokens.some((t) => !Number.isInteger(t) || t < 0 || t >= this.vocab.length)) throw new RangeError("every token must index the vocabulary");    let x: number[][] = [];    let residualStream: number[][] = [];    for (let index = 0; index < tokens.length; index++) {      const token = tokens[index]!;      const embedding = this.embeddings[token]!;      const position = positional ? Kernels.positionalEncoding(index, this.dim) : Array.from({ length: this.dim }, () => 0);      const sum = embedding.map((value, i) => value + position[i]!);      x.push(sum);      residualStream.push(sum);      // ▸ a token's vector plus a code for where it sits      this.tracer?.emit("embed", { index, token: this.vocab[token], embedding: [...embedding], position: [...position], sum: [...sum] });    }    const attention: number[][][] = [];    const h: number[][] = x.map(() => []);    const headOutputs: number[][] = residualStream.map(() => []);    for (let head = 0; head < this.heads; head++) {      const slice = x.map((row) => row.slice(head * this.headDim, (head + 1) * this.headDim));      const slice = residualStream.map((row) => row.slice(head * this.headDim, (head + 1) * this.headDim));      const { weights, outputs } = this.heads_[head]!.attend(slice);      attention.push(weights);      for (let index = 0; index < tokens.length; index++) {        h[index]!.push(...outputs[index]!);        headOutputs[index]!.push(...outputs[index]!);        // ▸ one head looks around the sentence from this token        this.tracer?.emit("attend", { head, index, weights: [...weights[index]!], output: [...outputs[index]!] });      }    }    x = x.map((row, index) => this.residual("attention", index, row, this.apply(this.wo, h[index]!)));    residualStream = residualStream.map((row, index) => this.residual("attention", index, row, this.apply(this.wo, headOutputs[index]!)));    x = x.map((row, index) => this.normalize("attention", index, row));    residualStream = residualStream.map((row, index) => this.normalize("attention", index, row));    const f: number[][] = [];    const feedForward: number[][] = [];    for (let index = 0; index < tokens.length; index++) {      const hiddenRow = this.apply(this.w1, x[index]!).map((value) => this.relu.apply(value));      const hiddenRow = this.apply(this.w1, residualStream[index]!).map((value) => this.relu.apply(value));      const output = this.apply(this.w2, hiddenRow);      f.push(output);      feedForward.push(output);      // ▸ a small relu network works on each row alone      this.tracer?.emit("ffn", { index, hidden: [...hiddenRow], output: [...output] });    }    x = x.map((row, index) => this.residual("ffn", index, row, f[index]!));    residualStream = residualStream.map((row, index) => this.residual("ffn", index, row, feedForward[index]!));    x = x.map((row, index) => this.normalize("ffn", index, row));    residualStream = residualStream.map((row, index) => this.normalize("ffn", index, row));    const probabilities: number[][] = [];    for (let index = 0; index < tokens.length; index++) {      const logits = this.apply(this.wout, x[index]!);      const logits = this.apply(this.wout, residualStream[index]!);      const p = Kernels.softmax(logits, temperature);      const top = p.indexOf(Math.max(...p));      probabilities.push(p);      // ▸ the row becomes a distribution over the vocabulary      this.tracer?.emit("predict", { index, logits: [...logits], probabilities: [...p], top: this.vocab[top], probability: p[top] });    }    return { stream: x, attention, probabilities };    return { stream: residualStream, attention, probabilities };  }}

top =: ran
probability =: 0.39

Every row is done: the inspected token predicts ran at 39 percent, and the attention rows average 2.11 bits of entropy.

import { Mulberry32 } from "../random";import type { Tracer } from "../trace";import { Relu } from "./activation";import { Kernels } from "./kernels";import { SelfAttention } from "./self-attention";export interface BlockResult { /** The residual stream after the block, one row per token. */ stream: number[][]; /** attention[head][query][key]. */ attention: number[][][]; /** One probability row over the vocabulary per token. */ probabilities: number[][];}/** * One transformer block with a vocabulary head at every position: embed * each token and add its position code, let every head attend over its * slice of the stream, add the result back and normalize, push each row * through a small relu network, add and normalize again, and read a * distribution over the vocabulary off every row. Unmasked, so with the * position code off two identical tokens produce identical rows. The * weights are seeded, not trained: the lesson is the data flow. * * Time: O(n² · d + n · d · hidden) for n tokens. Space: O(n · d). */export class TransformerBlock { readonly embeddings: number[][]; private readonly heads_: SelfAttention[]; private readonly wo: number[][]; private readonly w1: number[][]; private readonly w2: number[][]; private readonly wout: number[][]; private readonly relu = new Relu(); readonly headDim: number; constructor( readonly dim: number, readonly heads: number, readonly hidden: number, readonly vocab: readonly string[], seed = 1, private readonly tracer?: Tracer, ) { if (!Number.isInteger(heads) || heads < 1 || dim % heads !== 0) throw new RangeError("the dimension must split evenly across the heads"); if (!Number.isInteger(hidden) || hidden < 1) throw new RangeError("the hidden width must be a whole number of at least one"); if (vocab.length < 2) throw new RangeError("a vocabulary needs at least two words"); this.headDim = dim / heads; const random = new Mulberry32(seed); const matrix = (rows: number, cols: number) => Array.from({ length: rows }, () => Array.from({ length: cols }, () => random.between(-1, 1))); this.embeddings = matrix(vocab.length, dim); this.heads_ = Array.from({ length: heads }, (_, head) => new SelfAttention(this.headDim, "seeded", seed * 101 + head + 1)); this.wo = matrix(dim, dim); this.w1 = matrix(hidden, dim); this.w2 = matrix(dim, hidden); this.wout = matrix(vocab.length, dim); } private apply(m: readonly (readonly number[])[], v: readonly number[]): number[] { return m.map((row) => row.reduce((s, weight, i) => s + weight * (v[i] ?? 0), 0)); } /** The residual connection: the sub-layer's delta added back onto the row. */ private residual(stage: string, index: number, before: readonly number[], delta: readonly number[]): number[] { const after = before.map((value, i) => value + delta[i]!); // ▸ the sub-layer's result is added onto the row, not written over it this.tracer?.emit("add", { stage, index, before: [...before], delta: [...delta], after: [...after] }); return after; } /** Layer norm: the row recentred and rescaled, with no learned gain or shift. */ private normalize(stage: string, index: number, row: readonly number[]): number[] { const { out, mean, variance } = Kernels.layerNorm(row); // ▸ the row is recentred to mean zero and rescaled to unit variance this.tracer?.emit("norm", { stage, index, mean, variance, after: [...out] }); return out; } forward(tokens: readonly number[], positional: boolean, temperature = 1): BlockResult { if (tokens.some((t) => !Number.isInteger(t) || t < 0 || t >= this.vocab.length)) throw new RangeError("every token must index the vocabulary"); let x: number[][] = []; let residualStream: number[][] = []; for (let index = 0; index < tokens.length; index++) { const token = tokens[index]!; const embedding = this.embeddings[token]!; const position = positional ? Kernels.positionalEncoding(index, this.dim) : Array.from({ length: this.dim }, () => 0); const sum = embedding.map((value, i) => value + position[i]!); x.push(sum); residualStream.push(sum); // ▸ a token's vector plus a code for where it sits this.tracer?.emit("embed", { index, token: this.vocab[token], embedding: [...embedding], position: [...position], sum: [...sum] }); } const attention: number[][][] = []; const h: number[][] = x.map(() => []); const headOutputs: number[][] = residualStream.map(() => []); for (let head = 0; head < this.heads; head++) { const slice = x.map((row) => row.slice(head * this.headDim, (head + 1) * this.headDim)); const slice = residualStream.map((row) => row.slice(head * this.headDim, (head + 1) * this.headDim)); const { weights, outputs } = this.heads_[head]!.attend(slice); attention.push(weights); for (let index = 0; index < tokens.length; index++) { h[index]!.push(...outputs[index]!); headOutputs[index]!.push(...outputs[index]!); // ▸ one head looks around the sentence from this token this.tracer?.emit("attend", { head, index, weights: [...weights[index]!], output: [...outputs[index]!] }); } } x = x.map((row, index) => this.residual("attention", index, row, this.apply(this.wo, h[index]!))); residualStream = residualStream.map((row, index) => this.residual("attention", index, row, this.apply(this.wo, headOutputs[index]!))); x = x.map((row, index) => this.normalize("attention", index, row)); residualStream = residualStream.map((row, index) => this.normalize("attention", index, row)); const f: number[][] = []; const feedForward: number[][] = []; for (let index = 0; index < tokens.length; index++) { const hiddenRow = this.apply(this.w1, x[index]!).map((value) => this.relu.apply(value)); const hiddenRow = this.apply(this.w1, residualStream[index]!).map((value) => this.relu.apply(value)); const output = this.apply(this.w2, hiddenRow); f.push(output); feedForward.push(output); // ▸ a small relu network works on each row alone this.tracer?.emit("ffn", { index, hidden: [...hiddenRow], output: [...output] }); } x = x.map((row, index) => this.residual("ffn", index, row, f[index]!)); residualStream = residualStream.map((row, index) => this.residual("ffn", index, row, feedForward[index]!)); x = x.map((row, index) => this.normalize("ffn", index, row)); residualStream = residualStream.map((row, index) => this.normalize("ffn", index, row)); const probabilities: number[][] = []; for (let index = 0; index < tokens.length; index++) { const logits = this.apply(this.wout, x[index]!); const logits = this.apply(this.wout, residualStream[index]!); const p = Kernels.softmax(logits, temperature); const top = p.indexOf(Math.max(...p)); probabilities.push(p); // ▸ the row becomes a distribution over the vocabulary this.tracer?.emit("predict", { index, logits: [...logits], probabilities: [...p], top: this.vocab[top], probability: p[top] }); } return { stream: x, attention, probabilities }; return { stream: residualStream, attention, probabilities }; }}