scaled_dot_product_attention

Builtin index · Statistics & ML

result = scaled_dot_product_attention(q, k, v, [mask=])

Input: q, k, v, mask · Output: result

What it is

Core attention computation: an (Nq, D) query matrix q, an (Nk, D) key matrix k, and an (Nk, Dv) value matrix v (k and v share the same row count Nk). Scores are q * kᵀ scaled by 1/sqrt(D) before the softmax, so the softmax does not saturate as D grows; an optional mask matrix (same shape as the score matrix, (Nq, Nk)) adds large negative values to positions that should not be attended to before the softmax.

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

scaled_dot_product_attention(q, k, v, [mask]) — softmax(Q K^T / sqrt(d_k)) V (Vaswani et al. 2017, eq. 1). q/k/v are (seq_len_q, d_k)/(seq_len_k, d_k)/ (seq_len_k, d_v) matrices (or Tensors wrapping them); an optional 0/1 mask of the raw scores' own (seq_len_q, seq_len_k) shape (1 = attend, 0 = block) is added as a large negative bias BEFORE the softmax — the standard "additive mask" convention, not zeroing probabilities after the fact (which would leave blocked entries still skewing the row's normalization). K^T uses tt (this session's new tracked transpose, so gradient flows back into k too); the two matmuls (Q@K^T, probs@V) reuse tb's "*", and so Matrix::matmul's own existing rayon parallelism above its size threshold — the actual CPU-parallel win for this op. True head-level parallelism (see multi_head_attention) can't extend down into this call too: every tracked op here mutates self.tape behind one shared &mut self.

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