Probability, Logs and Softmax

Table of Contents

1. A Prediction Is a Distribution
2. Expected Value: The Average You Would See
3. Why Logs: Products Become Sums
4. Surprise and Entropy
5. Softmax: From Scores to Probabilities
6. The Overflow Trap and the Max Trick
7. Temperature Is Division Before Softmax
8. Cross-Entropy: Punishing a Wrong Guess
9. Perplexity in One Paragraph
10. Putting It Together: One Step of Next-Word Training
11. Recipe: The stats Module
12. Try It Yourself
13. Quick Review: Common Questions

1. A Prediction Is a Distribution

In the last chapter the model predicted a number: an exam score. A language model predicts something different. After the words "I drink", there is no single right next word. Tea is likely, coffee is quite likely, water is possible, and "brick" is very unlikely. So the model does not answer with one word. It answers with a probability for every word it knows:

after "I drink":

  tea     ##########  0.5
  coffee  ######      0.3
  water   ####        0.2
  ...     (every other word, close to 0)

That list of probabilities is called a distribution. It follows two rules you know from Probability Basics: every probability is between 0 and 1, and all of them add up to 1, because something has to come next. Check it in code. The numbers are invented for the example:

// demo-dist.js
import { sum, expectedValue } from "./src/stats.js";

// a made-up next-word distribution after "I drink"
const next = { tea: 0.5, coffee: 0.3, water: 0.2 };
const probs = Object.values(next);
console.log("sum:", sum(probs));

// a fair die: every face has probability 1/6
const faces = [1, 2, 3, 4, 5, 6];
const fair = faces.map(() => 1 / 6);
console.log("die average:", expectedValue(faces, fair));

// a lottery ticket: win 100 with probability 0.01, otherwise 0
console.log("ticket worth:", expectedValue([0, 100], [0.99, 0.01]));

// output
sum: 1
die average: 3.5
ticket worth: 1

All three lines are stable. This chapter uses src/stats.js, which you will save in the recipe at the end. In the next chapters you will see the same file again; the mean, variance and std functions in it are used in the last chapter.

Say it out loud: "A language model does not output one word. It outputs a probability for every word, and they add up to 1."

2. Expected Value: The Average You Would See

The last two lines above used expected value. It answers: if I could repeat this many times, what would the average result be? Multiply each outcome by its probability, then add:

expected value = sum of (value * probability)

A fair die gives (1 + 2 + 3 + 4 + 5 + 6) / 6 = 3.5, even though you can never roll a 3.5. The expected value need not be a possible outcome. It is the long-run average. The lottery ticket pays 100 with chance 0.01, so on average it is worth 1. If the ticket costs 2, you lose on average.

Why does this matter for AI? A loss is an expected value. When we train, we do not care about one example. We care about the average loss over the examples the model will see. The mean squared error in the last chapter was exactly an average of errors. Everything below ends in an average like that.

3. Why Logs: Products Become Sums

Logarithms scare people, but in AI they do one humble job: they turn multiplication into addition. In Functions, Growth and Logarithms you saw that log(a * b) = log(a) + log(b). We use the natural log, Math.log, throughout this chapter.

Why would we want that? Because probabilities of several things happening together are multiplied. The chance that a model gives the right word at each of 400 positions is p1 * p2 * ... * p400. Each number is below 1, so the product gets tiny, fast. Watch:

// demo-logs.js
// a log turns a product into a sum
console.log(Math.log(0.5 * 0.5));
console.log(Math.log(0.5) + Math.log(0.5));

// multiply 400 small probabilities of 0.1
let product = 1;
for (let i = 0; i < 400; i++) product *= 0.1;
console.log("product:", product);

// add up their logs instead
let logSum = 0;
for (let i = 0; i < 400; i++) logSum += Math.log(0.1);
console.log("log sum:", logSum.toFixed(3));

// a few values of the natural log
console.log(Math.log(1), Math.log(0.5).toFixed(4), Math.log(0.01).toFixed(4), Math.log(0));

// output
-1.3862943611198906
-1.3862943611198906
product: 0
log sum: -921.034
0 -0.6931 -4.6052 -Infinity

