The Math Behind Transformers, in Java

A transformer is the kind of neural network behind modern language models. From the outside it looks complicated. On the inside it repeats a small set of arithmetic operations billions of times: multiplying and adding lists of numbers, applying a few simple functions, and rescaling numbers to a sensible size. This lesson implements each operation in plain Java on float[] arrays. The vectors have only 3 or 4 numbers, so every result can be checked by hand.

This matters because the code inside an inference engine (the program that runs a model, such as llama.cpp) is made of exactly these loops. If you can read them, you can estimate how much work a model does for each token, see why some calculations overflow, and understand why two correct programs can print slightly different numbers. The lesson assumes only school arithmetic. Each term is explained the first time it appears.

After the lesson you will be able to write these operations yourself, read the inner loops of an engine, count the work in a layer, and explain why float results change when the order of additions changes.

Lesson 5 used a few of these without explanation (softmax, and the product of an embedding with a matrix). This lesson defines them. Lessons 9 and 10 combine them into attention and the transformer block.

Example 1: vectors as float[]

A vector is an ordered list of numbers. In Java it is a float[]. Inside a model, the state of one token (its hidden state, the numbers the model currently holds for that token) is a vector with a few hundred to a few thousand numbers. This lesson uses 3 or 4 numbers so a result fits on one line.

Two operations appear everywhere. Adding two vectors means adding the numbers at the same index. Scaling a vector means multiplying every number by one number. The program below implements both, then combines them. Look at the last line of output, which averages two vectors.

import java.util.Arrays;

public class Vectors {
    static float[] add(float[] a, float[] b) {
        float[] out = new float[a.length];
        for (int i = 0; i < a.length; i++) out[i] = a[i] + b[i];
        return out;
    }

    static float[] scale(float[] a, float s) {
        float[] out = new float[a.length];
        for (int i = 0; i < a.length; i++) out[i] = a[i] * s;
        return out;
    }

    public static void main(String[] args) {
        float[] a = {1, 2, 3};
        float[] b = {0.5f, -1, 4};
        System.out.println("a + b   = " + Arrays.toString(add(a, b)));
        System.out.println("a * 2   = " + Arrays.toString(scale(a, 2)));
        System.out.println("0.5a + 0.5b = " + Arrays.toString(add(scale(a, 0.5f), scale(b, 0.5f))));
    }
}

a + b   = [1.5, 1.0, 7.0]
a * 2   = [2.0, 4.0, 6.0]
0.5a + 0.5b = [0.75, 0.5, 3.5]

Adding (1, 2, 3) and (0.5, -1, 4) gives (1.5, 1, 7), one addition per index. Scaling (1, 2, 3) by 2 doubles every number. The last line uses both operations: it takes half of each vector and adds the halves. Half of each, added together, is the average of the two vectors, (0.75, 0.5, 3.5).

Averaging vectors is the first form of "mixing information" that lesson 9 uses. A float[] of length n is all a vector needs in Java. A float has 32 bits (4 bytes), so a vector of 4,096 numbers takes 16 KiB (kibibytes, units of 1,024 bytes).

Example 2: the dot product

A model often needs one number that says how well two vectors match. The dot product gives that number. Multiply the numbers at each index, then add the results: a . b = a0*b0 + a1*b1 + ....

The sign and size of the result carry meaning. It is large and positive when the vectors point the same way, zero when they are at a right angle (they share no direction), and negative when they point apart. Attention (lesson 9) uses this to score how well one token's query matches another token's key.

The program tests four pairs. Check each printed number against the arithmetic you can do in your head.

public class DotProduct {
    static float dot(float[] a, float[] b) {
        float sum = 0;
        for (int i = 0; i < a.length; i++) sum += a[i] * b[i];
        return sum;
    }

    public static void main(String[] args) {
        float[] x = {1, 0, 0};
        System.out.println("same direction : " + dot(x, new float[] {2, 0, 0}));
        System.out.println("right angle    : " + dot(x, new float[] {0, 5, 0}));
        System.out.println("opposite       : " + dot(x, new float[] {-3, 0, 0}));
        System.out.println("mixed          : " + dot(new float[] {1, 2, 3}, new float[] {4, -5, 6}));
    }
}

same direction : 2.0
right angle    : 0.0
opposite       : -3.0
mixed          : 12.0

For (1, 0, 0) and (2, 0, 0) the dot product is 1*2 + 0*0 + 0*0 = 2, so the vectors point the same way. With (0, 5, 0) it is 1*0 + 0*5 + 0*0 = 0, because the vectors are at a right angle. With (-3, 0, 0) it is -3, because the vectors point in opposite directions.

The last line computes the following by hand:

1*4 + 2*(-5) + 3*6 = 4 - 10 + 18 = 12

The sign tells the relation: positive means similar, zero means unrelated, negative means opposed. The size depends on two things at once, the match and the lengths of the vectors. A long vector gives a large dot product even for a weak match. The next example separates the two effects. The inner loop of dot is also where an inference engine spends most of its time, so its shape is worth remembering: one multiply and one add per index.

Example 3: length and cosine similarity

The dot product mixes direction with size. Doubling a vector doubles its dot product with everything, but its direction did not change. To compare direction alone, divide out the sizes.

The length (also called the norm) of a vector is the square root of its dot product with itself. For (3, 4) that is sqrt(9 + 16) = 5, the same distance you get from the Pythagorean theorem. Cosine similarity divides the dot product by both lengths, which leaves only direction. It is always between -1 (opposite) and 1 (same direction), with 0 for a right angle.

The program compares (3, 4) with four other vectors. Watch how the dot product changes with length while the cosine does not.

import java.util.Locale;

