multi_head_attention

Builtin index · Statistics & ML

result = multi_head_attention(x, heads, n_heads)

Input: x, heads, n_heads · Output: result

What it is

Runs several attention heads over the same (N, D) input x in parallel: heads is a tuple of n_heads (integer) Records, each with fields wq, wk, wv (per-head projection matrices) and wo (the head's output projection), and self-attention (q = k = v = x projected through each head's own weights) is computed per head, then the results are combined through each head's wo and summed.

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

multi_head_attention(x, weights, n_heads, [mask]) — splits into n_heads independent heads, each with its own wq/ wk/wv/wo projections, and combines them. weights is a Value::List of exactly n_heads Value::Records, each {wq, wk, wv, wo}: wq/wk/wv are (d_model, d_head) (x @ wq projects x's (seq_len, d_model) down to this head's own (seq_len, d_head) q/k/v), and wo is (d_head, d_model) — this head's own ROW-BLOCK slice of the usual combined output projection. Concatenating the heads' (seq_len, d_head) outputs and multiplying by one big (n_heads*d_head, d_model) Wo is the textbook formulation, but it needs a concat tracked primitive nothing else here needs (TapeEntry only carries two input slots, so a variadic concat was never on the table anyway). Splitting Wo into per-head row blocks (wo above) and SUMMING attn_head @ wo_head instead is the identical computation by the standard block-matmul identity ([A1 A2] @ [B1; B2] == A1@B1 + A2@B2) — so this needs only tb's existing "*"/"+", no new primitive and no new tensor_vjp rule. CPU parallelism: heads run in a sequential Rust for loop, not a rayon fan-out, because every iteration mutates self.tape (&mut self) — racing multiple heads across threads against one shared, non-atomic Vec<TapeEntry> would corrupt node numbering, the same constraint gru_forward's own doc comment documents for timesteps. The actual parallel win is one level down: each head's own Q@K^T/probs@V matmuls already run on qu-core's rayon thread pool once they cross Matrix::PARALLEL_MATMUL_THRESHOLD, same as every other tracked matmul in this interpreter.

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