Four things to see:

1. The first two lines are equal: the log of a product is the sum of the logs.
2. The real answer to the 400-fold product is 10^-400, far smaller than the smallest number a JavaScript float can hold. The computer rounds it to 0. This is called underflow, and the information is gone.
3. The sum of logs is -921.034: an ordinary number that is perfectly safe. Real models multiply thousands of probabilities, so they work with logs, never with the raw product.
4. Math.log of a number between 0 and 1 is negative. log(1) is 0, log(0.5) is -0.69, and log(0) is -Infinity. Probability close to 1 gives log close to 0. Probability close to 0 gives a large negative log.

All of this output is stable. Logs also help in a second way: a log squeezes huge ranges. The difference between probabilities 0.01 and 0.001 is tiny as plain numbers but clear on a log scale (-4.6 against -6.9), which is why a model that confidently gets something wrong is punished so hard, as you will see in section 8.

Say it out loud: "We use logs because they turn multiplying many small probabilities into adding, and because the product would be too tiny for a computer to store."

4. Surprise and Entropy

Imagine two weather forecasters. In a desert, the forecast "no rain" is right almost every day. When it finally rains, you are very surprised. In a rainy season it is a coin flip every day, so you are never very sure what comes next.

We can put a number on surprise: the surprise of an event with probability p is -log(p). A likely event (p near 1) has surprise near 0. A rare event has big surprise. Entropy is the average surprise of a distribution, weighted by how likely each outcome is. It measures how uncertain the distribution is:

low entropy:  tea 0.97  a 0.01  b 0.01  c 0.01   (almost sure)
high entropy: tea 0.25  a 0.25  b 0.25  c 0.25   (no idea)

The module's entropy uses log2, so the answer is in bits. A fair coin is exactly one bit: one yes-or-no question.

// demo-entropy.js
import { entropy } from "./src/stats.js";

const show = (name, probs) => console.log(name.padEnd(22), entropy(probs).toFixed(4), "bits");

show("fair coin", [0.5, 0.5]);
show("coin that is 90% heads", [0.9, 0.1]);
show("always heads", [1, 0]);
show("fair 4-sided die", [0.25, 0.25, 0.25, 0.25]);
show("one clear favourite", [0.97, 0.01, 0.01, 0.01]);

// output
fair coin              1.0000 bits
coin that is 90% heads 0.4690 bits
always heads           0.0000 bits
fair 4-sided die       2.0000 bits
one clear favourite    0.2419 bits

Read it from top to bottom. A fair coin is one bit. A biased coin is less uncertain, so less than one bit. A coin that is always heads has no uncertainty at all, zero bits. Four equally likely outcomes need two bits (two yes-or-no questions). A distribution with one clear favourite is nearly decided, so it is low. The rule of thumb: the flatter the distribution, the higher the entropy. All of these outputs are stable.

One honest detail: if you call entropy([1, 0]) without .toFixed, JavaScript may print -0, a harmless negative zero left over from the arithmetic.

The connection to language models: when a model is unsure, its next-word distribution is flat and has high entropy. When it is confident, the distribution is spiky and has low entropy. The temperature setting you met in Probabilities, Temperature and Sampling slides a model between those two states. You will see how in section 7.

5. Softmax: From Scores to Probabilities

A model's last layer does not output probabilities. It outputs raw scores for every word, one number each, called logits. They can be any numbers: positive, negative, large, small. They do not add up to anything. We need to turn them into a proper distribution.

The recipe is called softmax:

1. Take e to the power of each score (Math.exp). This makes every number positive, and larger scores much larger.
2. Divide each result by the total, so they add up to 1.

scores [2, 1, 0]
exp    [7.389, 2.718, 1.000]    total = 11.107
probs  [0.665, 0.245, 0.090]

The biggest score gets the biggest share, and every word keeps a chance, which is the "soft" in softmax (a "hard" max would give the winner everything).

// demo-softmax.js
import { softmax, sum } from "./src/stats.js";

const scores = [2, 1, 0];
const probs = softmax(scores);
console.log(probs.map((p) => p.toFixed(4)));
console.log("sum:", sum(probs));