public class Cosine {
    static float dot(float[] a, float[] b) {
        float sum = 0;
        for (int i = 0; i < a.length; i++) sum += a[i] * b[i];
        return sum;
    }

    static float length(float[] a) { return (float) Math.sqrt(dot(a, a)); }

    static float cosine(float[] a, float[] b) { return dot(a, b) / (length(a) * length(b)); }

    public static void main(String[] args) {
        float[] a = {3, 4};
        System.out.println("length of (3, 4) = " + length(a));
        float[][] others = {{6, 8}, {4, -3}, {-3, -4}, {30, 40.5f}};
        for (float[] b : others) {
            System.out.println(String.format(Locale.ROOT, "dot = %7.1f   cosine with (3, 4) = %6.3f   for (%.1f, %.1f)", dot(a, b), cosine(a, b), b[0], b[1]));
        }
    }
}

length of (3, 4) = 5.0
dot =    50.0   cosine with (3, 4) =  1.000   for (6.0, 8.0)
dot =     0.0   cosine with (3, 4) =  0.000   for (4.0, -3.0)
dot =   -25.0   cosine with (3, 4) = -1.000   for (-3.0, -4.0)
dot =   252.0   cosine with (3, 4) =  1.000   for (30.0, 40.5)

The length of (3, 4) is sqrt(9 + 16) = 5. The vector (6, 8) points the same way and is twice as long, so the dot product is 3*6 + 4*8 = 50 and the cosine is 1.000. The vector (4, -3) is at a right angle: dot 0, cosine 0. The vector (-3, -4) points the opposite way: dot -25, cosine -1.

The last vector, (30, 40.5), is ten times longer than (3, 4), so its dot product is large (252). Its direction is almost the same, so its cosine rounds to 1.000. This is the reason cosine is used to compare embeddings: it ignores how long the vectors are and reports only whether they point the same way.

Example 4: hand-made embeddings and their similarities

A computer cannot work with words directly, so a model maps each token to a vector. That vector is an embedding. The model learns the numbers during training, and tokens with similar meaning end up with vectors that point in similar directions.

This program uses 3-number embeddings written by hand. The three numbers loosely mean "animal", "vehicle", and "small". The numbers are invented for the lesson. A real embedding has hundreds or thousands of dimensions, and no person assigns them. The program prints the cosine similarity of every pair of words. Read the table by looking for which words have high values together.

import java.util.Locale;

public class Embeddings {
    static float cosine(float[] a, float[] b) {
        float d = 0, na = 0, nb = 0;
        for (int i = 0; i < a.length; i++) { d += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i]; }
        return d / (float) Math.sqrt(na * nb);
    }

    public static void main(String[] args) {
        String[] words = {"cat", "kitten", "dog", "truck", "scooter"};
        float[][] e = {{0.9f, 0.0f, 0.5f}, {0.9f, 0.0f, 0.9f}, {0.9f, 0.1f, 0.3f}, {0.0f, 0.9f, 0.1f}, {0.0f, 0.8f, 0.7f}};
        System.out.print("          ");
        for (String w : words) System.out.print(String.format("%8s", w));
        System.out.println();
        for (int i = 0; i < words.length; i++) {
            System.out.print(String.format("%-10s", words[i]));
            for (int j = 0; j < words.length; j++) System.out.print(String.format(Locale.ROOT, "%8.2f", cosine(e[i], e[j])));
            System.out.println();
        }
    }
}

               cat  kitten     dog   truck scooter
cat           1.00    0.96    0.98    0.05    0.32
kitten        0.96    1.00    0.89    0.08    0.47
dog           0.98    0.89    1.00    0.14    0.29
truck         0.05    0.08    0.14    1.00    0.82
scooter       0.32    0.47    0.29    0.82    1.00

The table is symmetric, because the similarity of A to B equals the similarity of B to A. Every word has similarity 1.00 with itself. The words cat, kitten, and dog are close to each other (0.89 to 0.98), because they all have a large "animal" number. The words truck and scooter are close to each other (0.82) and far from the animals (0.05 to 0.47).

The value for kitten and scooter (0.47) is the highest between the two groups. Both words have a large "small" number, so they share some direction even though one is an animal and one is a vehicle.

A table like this is how embedding similarity search works. A program turns each text into a vector and ranks the other texts by cosine. Inside a model, the same geometry lets later layers treat similar tokens in similar ways.

Example 5: matrix times vector

A matrix is a table of numbers with rows and columns. A model has to compute many dot products at once: one vector against many other vectors. A matrix-vector product does this. If the matrix has R rows and C columns and the vector has C numbers, the result is a vector with R numbers. Output number r is the dot product of row r with the vector.

The shapes must match: the number of columns in the matrix equals the length of the vector. The weights of a model are stored as matrices, so for each token the model generates, almost all the work is matrix-vector products.

The program multiplies a 2-row, 3-column matrix by a vector of 3 numbers. It then calls the function again with a vector of the wrong length to show what a shape error looks like.

import java.util.Arrays;

public class MatVec {
    static float[] matVec(float[][] m, float[] x) {
        if (m[0].length != x.length) {
            throw new IllegalArgumentException("matrix has " + m[0].length + " columns but vector has " + x.length + " numbers");
        }
        float[] y = new float[m.length];
        for (int r = 0; r < m.length; r++) {
            for (int c = 0; c < x.length; c++) y[r] += m[r][c] * x[c];
        }
        return y;
    }

    public static void main(String[] args) {
        float[][] m = {{1, 0, 2}, {0, 1, -1}};
        float[] x = {3, 4, 5};
        System.out.println("(2 x 3) times (3) = " + Arrays.toString(matVec(m, x)));
        try {
            matVec(m, new float[] {1, 2});
        } catch (IllegalArgumentException e) {
            System.out.println("error: " + e.getMessage());
        }
    }
}

