trendingzones

MATHS FOR AI/ML — INTRODUCTION

Probability & Softmax: How an LLM Actually Picks the Next Word

Every time a language model produces one word, it has already scored every word in its entire vocabulary — tens of thousands of candidates — and picked one. This page is about exactly how it goes from “a big list of raw scores” to “one chosen word,” using real arithmetic instead of hand-waving.

The Quick Answer

A language model’s raw output for the next word is a list of numbers called logits — one score per possible word, and they can be any real number, including negative ones. A formula called softmax converts that list of logits into a list of real probabilities: every value becomes positive, and the whole list adds up to exactly 100%. Once the model has that probability distribution, a sampling strategy — always taking the top choice, or picking randomly weighted by probability — decides which single word actually gets produced.

What’s a Logit, and Why Isn’t It Already a Probability?

As covered in How LLMs Work, each forward pass through a language model produces one raw score per vocabulary entry — a logit. These numbers come straight out of the model’s final layer, and there’s nothing that forces them to behave like probabilities: they can be any size, and they can be negative. A logit of −3 for “banana” and a logit of 5 for “the” just mean “the model currently favors ‘the’ a lot more than ‘banana’” — nothing about those numbers sums to anything meaningful yet.

To turn that list of arbitrary scores into something you can actually treat as “the probability of each word,” you need a function with two guarantees: every output has to be positive (probabilities can’t be negative), and every output has to sum to 1 across the whole vocabulary. That function is softmax.

Softmax, Step by Step

The softmax formula does exactly two things to a list of logits: it exponentiates every value (raises the mathematical constant e≈ 2.718 to the power of that logit), then divides each exponentiated value by the sum of all of them. Try it on a few small examples below — the exact same arithmetic a model runs on 100,000+ logits at once, just shrunk down to 3 words so every step is visible:

Why Not Just Divide by the Sum?

It might seem simpler to skip the exponentiation step entirely — just take each logit and divide it by the sum of all the logits. That works fine as long as every logit is positive. But real logits are routinely negative, and plain division breaks the moment one shows up: a negative logit divided by a positive sum gives a negative “probability,” which is meaningless. Exponentiation fixes this structurally — e raised to any real number, however negative, is always positive. Compare both methods on the same logits, including a negative one:

From Probabilities to One Word: Greedy vs. Sampling

Softmax hands back a full probability distribution over every candidate word — it doesn’t pick one by itself. That’s the job of a separate sampling strategy. The simplest option, greedy decoding, always takes the single highest-probability word, every time, with zero randomness. The alternative, sampling, treats the distribution as genuinely random: a word with 60% probability gets picked about 60% of the time, not always. Try both below on the same distribution:

This exact choice is why the same prompt to an LLM can produce a different answer every time you ask it (when sampling is on), or the exact same answer every time (when decoding is greedy, or temperature is set to 0) — a behavior covered conceptually in Prompt Engineering.

Restricting the Candidates: Top-k and Top-p

Pure random sampling over the full vocabulary has a real downside: even a word with a tiny 0.1% probability can occasionally get picked, and over a long response those rare picks add up to noticeably worse text. Two common fixes narrow the field before sampling: top-k keeps only the k highest-probability words and ignores the rest, while top-p (also called nucleus sampling) keeps the smallest set of top words whose probabilities add up to at least p. Toggle between them below on the same 6-word distribution:

Fun Fact

OpenAI’s own API documentation describes the temperature parameter — the setting that reshapes this exact probability distribution before sampling — as ranging from 0 to 2, where values like 0.2 make output “more focused and deterministic” and values like 0.8 make it “more random.” Every one of those settings is really just adjusting the same softmax math covered above, before a single word is ever picked.

Source: OpenAI API Reference — Chat Completions

Test Yourself

What does softmax actually do to a list of raw scores (logits)?

Why can’t a language model just divide each raw score by the sum of all scores instead of using softmax?

What is the difference between greedy decoding and sampling?

Brainstorm

These don’t have a single correct answer — they’re here to make you pause and think before checking one way to reason about each one.

If a model’s raw scores for "cat," "dog," and "bird" were 5, 5, and 5 (all exactly equal), what would you expect the resulting probabilities to look like, and why?

Why might a chatbot that always uses greedy decoding start repeating the same phrase over and over in a long response?

Two candidate next words have probabilities of 60% and 40% after softmax. If you sampled from this distribution 10 times, would you expect exactly 6 and 4? Why or why not?

Why do you think softmax uses exponentiation (exp) specifically, instead of some other way to make all values positive, like just taking the absolute value of each logit?

If you were designing a customer support chatbot, would you want it to use greedy decoding, sampling, or something else? What would you want to weigh in that decision?

Ready for the actual temperature formula (dividing logits before softmax runs), a real entropy calculation, and why top-k and top-p can disagree on which words survive? Continue to the Intermediate level →