Attention in pure Java, with real numbers

I built jamna.nn, a small transformer in Java with no libraries, not even for math, because I wanted to see every number a language model computes. Everything in a transformer serves one line from Attention Is All You Need:

Attention(Q, K, V) = softmax( Q·Kᵀ / √d ) · V

Formulas like this are easy to nod along to and hard to really understand. So let's run it on three tokens, with 2-dimensional vectors, until every number has a reason.

1. Each token gets three roles

A token's embedding is multiplied by three learned weight matrices, W_Q, W_K and W_V, giving three vectors with three jobs:

• query: what this token is looking for
• key: what this token offers to others
• value: what it passes on if it gets chosen

With small hand-picked weights, our three tokens become:

Why three separate vectors and not one? Because "what I'm looking for" and "what I contain" are different questions. A verb looks for its subject, but it doesn't offer itself as a subject. Separate matrices let the model learn the two independently.

2. Score: how well does each key answer token 1's query?

A dot product is large when two vectors point the same way, so it measures "match":

q₁·k₁ = 1.0×1.0 + 0.5×0.0 = 1.0
q₁·k₂ = 1.0×0.5 + 0.5×1.0 = 1.0
q₁·k₃ = 1.0×1.5 + 0.5×1.0 = 2.0

Token 3's key matches best. Note that token 1 also scores itself. Attention is not "look at the others", it's "look at everyone, including me".

3. Scale by √d, the step everyone skips

Here d = 2, so √d ≈ 1.414 and the scores become [0.707, 0.707, 1.414].

Why? If each component of q and k has variance about 1, their dot product is a sum of d such terms, so its variance is about d. In a real model with d = 64 per head, raw scores easily reach ±15. Feed that to softmax and almost all of the weight lands on one token, the other weights are close to zero, and so are their gradients. The model stops learning which other tokens matter. Dividing by √d brings the variance back to about 1, so softmax stays in its useful range.

4. Softmax: scores become a probability split

e^0.707 = 2.028    e^0.707 = 2.028    e^1.414 = 4.113
                                     sum = 8.169

weights = [2.028/8.169, 2.028/8.169, 4.113/8.169]
        = [0.248, 0.248, 0.503]

The exponential makes every weight positive and stretches differences: a score gap of 0.707 became a 2× difference in weight. Doing the same for all three tokens gives the full attention matrix. Each row is one token deciding where to look:

5. Blend the values

Token 1's new vector is the weighted sum of everyone's values:

0.248 × [0.5, 1.0] = [0.124, 0.248]
0.248 × [1.0, 0.0] = [0.248, 0.000]
0.503 × [1.5, 1.0] = [0.755, 0.503]
                     ---------------
new token 1        = [1.127, 0.751]

That's all attention is. Token 1 started as [0.5, 1.0] and is rewritten as a mix in which half the content comes from token 3. Stack many layers of this, each with many heads, and you have the mechanism that lets a word's meaning depend on its context.

The same thing in Java

static double[] attend(double[] q, double[][] keys, double[][] values) {
    int n = keys.length, d = q.length;
    double[] w = new double[n];
    double max = Double.NEGATIVE_INFINITY;

    for (int j = 0; j < n; j++) {               // steps 2–3: score and scale
        double dot = 0;
        for (int i = 0; i < d; i++) dot += q[i] * keys[j][i];
        w[j] = dot / Math.sqrt(d);
        max = Math.max(max, w[j]);
    }

    double sum = 0;                               // step 4: softmax
    for (int j = 0; j < n; j++) {
        w[j] = Math.exp(w[j] - max);
        sum += w[j];
    }

    double[] out = new double[values[0].length];  // step 5: blend values
    for (int j = 0; j < n; j++)
        for (int i = 0; i < out.length; i++)
            out[i] += (w[j] / sum) * values[j][i];
    return out;
}

attend(new double[]{1.0, 0.5}, K, V) with the tables above returns [1.127, 0.751], the same as the hand calculation.

Two lines deserve a note:

• Math.exp(w[j] - max). Subtracting the maximum score doesn't change the result, because softmax is unchanged when the same constant is subtracted from every score. It does stop exp from overflowing to Infinity when scores are large.
• What's missing for a GPT. In a decoder, a causal mask sets the scores of later tokens to −∞ before softmax, which makes their weight exactly 0. A token can't attend to the future it's being trained to predict.

References

• Vaswani et al., Attention Is All You Need (2017)
• jamna.nn, my transformer in pure Java
• jamna.nn documentation, an interactive explainer that walks through embeddings, attention and softmax with the real trained weights