(2 x 3) times (3) = [13.0, -1.0]
error: matrix has 3 columns but vector has 2 numbers

The first row of the matrix, (1, 0, 2), dotted with (3, 4, 5) gives 3 + 0 + 10 = 13. The second row, (0, 1, -1), gives 0 + 4 - 5 = -1. The matrix has 2 rows and 3 columns, so it maps a vector of 3 numbers to a vector of 2 numbers. The shape changes from 3 to 2.

The second call fails because a vector of 2 numbers does not fit a matrix with 3 columns. Shape mismatches are the most common bug when writing model code by hand. Checking shapes at the start of a function gives a clear error message at the point of the mistake, and not a confusing wrong number later.

Example 6: matrix times matrix on a flat array

Multiplying two matrices repeats the matrix-vector product once per column of the second matrix. The product of an M x K matrix and a K x N matrix is an M x N matrix. Its entry (i, j) is the dot product of row i of the first matrix and column j of the second.

Real engines do not store a matrix as float[][], which keeps one separate array per row. They use one flat float[] with the rows placed one after another. This is row-major order. A flat array is one contiguous block of memory, so the CPU reads it quickly, and no pointer per row is needed. In a matrix with cols columns, the element at row i and column j is at index i * cols + j.

The program multiplies a 2x3 matrix by a 3x2 matrix, both stored flat. Check the index arithmetic in the first output line and the first product entry in the last line.

import java.util.Arrays;

public class MatMul {
    static float[] matMul(float[] a, float[] b, int m, int k, int n) {
        float[] c = new float[m * n];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                float sum = 0;
                for (int p = 0; p < k; p++) sum += a[i * k + p] * b[p * n + j];
                c[i * n + j] = sum;
            }
        }
        return c;
    }

    public static void main(String[] args) {
        float[] a = {1, 2, 3, 4, 5, 6};          // 2 x 3
        float[] b = {7, 8, 9, 10, 11, 12};       // 3 x 2
        System.out.println("A (2x3) flat = " + Arrays.toString(a) + "  element (1,2) at index " + (1 * 3 + 2) + " = " + a[1 * 3 + 2]);
        float[] c = matMul(a, b, 2, 3, 2);
        System.out.println("C (2x2) flat = " + Arrays.toString(c));
        System.out.println("check C(0,0) = 1*7 + 2*9 + 3*11 = " + (1 * 7 + 2 * 9 + 3 * 11));
    }
}

A (2x3) flat = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]  element (1,2) at index 5 = 6.0
C (2x2) flat = [58.0, 64.0, 139.0, 154.0]
check C(0,0) = 1*7 + 2*9 + 3*11 = 58

The flat array [1, 2, 3, 4, 5, 6] holds the matrix [[1, 2, 3], [4, 5, 6]] row by row. The element in row 1, column 2 is at index 1 * 3 + 2 = 5, and its value is 6.

The product has shape 2 by 2, and its flat values are [58, 64, 139, 154]. The check line computes the first entry by hand:

C(0,0) = 1*7 + 2*9 + 3*11 = 7 + 18 + 33 = 58

This uses row 0 of A and column 0 of B. Column 0 of B sits at indices 0, 2, and 4 of the flat array, so the code reads B with a step of n = 2.

Three nested loops do M * N * K multiply-adds. Real engines reorder and split these loops into tiles so that the CPU caches (small fast memories next to the processor, covered in lesson 16) are used well. The reordered loops compute the same numbers, up to the rounding effect shown in Example 15.

Example 7: a linear layer and its parameter count

A linear layer is the main building block of a model. It computes y = W x + b. The weight matrix W has out rows and in columns. The vector b is the bias, out numbers added to the result. The weights and the bias are the parameters of the layer: the numbers that training adjusts. A layer has out * in + out parameters. Many transformer layers use no bias, and then they have out * in.

The program applies one small layer to a vector, then counts the parameters of layer sizes that appear in real models. Look at how fast the counts grow with the sizes.

import java.util.Arrays;

public class Linear {
    static float[] linear(float[] w, float[] b, float[] x, int out, int in) {
        float[] y = new float[out];
        for (int r = 0; r < out; r++) {
            float sum = b == null ? 0 : b[r];
            for (int c = 0; c < in; c++) sum += w[r * in + c] * x[c];
            y[r] = sum;
        }
        return y;
    }

    public static void main(String[] args) {
        float[] w = {1, 2, 3, -1, 0, 1};     // 2 outputs, 3 inputs
        float[] b = {0.5f, -0.5f};
        System.out.println("y = " + Arrays.toString(linear(w, b, new float[] {1, 1, 1}, 2, 3)));

        long[][] sizes = {{2, 3}, {896, 896}, {4096, 4096}, {4096, 14336}};
        for (long[] s : sizes) {
            long out = s[0], in = s[1];
            System.out.println(String.format("out %5d, in %5d: with bias %,d parameters, without bias %,d", out, in, out * in + out, out * in));
        }
    }
}

y = [6.5, -0.5]
out     2, in     3: with bias 8 parameters, without bias 6
out   896, in   896: with bias 803,712 parameters, without bias 802,816
out  4096, in  4096: with bias 16,781,312 parameters, without bias 16,777,216
out  4096, in 14336: with bias 58,724,352 parameters, without bias 58,720,256

For the input (1, 1, 1), the first output is 1 + 2 + 3 + 0.5 = 6.5, and the second is -1 + 0 + 1 - 0.5 = -0.5. The parameter counts follow the formula out * in (+ out). The small layer has 2 * 3 = 6 weights and 2 biases, so 8 with a bias.