// adding the same number to every score changes nothing
console.log(softmax([1, 2, 3]).map((p) => p.toFixed(4)));
console.log(softmax([101, 102, 103]).map((p) => p.toFixed(4)));

// output
[ '0.6652', '0.2447', '0.0900' ]
sum: 0.9999999999999999
[ '0.0900', '0.2447', '0.6652' ]
[ '0.0900', '0.2447', '0.6652' ]

Notice three things. The probabilities match the hand calculation. The sum prints 0.9999999999999999 instead of 1: floating point rounding, the same kind of gotcha as in chapter 2, and nothing to worry about. And [1, 2, 3] and [101, 102, 103] give identical results. Softmax only cares about the differences between scores, not their size. That fact is the key to the next section. All output here is stable.

6. The Overflow Trap and the Max Trick

Because softmax uses Math.exp, big scores are dangerous. exp grows extremely fast. A float can hold numbers up to about 1.8e308, and exp passes that limit at an input of around 710:

// demo-overflow.js
import { softmax, sum } from "./src/stats.js";

// the "obvious" softmax, without the max trick
function naiveSoftmax(scores) {
  const exps = scores.map((x) => Math.exp(x));
  const total = sum(exps);
  return exps.map((e) => e / total);
}

console.log(Math.exp(709));
console.log(Math.exp(710));
console.log("naive: ", naiveSoftmax([1000, 999, 998]));
console.log("stable:", softmax([1000, 999, 998]).map((p) => p.toFixed(4)));

// output
8.218407461554972e+307
Infinity
naive:  [ NaN, NaN, NaN ]
stable: [ '0.6652', '0.2447', '0.0900' ]

The naive version turns the scores into Infinity, then computes Infinity / Infinity, which is NaN ("not a number"). The whole prediction is garbage, and one NaN in a network spreads to every number it touches.

The fix uses the fact from section 5: softmax does not change when you add the same number to every score. So subtract the largest score first. The biggest becomes 0, so exp(0) = 1, and everything else is a small negative, so exp of it is between 0 and 1. Nothing can overflow:

scores         [1000, 999, 998]
minus the max  [   0,  -1,  -2]
exp            [ 1.000, 0.368, 0.135]   <- all safe numbers

That is why softmax in src/stats.js starts with Math.max(...logits). It is a one-line trick, and every real library does the same, because it is the difference between working and NaN. The numbers above are stable.

Say it out loud: "Softmax turns scores into probabilities. We subtract the biggest score first so the exponentials cannot overflow."

7. Temperature Is Division Before Softmax

Now you can see exactly what temperature does. Before softmax, divide every score by a number T. A small T makes the differences bigger. A large T shrinks them. Same scores as before, four temperatures:

// demo-temperature.js
import { softmax, entropy } from "./src/stats.js";

const scores = [2, 1, 0];
for (const t of [0.5, 1, 2, 5]) {
  const probs = softmax(scores.map((s) => s / t));
  console.log(`T = ${t}:`, probs.map((p) => p.toFixed(3)), "entropy", entropy(probs).toFixed(3));
}

// output
T = 0.5: [ '0.867', '0.117', '0.016' ] entropy 0.636
T = 1: [ '0.665', '0.245', '0.090' ] entropy 1.201
T = 2: [ '0.506', '0.307', '0.186' ] entropy 1.472
T = 5: [ '0.402', '0.329', '0.269' ] entropy 1.566

At T = 0.5 the top word gets 86.7%: sharp, low entropy, so the model nearly always picks the favourite. At T = 5 the three words are close to equal: flat, high entropy, so sampling wanders and gets more varied (and more likely to say something odd). T = 1 leaves the scores alone. Stable output. The temperature chapter explained this as a dial; now you know it is one division. A common rule of thumb: low temperature for facts and code, higher for brainstorming.

8. Cross-Entropy: Punishing a Wrong Guess

Last chapter we needed a loss for predicting a score. Now we need a loss for predicting a word. The training text gives us the right answer: the word that really came next. The model gave that word some probability p. We want p to be high.

The loss is simply:

