transformer_block

Builtin index · Statistics & ML

result = transformer_block(x, weights)

Input: x, weights · Output: result

What it is

One transformer block on an (N, D) input x: multi-head self-attention plus a two-layer feed-forward network, each wrapped in a residual connection and a layer_norm. weights is a Record with fields heads (the tuple multi_head_attention expects), ln1_gamma/ln1_beta and ln2_gamma/ln2_beta (length-D layer-norm parameters for the two sublayers), and w1/b1/w2/b2 (the feed-forward layer's weights and biases).

Overloads

One call shape — this function is not overloaded.

Examples

From the documentation, and run by the build -- the output and the figure below are what the engine actually produced. Simplest first.

x = [1.0, 2.0, 3.0; 10.0, 20.0, 30.0]
normed = layer_norm(x, [1.0, 1.0, 1.0], [0.0, 0.0, 0.0])
print("layer_norm (scale-invariant per row): {normed}")

q = [1.0, 0.5; -0.5, 1.0; 0.2, -0.3]
k = [1.0, 0.0; 0.0, 1.0; 0.5, 0.5]
v = [1.0, 2.0; 3.0, 4.0; 5.0, 6.0]
attn = scaled_dot_product_attention(q, k, v)
print("attention output: {round(attn, 3)}")

eye2 = [1.0, 0.0; 0.0, 1.0]
head = {wq = eye2, wk = eye2, wv = eye2, wo = eye2}
mh = multi_head_attention(q, (head,), 1)
plain = scaled_dot_product_attention(q, q, q)
print("multi_head_attention with 1 identity head matches plain self-attention: {sum(abs(mh - plain)) < 1e-9}")

weights = {heads = (head,), ln1_gamma = [1.0, 1.0], ln1_beta = [0.0, 0.0],
           w1 = [0.2, -0.1; 0.1, 0.3], b1 = [0.0, 0.0],
           w2 = [0.1, -0.2; 0.3, 0.1], b2 = [0.0, 0.0],
           ln2_gamma = [1.0, 1.0], ln2_beta = [0.0, 0.0]}
block_out = transformer_block(q, weights)
print("transformer_block: {round(block_out, 3)}")

layer_norm (scale-invariant per row): [-1.224736, 0, 1.224736; -1.224745, 0, 1.224745]
attention output: [2.872, 3.872; 3.25, 4.25; 2.872, 3.872]
multi_head_attention with 1 identity head matches plain self-attention: true
transformer_block: [1, -1; -1, 1; 1, -1]

from statistics-ml

How it works

transformer_block(x, weights) — one post-norm Transformer encoder block (Vaswani et al. 2017, §3.1): multi-head self- attention, a residual add + layer_norm, a two-layer ReLU feedforward (reusing the EXISTING dense/relu builtins — no new layer primitive needed here either), then a second residual add + layer_norm: text a = multi_head_attention(x, weights.heads, len(weights.heads)) h1 = layer_norm(x + a, weights.ln1_gamma, weights.ln1_beta) f = dense(relu(dense(h1, weights.w1, weights.b1)), weights.w2, weights.b2) h2 = layer_norm(h1 + f, weights.ln2_gamma, weights.ln2_beta) weights is a Value::Record with fields heads (the multi_head_attention weights list — n_heads is read straight off its length, never a separate argument), ln1_gamma/ln1_beta/ln2_gamma/ln2_beta, and w1/b1/ w2/b2 (the feedforward's two dense layers). Not its own tape op: every step below calls an already-tracked builtin, so backprop through the whole block falls out of the tape's ordinary reverse walk with no bespoke tensor_vjp rule — the same "compose, don't reimplement" argument dense/ gru_cell_compute already make.

From the note above this builtin in the interpreter, where the reasoning was written down as it was built.