The larger layers use sizes taken from real models. The number 896 is the hidden width of Qwen2-0.5B (the length of the vector for each token). The numbers 4,096 and 14,336 are the hidden width of Llama 3 8B and the wider inner size used inside each of its layers (lesson 10 reads such a configuration). A square layer of 896 by 896 has 802,816 weights. A layer of 4,096 by 4,096 has 16,777,216 weights, about 16.8 million. A layer with 14,336 inputs and 4,096 outputs has 58.7 million.

Llama 3 8B has 32 layers with 7 weight matrices each, which is 224 weight matrices. Each has between 4 and 59 million parameters. This is why the work per token is dominated by matrix-vector products.

Example 8: activation functions

Suppose a model were built only from linear layers. Stacking linear layers gives another linear layer, no matter how many are stacked, so the extra layers would add nothing. An activation function fixes this. It is a simple non-linear function applied to each number between linear layers. It is what lets a stack of layers represent complicated behavior.

Three functions appear in this course:

Phi(x) in the GELU row is the probability that a standard bell-curve random number is below x. The sigmoid function squeezes any number into the range 0 to 1. The MLP (multi-layer perceptron) is the part of each transformer layer made of linear layers and an activation function.

The program prints all three functions for inputs from -3 to 3. Compare the columns at negative inputs and at positive inputs.

import java.util.Locale;

public class Activations {
    static double relu(double x) { return Math.max(0, x); }

    static double gelu(double x) {
        return 0.5 * x * (1 + Math.tanh(Math.sqrt(2 / Math.PI) * (x + 0.044715 * x * x * x)));
    }

    static double silu(double x) { return x / (1 + Math.exp(-x)); }

    public static void main(String[] args) {
        System.out.println("     x     ReLU     GELU     SiLU");
        for (double x = -3; x <= 3; x += 1) {
            System.out.println(String.format(Locale.ROOT, "%6.1f %8.3f %8.3f %8.3f", x, relu(x), gelu(x), silu(x)));
        }
    }
}

     x     ReLU     GELU     SiLU
  -3.0    0.000   -0.004   -0.142
  -2.0    0.000   -0.045   -0.238
  -1.0    0.000   -0.159   -0.269
   0.0    0.000    0.000    0.000
   1.0    1.000    0.841    0.731
   2.0    2.000    1.955    1.762
   3.0    3.000    2.996    2.858

ReLU is 0 for every negative input and passes positive input through unchanged. GELU and SiLU are smooth curves with a similar shape:

• For large negative input, both are close to 0 (GELU -0.004 and SiLU -0.142 at -3).
• For negative input in general, both stay slightly below zero (the lowest value is near -0.17 for GELU and near -0.28 for SiLU).
• For large positive input, both approach the identity (GELU 2.996 and SiLU 2.858 at 3).

At 1, GELU and SiLU give 0.841 and 0.731 where ReLU gives 1. The small differences matter, because a model is trained with one function and must run with that same function. An engine cannot swap SiLU for GELU and still get correct answers. The formula x * sigmoid(x) for SiLU is the one lesson 10 uses inside the Llama-style MLP.

Example 9: softmax

A model's last layer produces one raw score per possible next token. These scores are called logits, and they can be any real numbers. To choose a token, the model needs probabilities: positive numbers that add up to 1. Softmax converts logits into probabilities. It raises e (about 2.718) to the power of each score, then divides each result by the sum of all of them: p_i = e^(z_i) / sum_j e^(z_j). A larger score gets a larger probability.

Take the logits (2.0, 1.0, 0.1). By hand:

e^2.0 = 7.389    e^1.0 = 2.718    e^0.1 = 1.105    sum = 11.212
p = (7.389 / 11.212, 2.718 / 11.212, 1.105 / 11.212) = (0.659, 0.242, 0.099)

Lesson 5 used softmax on the output of a model. Attention (lesson 9) uses it to turn match scores into weights.

The program computes softmax for these logits and checks that the sum is 1. The last line adds 100 to every logit and runs softmax again. Compare the two results.

import java.util.Arrays;
import java.util.Locale;

public class Softmax {
    static double[] softmax(double[] z) {
        double sum = 0;
        double[] p = new double[z.length];
        for (int i = 0; i < z.length; i++) { p[i] = Math.exp(z[i]); sum += p[i]; }
        for (int i = 0; i < z.length; i++) p[i] /= sum;
        return p;
    }

    public static void main(String[] args) {
        double[] z = {2.0, 1.0, 0.1};
        double[] p = softmax(z);
        StringBuilder sb = new StringBuilder();
        double total = 0;
        for (double v : p) { sb.append(String.format(Locale.ROOT, "%.4f ", v)); total += v; }
        System.out.println("logits " + Arrays.toString(z) + " -> probabilities " + sb.toString().trim());
        System.out.println(String.format(Locale.ROOT, "sum = %.4f", total));
        double[] shifted = {102.0, 101.0, 100.1};
        System.out.println("adding 100 to every logit: " + Arrays.toString(Arrays.stream(softmax(shifted)).mapToObj(v -> String.format(Locale.ROOT, "%.4f", v)).toArray()));
    }
}

logits [2.0, 1.0, 0.1] -> probabilities 0.6590 0.2424 0.0986
sum = 1.0000
adding 100 to every logit: [0.6590, 0.2424, 0.0986]

The logits (2.0, 1.0, 0.1) become probabilities (0.659, 0.242, 0.099), matching the hand calculation. The largest logit gets the largest share, and the numbers add up to 1.

The last line adds 100 to every logit and gets exactly the same probabilities. Adding a constant multiplies every e^z by the same factor, and the division cancels that factor. So softmax depends on the differences between logits, and not on their absolute values. The next example uses this fact to fix a real problem.