loss = -log(p)          where p = probability the model gave to the right answer

This is cross-entropy for one example. It uses exactly the surprise from section 4: "how surprised was the model by the truth?" See how it behaves:

// demo-cross-entropy.js
import { softmax, crossEntropy } from "./src/stats.js";

const probs = softmax([2, 1, 0]); // 0.665, 0.245, 0.090
for (let correct = 0; correct < 3; correct++) {
  console.log(`right answer is #${correct}: loss =`, crossEntropy(probs, correct).toFixed(4));
}

console.log("confident and right:", crossEntropy([0.99, 0.005, 0.005], 0).toFixed(4));
console.log("confident and wrong:", crossEntropy([0.01, 0.5, 0.49], 0).toFixed(4));
console.log("certain and wrong:  ", crossEntropy([0, 1], 0));
console.log("pure guess, 4 options:", crossEntropy([0.25, 0.25, 0.25, 0.25], 0).toFixed(4));
console.log("ln(4):                ", Math.log(4).toFixed(4));

// output
right answer is #0: loss = 0.4076
right answer is #1: loss = 1.4076
right answer is #2: loss = 2.4076
confident and right: 0.0101
confident and wrong: 4.6052
certain and wrong:   Infinity
pure guess, 4 options: 1.3863
ln(4):                 1.3863

The rules fall out of the numbers:

• High probability on the truth gives low loss. 0.99 on the right word costs only 0.0101. Perfect confidence (1.0) would cost 0.
• Low probability on the truth gives high loss. Giving the right word just 0.01 costs 4.6052, over 400 times worse.
• Certainty in the wrong answer is infinitely bad. A probability of exactly 0 gives -log(0), which is Infinity. Real systems keep probabilities away from 0 for this reason.
• A pure guess has a known cost: with 4 equally likely options the loss is ln(4) = 1.3863. A model that cannot beat that is learning nothing.

Notice that cross-entropy only looks at the probability of the right word. It does not matter how the wrong ones are spread. The model is rewarded for pushing probability onto the truth, and softmax guarantees that pushing one up pushes the others down.

This is also why logs earn their place: taking -log turns the product of probabilities across a whole text into a sum of per-word losses. Average those, and you have the single loss number a model reports. It is the same loop as the last chapter, with a different loss function: gradient descent nudges the model's numbers so the average of -log(p) falls. All output here is stable.

The module's entropy is in bits (log2) but crossEntropy uses the natural log. Both are fine, they differ by a fixed constant factor, and you only compare numbers measured the same way.

Say it out loud: "The loss for next-word prediction is minus the log of the probability the model gave to the right word. Right and confident is near zero; wrong and confident is huge."

9. Perplexity in One Paragraph

Loss numbers like 1.3863 are hard to feel. Perplexity turns them into something friendlier: take e to the power of the average cross-entropy. The result reads as "how many equally likely options the model is effectively choosing between." A perplexity of 4 means the model is as unsure as if it were picking from 4 equal choices at every step. Lower is better.

// demo-perplexity.js
import { crossEntropy, mean } from "./src/stats.js";

// probability each model gave to the real next word, over a test text
const perplexity = (probsOfTruth) =>
  Math.exp(mean(probsOfTruth.map((p) => crossEntropy([p], 0))));

console.log("guessing 1 of 4:", perplexity([0.25, 0.25, 0.25]).toFixed(3));
console.log("good model:     ", perplexity([0.9, 0.8, 0.95]).toFixed(3));
console.log("shaky model:    ", perplexity([0.5, 0.1, 0.02, 0.3]).toFixed(3));

// output
guessing 1 of 4: 4.000
good model:      1.135
shaky model:     7.598

(crossEntropy([p], 0) is just a shortcut for -log(p).) The pure guesser scores exactly 4, the good model is close to 1 (almost always sure and right), and the shaky model is around 7.6. The probabilities here are invented. Perplexity is a common way to compare language models on the same text. It is only meaningful when both models are measured on the same text with the same vocabulary, so be careful comparing numbers from different setups. Stable output.

10. Putting It Together: One Step of Next-Word Training

