Hard attention · transformers

Where should I look

The core move inside a transformer, worked out for one word.

Attention lets one word decide how much to listen to each of the others. The word asks a question (its query vector), every word offers a key vector, and the closer they match, the more weight that word gets.

Write attentionWeights(query, keys): query is an array of d numbers and keys is an array of key vectors, each also d long. Score each key as the dot product with query divided by the square root of d. Then turn the scores into weights with softmax (subtract the largest score first so nothing overflows). Return the weights, which add up to 1. With no keys, return [].

Call show(weights) and the stage draws one bar per key.

Example: attentionWeights([1, 0], [[1, 0], [0, 1]]) returns about [0.67, 0.33].

Tests you can see

attentionWeights([1,0], [[1,0],[0,1]])[0.6697615493266569,0.3302384506733431]
attentionWeights([1,2,0,1], [[1,0,0,0],[0,1,0,1],[2,2,1,0]])[0.0628900132458675,0.17095278019779026,0.7661572065563422]

More tests are hidden. The server runs them all before a time counts.

The idea behind this, in a two-minute live page

Call show() while your code works and the run plays back here.

javascriptattentionWeights()

function attentionWeights(query, keys) {
  // dot product, scale by sqrt(d), softmax
  return keys.map(() => 1 / keys.length);
}

function attentionWeights(query, keys) {
  // dot product, scale by sqrt(d), softmax
  return keys.map(() => 1 / keys.length);
}

Tab indents · Cmd/Ctrl+Enter runs