Example 10: overflow and the stable softmax

A float cannot hold numbers above about 3.4e38, and e^89 is already larger than that. A model can produce logits of 100 or more. With such logits the direct formula computes infinity divided by infinity, which is NaN ("not a number", the value Java gives to an undefined result). One NaN spoils everything computed after it.

Example 9 shows the way out. Subtracting the same number m from every logit does not change the result: e^(z_i - m) / sum_j e^(z_j - m) equals e^(z_i) / sum_j e^(z_j) for any m. If m is the largest logit, every exponent is 0 or negative, so no value can overflow. This is the stable softmax. The softmax in ggml (the tensor library inside llama.cpp, function ggml_vec_soft_max_f32 in ggml/src/ggml-cpu/vec.cpp) subtracts the maximum before it exponentiates.

The program computes both versions on logits of 1000, 1001, and 1002, and prints the largest float for reference.

import java.util.Arrays;

public class StableSoftmax {
    static float[] naive(float[] z) {
        float sum = 0;
        float[] p = new float[z.length];
        for (int i = 0; i < z.length; i++) { p[i] = (float) Math.exp(z[i]); sum += p[i]; }
        for (int i = 0; i < z.length; i++) p[i] /= sum;
        return p;
    }

    static float[] stable(float[] z) {
        float max = Float.NEGATIVE_INFINITY;
        for (float v : z) max = Math.max(max, v);
        float sum = 0;
        float[] p = new float[z.length];
        for (int i = 0; i < z.length; i++) { p[i] = (float) Math.exp(z[i] - max); sum += p[i]; }
        for (int i = 0; i < z.length; i++) p[i] /= sum;
        return p;
    }

    public static void main(String[] args) {
        float[] z = {1000f, 1001f, 1002f};
        System.out.println("largest float = " + Float.MAX_VALUE);
        System.out.println("naive  : " + Arrays.toString(naive(z)));
        System.out.println("stable : " + Arrays.toString(stable(z)));
        System.out.println("stable on {0, 1, 2} : " + Arrays.toString(stable(new float[] {0f, 1f, 2f})));
    }
}

largest float = 3.4028235E38
naive  : [NaN, NaN, NaN]
stable : [0.09003057, 0.24472848, 0.66524094]
stable on {0, 1, 2} : [0.09003057, 0.24472848, 0.66524094]

The largest float is about 3.4e38, and e^1000 is far beyond it. The direct formula gives infinity for every entry, and infinity divided by infinity is NaN. All three outputs are NaN.

The stable version subtracts the maximum (1002) first, so the exponents become -2, -1, and 0. The result for {1000, 1001, 1002} equals the result for {0, 1, 2}, which is (0.090, 0.245, 0.665). The two inputs differ by a constant, so their softmax results are the same.

An engine that skipped this step would produce NaN logits for a confident model, and the next token would be garbage. Fast attention implementations, covered in the attention-optimizations lesson later in this course, apply the same trick one block of numbers at a time.

Example 11: log-softmax and cross-entropy loss

Training needs a number that says how wrong a prediction was. The usual choice is based on the probability the model gave to the correct next token. If that probability is high, the model was right and the error should be small. If it is low, the error should be large. The cross-entropy loss for one prediction is -log p_target, where target is the index of the correct token (lesson 5). Here log is the natural logarithm.

By hand, with p = 0.5, the loss is -log(0.5) = 0.693. With p = 0.25 it is -log(0.25) = 1.386. With p = 1 it is 0.

A direct way to compute this is to run softmax and then take the log. That loses precision when a probability is tiny. Log-softmax computes the logarithm in one step, with the same maximum m as before: log p_i = z_i - m - log(sum_j e^(z_j - m)).

The program computes the loss both ways for three predictions of a 4-token vocabulary, where token 0 is the correct one: a confident and right prediction, an unsure one, and a confident and wrong one. The last line shows a case where only log-softmax works.

import java.util.Locale;

public class CrossEntropy {
    static double[] logSoftmax(double[] z) {
        double max = Double.NEGATIVE_INFINITY;
        for (double v : z) max = Math.max(max, v);
        double sum = 0;
        for (double v : z) sum += Math.exp(v - max);
        double logSum = Math.log(sum);
        double[] out = new double[z.length];
        for (int i = 0; i < z.length; i++) out[i] = z[i] - max - logSum;
        return out;
    }

    public static void main(String[] args) {
        double[][] cases = {{5, 0, 0, 0}, {0.1, 0, 0.2, 0}, {0, 5, 0, 0}};
        String[] names = {"confident, right", "unsure", "confident, wrong"};
        int target = 0;
        for (int c = 0; c < cases.length; c++) {
            double[] lp = logSoftmax(cases[c]);
            double sum = 0;
            for (double v : cases[c]) sum += Math.exp(v);
            double viaSoftmax = -Math.log(Math.exp(cases[c][target]) / sum);
            System.out.println(String.format(Locale.ROOT, "%-18s p(target) = %.4f  loss = %.4f  (via softmax: %.4f)", names[c], Math.exp(lp[target]), -lp[target], viaSoftmax));
        }
        double[] big = {-200, 0};
        System.out.println(String.format(Locale.ROOT, "log p of the unlikely token with logits {-200, 0}: log-softmax %.2f", logSoftmax(big)[0]));
    }
}

confident, right   p(target) = 0.9802  loss = 0.0200  (via softmax: 0.0200)
unsure             p(target) = 0.2554  loss = 1.3648  (via softmax: 1.3648)
confident, wrong   p(target) = 0.0066  loss = 5.0200  (via softmax: 5.0200)
log p of the unlikely token with logits {-200, 0}: log-softmax -200.00

