Skip to main content

Medium attention · transformers

No peeking ahead

Why a model writing word five is not allowed to see word six.

A model that writes text one word at a time must not look at words it has not written yet. A causal mask enforces that: word i may only attend to words 0 through i.

Write causalAttention(scores): scores is an n by n array of numbers, where scores[i][j] is how much word i wants to look at word j. Return a new n by n array where row i is the softmax of scores[i][0..i] and every later position (j greater than i) is exactly 0. Each row then adds up to 1. An empty input returns [].

Call show(mask) with the n by n mask of 1s (allowed) and 0s (blocked) and the stage draws it as a grid.

Example: causalAttention([[1, 5], [0, 0]]) returns [[1, 0], [0.5, 0.5]].

Tests you can see

  • causalAttention([[1,5],[0,0]])[[1,0],[0.5,0.5]]
  • causalAttention([[0,9,9],[1,2,9],[0,1,2]])[[1,0,0],[0.2689414213699951,0.7310585786300049,0],[0.09003057317038046,0.24472847105479764,0.6652409557748218]]

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.

javascriptcausalAttention()

Tab indents · Cmd/Ctrl+Enter runs