Machine learning

Attention

The one operation underneath every language model. Each word builds a query, every other word offers a key, and the match decides who gets listened to. Everything else in a transformer is plumbing around this.

EXP-023

Move the cursor over the canvas — it is part of the simulation.

Drag your finger across the canvas — it is part of the simulation.

What to look for

Hover over a row to isolate one word and see the arcs of who it is looking at. Then drag sharpness down toward 0.2: the attention spreads until every word listens to everything equally, which is the same as listening to nothing. Push it up and it collapses to a single hard pointer.

The maths

Atención(Q,K,V) = softmax(QKᵀ/√d)·V

The whole thing, in one line. QKᵀ is every query dotted with every key, giving a square matrix of raw affinities. Softmax turns each row into a set of weights that sum to one, and those weights mix the values.

softmax(x)ᵢ = e^xᵢ / Σⱼ e^xⱼ

Turns any list of numbers into a probability distribution. It is computed here after subtracting the row maximum, which changes nothing mathematically and prevents e^x from overflowing.

÷ √d

The detail people skip. Dot products of d-dimensional vectors grow like √d, so without this the softmax input gets large, the exponential saturates, and every row becomes almost one-hot with almost no gradient. Dividing by √d keeps the variance at 1 and the model trainable.

multi-cabeza

The same sentence is projected several times with different weights, so each head can attend to a different kind of relation: one tracks syntax, another tracks what a pronoun refers to. Switching heads here changes the projection matrix and nothing else.

How it is built

Honest about what this is: the sentence and its four-dimensional vectors are hand-written, not learned. The dimensions mean animal, furniture, action and reference, and the four heads are fixed permutation-style matrices. It is a mechanism demo, not a model. The arithmetic, though, is exactly what runs inside a real transformer: same projections, same scaled dot product, same softmax. I checked that every row of the matrix sums to 1.000000.

The code, step by step

Cut straight from lab.js. These are the real lines that run above, not a simplified version.

01

A sentence with hand-made meaning