A confident and correct prediction (p = 0.9802) costs a loss of 0.0200. An unsure prediction (p = 0.2554) costs 1.3648. A confident and wrong prediction (p = 0.0066) costs 5.0200. The loss grows quickly as the probability of the right answer shrinks, which is why training pushes hardest on the predictions that are most wrong.

The two routes agree for ordinary values. The last line shows why log-softmax exists. For logits {-200, 0}, the probability of the first token is e^-200, about 1e-87. That is too small for a float: it would round to 0, and log(0) is negative infinity. Log-softmax never forms the tiny probability and returns -200.00 exactly.

The same quantity, averaged over a text, is the cross-entropy and perplexity of lesson 5.

Example 12: embedding lookup

A model starts with token ids, which are whole numbers, and needs vectors. An embedding table is a matrix with one row per token in the vocabulary (the set of all tokens the model knows). The embedding of token t is row t. A table for V tokens and d dimensions has V * d parameters.

The same result comes from multiplying a one-hot vector by the table. A one-hot vector has zeros everywhere and a single 1 at position t. The product picks out row t because every other row is multiplied by zero. This shows that a lookup is a special matrix product. Engines use the lookup, because multiplying by a vector of zeros wastes work.

The program builds a small table, looks up three tokens both ways, and compares the results. The last line counts the parameters of a real vocabulary size.

import java.util.Arrays;

public class EmbeddingLookup {
    public static void main(String[] args) {
        int vocab = 6, dims = 4;
        float[] table = new float[vocab * dims];
        for (int i = 0; i < table.length; i++) table[i] = i / 10f;       // row t holds t*0.4 + (0 .. 0.3)
        int[] tokens = {2, 0, 5};
        for (int t : tokens) {
            float[] row = Arrays.copyOfRange(table, t * dims, (t + 1) * dims);
            float[] viaOneHot = new float[dims];
            for (int v = 0; v < vocab; v++) {
                float hot = v == t ? 1f : 0f;
                for (int d = 0; d < dims; d++) viaOneHot[d] += hot * table[v * dims + d];
            }
            System.out.println("token " + t + " -> lookup " + Arrays.toString(row) + "  one-hot product equal: " + Arrays.equals(row, viaOneHot));
        }
        long realVocab = 151_936, realDims = 896;
        System.out.println(String.format("a table for V = %,d and d = %d has %,d parameters", realVocab, realDims, realVocab * realDims));
    }
}

token 2 -> lookup [0.8, 0.9, 1.0, 1.1]  one-hot product equal: true
token 0 -> lookup [0.0, 0.1, 0.2, 0.3]  one-hot product equal: true
token 5 -> lookup [2.0, 2.1, 2.2, 2.3]  one-hot product equal: true
a table for V = 151,936 and d = 896 has 136,134,656 parameters

Looking up row 2 of the table gives [0.8, 0.9, 1.0, 1.1]. Multiplying a one-hot vector by the table gives the same numbers for all three tokens, and Arrays.equals confirms it. The lookup copies d numbers and does no arithmetic. The one-hot product does V * d multiplications, almost all of them by zero.

The last line counts the table of a real vocabulary: 151,936 tokens times 896 dimensions is 136,134,656, about 136 million parameters. That is more than a quarter of the whole 0.5-billion-parameter model it belongs to. For small models, the embedding table is therefore a large share of the file. Many models also reuse it as the output matrix, a technique called tied embeddings (lesson 10).

Example 13: LayerNorm and RMSNorm

The numbers inside a deep network can grow or shrink from layer to layer. If they grow too large, calculations overflow. If they shrink too far, information is lost. A normalization step rescales a vector to a standard size before the next layer uses it.

Two versions are common. LayerNorm subtracts the mean of the vector's numbers, then divides by their standard deviation (a measure of how spread out the numbers are): (x - mean) / sqrt(variance + eps). RMSNorm skips the mean and divides by the root mean square, which is the square root of the average of the squared numbers: x / sqrt(mean(x^2) + eps). In both, eps is a tiny constant that avoids division by zero. In both, a learned weight vector with the same length as x multiplies the result.

Llama, Qwen, and many recent models use RMSNorm. It does not compute the mean, so it needs one pass over the numbers fewer than LayerNorm.

By hand, for x = (1, 2, 3, 4):

LayerNorm: mean = 2.5, deviations = (-1.5, -0.5, 0.5, 1.5)
           variance = (2.25 + 0.25 + 0.25 + 2.25) / 4 = 1.25, std = 1.118
           result = (-1.5, -0.5, 0.5, 1.5) / 1.118
RMSNorm:   mean of squares = (1 + 4 + 9 + 16) / 4 = 7.5, root = 2.7386
           result = (1, 2, 3, 4) / 2.7386

The program computes both with all weights set to 1, checks the size after RMSNorm, and feeds in a vector that is 100 times larger.

import java.util.Arrays;
import java.util.Locale;

public class Norms {
    static float[] layerNorm(float[] x, float[] w, float eps) {
        float mean = 0;
        for (float v : x) mean += v;
        mean /= x.length;
        float var = 0;
        for (float v : x) var += (v - mean) * (v - mean);
        var /= x.length;
        float inv = 1f / (float) Math.sqrt(var + eps);
        float[] y = new float[x.length];
        for (int i = 0; i < x.length; i++) y[i] = (x[i] - mean) * inv * w[i];
        return y;
    }

    static float[] rmsNorm(float[] x, float[] w, float eps) {
        float ss = 0;
        for (float v : x) ss += v * v;
        float inv = 1f / (float) Math.sqrt(ss / x.length + eps);
        float[] y = new float[x.length];
        for (int i = 0; i < x.length; i++) y[i] = x[i] * inv * w[i];
        return y;
    }