Here is one training example in plain steps. It joins every idea in the last two chapters:

text: "I drink ___"        true next word: "tea"

1. model computes scores (logits):     [ 2.0, 1.0, 0.0 ]   (tea, coffee, water)
2. softmax turns them into probs:      [ 0.665, 0.245, 0.090 ]
3. loss = -log(prob of "tea"):         0.4076
4. gradient of the loss w.r.t. every number in the model
5. nudge each number a little against its gradient: loss goes down

Steps 1 to 3 you now have in code. Steps 4 and 5 are chapters 4 and 5. Repeat over billions of words and the probability given to the right word rises on average, which is what "the model got better at predicting text" means. For a real model, the scores come from a deep network with many layers (you will build a tiny one next), and the list is as long as the vocabulary (see Tokens and Tokenizers). That is a simplification we are making honestly: the toy example has three words, a real one has tens of thousands of tokens.

11. Recipe: The stats Module

Save this as src/stats.js. It has no imports:

export const sum = (xs) => xs.reduce((a, b) => a + b, 0);
export const mean = (xs) => sum(xs) / xs.length;
export const variance = (xs) => {
  const m = mean(xs);
  return mean(xs.map((x) => (x - m) ** 2));
};
export const std = (xs) => Math.sqrt(variance(xs));

// Rescale to mean 0, standard deviation 1.
export function standardize(xs) {
  const m = mean(xs);
  const s = std(xs);
  return xs.map((x) => (s === 0 ? 0 : (x - m) / s));
}

// Stable softmax: subtract the max first so exp() never overflows.
export function softmax(logits) {
  const max = Math.max(...logits);
  const exps = logits.map((x) => Math.exp(x - max));
  const total = sum(exps);
  return exps.map((e) => e / total);
}

export const entropy = (probs) =>
  -sum(probs.map((p) => (p > 0 ? p * Math.log2(p) : 0)));

// Loss for one prediction: -log(probability given to the right answer).
export const crossEntropy = (probs, correctIndex) => -Math.log(probs[correctIndex]);

export const expectedValue = (values, probs) =>
  sum(values.map((v, i) => v * probs[i]));

It exports sum, mean, variance, std, standardize (rescales to mean 0 and standard deviation 1, used in the last chapter), softmax (stable, subtracts the max), entropy (in bits), crossEntropy(probs, correctIndex) (natural log, as -log(probs[correctIndex])) and expectedValue(values, probs).

12. Try It Yourself

1. A fair die's entropy. Compute entropy for six equal probabilities of 1/6. You should see about 2.585 bits, and the same as Math.log2(6). Why is that more than the 2 bits of a 4-sided die?
2. Sort by surprise. For the distribution [0.5, 0.25, 0.15, 0.1], print -Math.log2(p) for each. Which outcome is most surprising? Then compute the weighted average and compare it to entropy.
3. Break softmax. Call softmax([-1000, -1001, -1002]) and the naiveSoftmax from demo-overflow.js. Which one fails now, and what does the naive version return? (Hint: exp of a very negative number underflows to 0.)
4. Find your temperature. Use softmax(scores.map(s => s / t)) with scores = [3, 2, 1, 0] and find a t where the top word gets about 50%. A guess-and-check loop is fine.
5. Loss of a whole text. Three predictions, [0.7, 0.2, 0.1], [0.1, 0.8, 0.1] and [0.3, 0.3, 0.4], have the right answers at positions 0, 1 and 2. Compute the average cross-entropy. You should see about 0.4987.
6. Overconfidence. Compare crossEntropy([0.5, 0.5], 0) with crossEntropy([0.99, 0.01], 1). Which model deserves the bigger penalty, and what does that say about confident wrong answers?

13. Quick Review: Common Questions

Questions beginners often ask about this chapter. Try answering each one yourself first, then click the question to open the answer.

Where to go next. You can now explain how a model's raw scores become probabilities, and how a single number, the cross-entropy, says how wrong the model is. You have also seen every piece of the training loop except the model itself. In the next chapter you build that: a tiny neural network with a few neurons, trained from scratch on a problem a straight line cannot solve: Build a Neural Network.