const TOK = ['the', 'cat', 'sat', 'on', 'the', 'mat', 'because', 'it', 'was', 'tired'];
// Las magnitudes importan: con vectores cerca de cero todos los productos
// punto se parecen, el softmax sale casi plano y la matriz se ve muerta.
// La primera versión tenía ese problema. Los embeddings reales tampoco
// son diminutos.
const E = [
  [0.2, 0.2, 0.1, 0.3], [2.4, 0.1, 0.2, 0.6], [0.3, 0.2, 2.3, 0.2],
  [0.1, 0.8, 0.5, 0.1], [0.2, 0.2, 0.1, 0.3], [0.2, 2.4, 0.1, 0.5],
  [0.1, 0.1, 0.5, 1.0], [1.6, 0.5, 0.1, 2.2], [0.2, 0.1, 1.3, 0.5],
  [1.4, 0.1, 0.9, 0.3]

Ten tokens and a four-dimensional vector each, with axes that mean animal, furniture, action and reference. These are written by hand, not learned — being clear about that is the point. A real model has thousands of dimensions and no axis means anything nameable.

02

Queries, keys, softmax

const proj = v => W.map(r => r.reduce((a, x, i) => a + x * v[i], 0));
const Q = E.map(proj), K = E.map(proj);
// A = softmax(Q·Kᵀ / √d)
const A = Q.map(q => {
  const sc = K.map(k => k.reduce((a, x, i) => a + x * q[i], 0) / Math.sqrt(4) / P.temp);
  const mx = Math.max(...sc), ex = sc.map(v => Math.exp(v - mx));
  const sum = ex.reduce((a, b) => a + b, 0);
  return ex.map(v => v / sum);

Project every token twice, once as a query and once as a key, then dot each query with every key. Dividing by √d keeps the numbers in the range where the exponential still has a gradient. Subtracting the row max before exponentiating changes nothing mathematically and stops e^x from overflowing — the standard stable softmax.

03

Draw who is listening to whom

if (focus >= 0) {
  const by = gy + cell * n + h * 0.10, bx = gx;
  for (let j = 0; j < n; j++) {
    const a = A[focus][j];
    if (a < 0.02) continue;
    const x1 = bx + focus * cell + cell / 2, x2 = bx + j * cell + cell / 2;
    ctx.strokeStyle = `rgba(${acc},${Math.min(0.9, a * 1.6).toFixed(3)})`;
    ctx.lineWidth = 0.5 + a * 5;
    ctx.beginPath();
    ctx.moveTo(x1, by);
    ctx.quadraticCurveTo((x1 + x2) / 2, by - Math.abs(x2 - x1) * 0.45 - 12, x2, by);

One row of the matrix is one word's distribution of attention, and it sums to one. Drawing it as arcs whose thickness is the weight makes the sentence readable in a way the grid alone is not: you can see "it" reaching back for "cat".

Why it matters

This replaced recurrence in 2017 and everything since is built on it. The reason it won is not accuracy, it is shape: every pair of positions is compared in one step, with no sequential dependency, so the whole thing parallelises across a GPU. Recurrent networks could not, and that is the entire story of the last decade of scale.

The whole thing

Everything above, in one piece. This is the entire experiment as it lives in lab.js.

FULL

lab.js · attention

{ id: 'attention',
      name: L('Attention', 'Atención'),
      note: L('What a transformer does to a sentence: every word weighs every other one.',
              'Lo que un transformer le hace a una frase: cada palabra pesa a todas las demás.'),
      params: [
        { key: 'temp', label: L('Softmax sharpness', 'Nitidez del softmax'), min: 0.2, max: 4, step: 0.05, def: 1 },
        { key: 'head', label: L('Head', 'Cabeza'), min: 0, max: 3, step: 1, def: 0 },
        { key: 'row',  label: L('Focus token', 'Token en foco'), min: -1, max: 9, step: 1, def: -1 }
      ],
      make() {
        // Frase fija con vectores hechos a mano en 4 dimensiones interpretables:
        // [animal, mueble, acción, referencia]. No es un modelo entrenado y no
        // pretende serlo: enseña el MECANISMO, que es lo que casi nadie ve.
        const TOK = ['the', 'cat', 'sat', 'on', 'the', 'mat', 'because', 'it', 'was', 'tired'];
        // Las magnitudes importan: con vectores cerca de cero todos los productos
        // punto se parecen, el softmax sale casi plano y la matriz se ve muerta.
        // La primera versión tenía ese problema. Los embeddings reales tampoco
        // son diminutos.
        const E = [
          [0.2, 0.2, 0.1, 0.3], [2.4, 0.1, 0.2, 0.6], [0.3, 0.2, 2.3, 0.2],
          [0.1, 0.8, 0.5, 0.1], [0.2, 0.2, 0.1, 0.3], [0.2, 2.4, 0.1, 0.5],
          [0.1, 0.1, 0.5, 1.0], [1.6, 0.5, 0.1, 2.2], [0.2, 0.1, 1.3, 0.5],
          [1.4, 0.1, 0.9, 0.3]
        ];
        // Cuatro cabezas: cada una proyecta con pesos distintos y por eso
        // atiende a relaciones distintas sobre la misma frase.
        const HEAD = [
          [[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]],
          [[0,0,0,1],[1,0,0,0],[0,1,0,0],[0,0,1,0]],
          [[0,1,0,0],[0,0,1,0],[0,0,0,1],[1,0,0,0]],
          [[1,0,0,1],[0,1,1,0],[1,0,1,0],[0,1,0,1]]
        ];
        return {
          step(ctx, w, h, t, acc, P, M) {
            const n = TOK.length, W = HEAD[P.head | 0];
            const proj = v => W.map(r => r.reduce((a, x, i) => a + x * v[i], 0));
            const Q = E.map(proj), K = E.map(proj);
            // A = softmax(Q·Kᵀ / √d)
            const A = Q.map(q => {
              const sc = K.map(k => k.reduce((a, x, i) => a + x * q[i], 0) / Math.sqrt(4) / P.temp);
              const mx = Math.max(...sc), ex = sc.map(v => Math.exp(v - mx));
              const sum = ex.reduce((a, b) => a + b, 0);
              return ex.map(v => v / sum);
            });

            const focus = M.in ? Math.min(n - 1, Math.floor((M.y / h) * n)) : (P.row | 0);
            const gy = h * 0.14, gx = w * 0.30, cell = Math.min((w * 0.42) / n, (h * 0.62) / n);

            // Matriz de atención
            for (let i = 0; i < n; i++) {
              for (let j = 0; j < n; j++) {
                const a = A[i][j];
                const on = focus < 0 || focus === i;
                ctx.fillStyle = `rgba(${acc},${(a * (on ? 0.95 : 0.12)).toFixed(3)})`;
                ctx.fillRect(gx + j * cell, gy + i * cell, cell - 1, cell - 1);
              }
            }
            ctx.strokeStyle = `rgba(${acc},0.2)`; ctx.lineWidth = 1;
            ctx.strokeRect(gx, gy, cell * n, cell * n);

            ctx.font = '11px monospace';
            for (let i = 0; i < n; i++) {
              const on = focus < 0 || focus === i;
              ctx.fillStyle = `rgba(${acc},${on ? 0.95 : 0.35})`;
              ctx.textAlign = 'right';
              ctx.fillText(TOK[i], gx - 8, gy + i * cell + cell * 0.7);   // consultas
              ctx.textAlign = 'left';
              ctx.save();
              ctx.translate(gx + i * cell + cell * 0.7, gy - 10);
              ctx.rotate(-Math.PI / 4);
              ctx.fillStyle = `rgba(${acc},0.6)`;
              ctx.fillText(TOK[i], 0, 0);                                 // claves
              ctx.restore();
            }
            ctx.textAlign = 'left';

            // Arcos de la fila en foco: a quién mira esa palabra
            if (focus >= 0) {
              const by = gy + cell * n + h * 0.10, bx = gx;
              for (let j = 0; j < n; j++) {
                const a = A[focus][j];
                if (a < 0.02) continue;
                const x1 = bx + focus * cell + cell / 2, x2 = bx + j * cell + cell / 2;
                ctx.strokeStyle = `rgba(${acc},${Math.min(0.9, a * 1.6).toFixed(3)})`;
                ctx.lineWidth = 0.5 + a * 5;
                ctx.beginPath();
                ctx.moveTo(x1, by);
                ctx.quadraticCurveTo((x1 + x2) / 2, by - Math.abs(x2 - x1) * 0.45 - 12, x2, by);
                ctx.stroke();
              }
              for (let j = 0; j < n; j++) {
                ctx.fillStyle = `rgba(${acc},${j === focus ? 1 : 0.5})`;
                ctx.font = '11px monospace';
                ctx.fillText(TOK[j], bx + j * cell + 2, by + 16);
              }
            }

            ctx.font = '11px monospace';
            ctx.fillStyle = `rgba(${acc},0.55)`;
            ctx.fillText('A = softmax(Q·Kᵀ / √d)      rows are queries, columns are keys', 12, h - 28);
            ctx.fillStyle = `rgba(${acc},0.95)`;
            ctx.fillText('head ' + (P.head | 0) + '   sharpness ' + P.temp.toFixed(2) +
              (focus >= 0 ? '   focus: "' + TOK[focus] + '"' : '   (hover a row)'), 12, h - 12);
          }
        };
      }
    },