    static String fmt(float[] a) {
        StringBuilder sb = new StringBuilder("[");
        for (float v : a) sb.append(String.format(Locale.ROOT, "%.4f ", v));
        return sb.toString().trim() + "]";
    }

    public static void main(String[] args) {
        float[] x = {1, 2, 3, 4};
        float[] ones = {1, 1, 1, 1};
        float[] ln = layerNorm(x, ones, 1e-5f), rms = rmsNorm(x, ones, 1e-5f);
        System.out.println("x         = " + Arrays.toString(x));
        System.out.println("LayerNorm = " + fmt(ln));
        System.out.println("RMSNorm   = " + fmt(rms));
        float sq = 0;
        for (float v : rms) sq += v * v;
        System.out.println(String.format(Locale.ROOT, "mean of squares after RMSNorm = %.4f", sq / 4));
        System.out.println("RMSNorm of 100 * x = " + fmt(rmsNorm(new float[] {100, 200, 300, 400}, ones, 1e-5f)));
    }
}

x         = [1.0, 2.0, 3.0, 4.0]
LayerNorm = [-1.3416 -0.4472 0.4472 1.3416]
RMSNorm   = [0.3651 0.7303 1.0954 1.4606]
mean of squares after RMSNorm = 1.0000
RMSNorm of 100 * x = [0.3651 0.7303 1.0954 1.4606]

LayerNorm turns (1, 2, 3, 4) into (-1.3416, -0.4472, 0.4472, 1.3416), matching the hand calculation: the mean (2.5) is subtracted, and the result is divided by the standard deviation (about 1.118). The output has mean 0 and variance 1.

RMSNorm gives (0.3651, 0.7303, 1.0954, 1.4606). The vector is divided by its root mean square (sqrt(7.5) = 2.7386), and the output has a mean of squares of exactly 1.0000. RMSNorm does not move the center of the numbers.

The last line scales the input by 100 and gets the same output. This is the purpose of normalization: the next layer sees numbers of the same size no matter how large the earlier numbers became. In real layers, the learned weight vector then stretches each dimension. This lesson sets it to 1 to keep the numbers readable.

Example 14: counting FLOPs

To judge how fast a model can run, you need to know how much arithmetic it does. A FLOP is one floating-point operation, such as one add or one multiply. Each step of a dot product is a multiply followed by an add, so it counts as 2 FLOPs.

A matrix-vector product with R rows and C columns does R * C multiplies and about R * C adds, so it costs 2 * R * C FLOPs. Every weight in the model takes part in exactly one multiply for each token. A model with P parameters therefore needs about 2 * P FLOPs to produce one token.

The program first counts the operations inside a real matrix-vector loop to check the formula. It then prints the per-token work for several model sizes. GFLOP means billions of operations, and TFLOP means trillions.

import java.util.Locale;

public class Flops {
    static long counted;

    static float[] matVec(float[] w, float[] x, int rows, int cols) {
        float[] y = new float[rows];
        for (int r = 0; r < rows; r++) {
            float sum = 0;
            for (int c = 0; c < cols; c++) {
                sum += w[r * cols + c] * x[c];
                counted += 2;
            }
            y[r] = sum;
        }
        return y;
    }

    public static void main(String[] args) {
        int rows = 64, cols = 32;
        matVec(new float[rows * cols], new float[cols], rows, cols);
        System.out.println("counted = " + counted + ", formula 2 * rows * cols = " + 2L * rows * cols);

        String[] names = {"0.5B", "1.5B", "7B", "8B", "70B"};
        double[] params = {0.5e9, 1.5e9, 7e9, 8e9, 70e9};
        for (int i = 0; i < names.length; i++) {
            double perToken = 2 * params[i];
            System.out.println(String.format(Locale.ROOT, "%-5s about %6.1f GFLOP per generated token; a 1000-token prompt needs %8.1f TFLOP", names[i], perToken / 1e9, perToken * 1000 / 1e12));
        }
    }
}

counted = 4096, formula 2 * rows * cols = 4096
0.5B  about    1.0 GFLOP per generated token; a 1000-token prompt needs      1.0 TFLOP
1.5B  about    3.0 GFLOP per generated token; a 1000-token prompt needs      3.0 TFLOP
7B    about   14.0 GFLOP per generated token; a 1000-token prompt needs     14.0 TFLOP
8B    about   16.0 GFLOP per generated token; a 1000-token prompt needs     16.0 TFLOP
70B   about  140.0 GFLOP per generated token; a 1000-token prompt needs    140.0 TFLOP

The counted operations (4,096) equal 2 * 64 * 32, as the formula says. For models, 2 * parameters gives the work per generated token: 1.0 GFLOP for a 0.5B model, 14 GFLOP for a 7B model, and 140 GFLOP for a 70B model.

Prefill is the step that processes a whole prompt before generation starts. Its work is the per-token work times the number of tokens. A 1,000-token prompt on a 7B model needs about 14 TFLOP.

This estimate ignores the work attention does over the context, which grows as the context grows (lesson 9, Example 12). It is still accurate enough to predict the speed of prefill on a machine with a known FLOP rate (lesson 17).

Example 15: summation order changes float results

Floating-point addition is not associative. That means (a + b) + c can differ from a + (b + c). The reason is that every add rounds the result to the nearest float, and a float keeps only about 7 decimal digits. Adding a small number to a large one drops the small number's low digits.

A consequence follows. If a program sums the same numbers in a different order, it can return a slightly different result. For example, a loop on one core and a loop split over four cores add in different orders. Both results are correct within rounding. This is the first source of nondeterminism (different output from the same input) in inference. Lesson 12 returns to it when it asks why two servers with the same seed can still produce different text.

The program sums one million float numbers in three orders and compares each with a double sum, which is accurate enough to serve as the reference. The last line shows the effect in its smallest form.

import java.util.Locale;
import java.util.Random;

public class SummationOrder {
    public static void main(String[] args) {
        Random rnd = new Random(1);
        float[] x = new float[1_000_000];
        double exact = 0;
        for (int i = 0; i < x.length; i++) {
            x[i] = (float) (rnd.nextGaussian() * 0.01 + 0.001);
            exact += x[i];
        }
        float forward = 0;
        for (int i = 0; i < x.length; i++) forward += x[i];
        float backward = 0;
        for (int i = x.length - 1; i >= 0; i--) backward += x[i];
        float[] acc = new float[4];
        for (int i = 0; i < x.length; i++) acc[i % 4] += x[i];
        float fourLanes = (acc[0] + acc[1]) + (acc[2] + acc[3]);
        System.out.println(String.format(Locale.ROOT, "exact (double)    = %.6f", exact));
        System.out.println(String.format(Locale.ROOT, "forward           = %.6f  error %+.6f", forward, forward - exact));
        System.out.println(String.format(Locale.ROOT, "backward          = %.6f  error %+.6f", backward, backward - exact));
        System.out.println(String.format(Locale.ROOT, "4 accumulators    = %.6f  error %+.6f", fourLanes, fourLanes - exact));
        System.out.println("forward == backward: " + (forward == backward) + ", forward == 4 lanes: " + (forward == fourLanes));
        float big = 1e8f;
        System.out.println("(1e8f + 3f) + 3f = " + ((big + 3f) + 3f) + ",  1e8f + (3f + 3f) = " + (big + (3f + 3f)));
    }
}

exact (double)    = 1004.090157
forward           = 1004.082092  error -0.008064
backward          = 1004.077026  error -0.013130
4 accumulators    = 1004.083862  error -0.006294
forward == backward: false, forward == 4 lanes: false
(1e8f + 3f) + 3f = 1.0E8,  1e8f + (3f + 3f) = 1.0000001E8

The double sum is the reference, 1004.090157. Summing the million float numbers forward gives 1004.082092. Summing backward gives 1004.077026. Summing with four interleaved accumulators gives 1004.083862. These are three different results, each within about 0.013 of the reference (0.001 percent). None is wrong, and none equals another.

The last line shows the effect in its simplest form. Adding 3 to 100,000,000 in float does nothing, because the gap between neighboring floats near 1e8 is 8, and 3 is less than half of that gap. Adding 3 + 3 = 6 first and then adding it to 1e8 does move the result, to 100,000,008.

A multi-threaded or vectorized engine sums in a different order than a single-threaded one, so its logits differ in the last digits. Usually this changes nothing. When two tokens have almost equal scores, it can change which one is picked.

Interview points

• A dot product multiplies the matching numbers of two vectors and adds the products. It is large when the vectors point the same way. Cosine similarity divides by both lengths, so it compares direction only.
• A matrix-vector product with R rows and C columns is R dot products and costs 2 * R * C FLOPs. Decoding one token is mostly such products with the weight matrices.
• Row-major flat storage puts element (i, j) of a matrix with cols columns at index i * cols + j. Contiguous rows make sequential reads fast.
• A linear layer y = W x + b has out * in + out parameters. Many transformer layers have no bias.
• Without an activation function, stacked linear layers collapse into one linear layer. SiLU and GELU are smooth alternatives to ReLU, and a model must run with the function it was trained with.
• Softmax does not change when a constant is added to all logits, so engines subtract the maximum logit to avoid overflow.
• Log-softmax computes z_i - max - log(sum(exp(z_j - max))) directly. This keeps tiny probabilities representable and gives the cross-entropy loss -log p_target.
• An embedding lookup selects one row of the table. It equals a one-hot vector times the embedding matrix, without the wasted multiplications.
• LayerNorm subtracts the mean and divides by the standard deviation. RMSNorm divides only by the root mean square and has no mean subtraction, so it needs fewer operations.
• A model with P parameters needs about 2 * P FLOPs per token, ignoring attention over the context.
• Float addition is not associative, so the summation order (threads, SIMD lanes, tile sizes) changes the last digits of a result.
• Matching logits between two engines means matching within a small tolerance, and not bit for bit.

Summary

• A vector is a float[]. Adding and scaling work index by index.
• dot(a, b) = sum(a_i * b_i), length(a) = sqrt(dot(a, a)), and cosine = dot / (length(a) * length(b)), which lies between -1 and 1.
• An embedding assigns each token a vector. Similar tokens have a high cosine similarity.
• A matrix times a vector gives one dot product per row. The number of columns must equal the length of the vector.
• Matrices are stored flat in row-major order. Element (i, j) is at index i * cols + j.
• A linear layer is y = W x + b with out * in + out parameters.
• ReLU is max(0, x). SiLU is x * sigmoid(x). GELU is x * Phi(x), often computed with a tanh approximation.
• Softmax is exp(z_i) / sum(exp(z_j)). Subtracting the maximum logit gives the same result without overflow. The logits {1000, 1001, 1002} give NaN without it.
• Cross-entropy for one prediction is -log p_target, computed with log-softmax.
• An embedding lookup selects row t of a V x d table, which has V * d parameters.
• LayerNorm is (x - mean) / sqrt(var + eps). RMSNorm is x / sqrt(mean(x^2) + eps). Both are followed by a learned weight.
• A matrix-vector product costs 2 * rows * cols FLOPs. A model costs about 2 * parameters FLOPs per token.
• Summing one million floats in three different orders gave three different results, because float addition is not associative.