Image by Editor

Modern AI systems like large language models (LLMs) make countless atomic decisions per second in the process of delivering outputs. At their core, though, they are nothing more than spectacular predictive — or guessing, if you prefer — machines. They predict the most likely next word to add to a response being generated, based on the words provided so far.

This article takes a gentle look at the math behind how these models actually make such choices. Before confidently generating a word, an LLM generates a set of “raw scores” (called logits) for every possible word, and one key challenge is turning these messy scores into clean, interpretable probabilities. This is what the softmax function does: let’s see how and why softmax is a core mechanism used in both classical machine learning models and state-of-the-art solutions.

How to Make Sense of Logits

Suppose an AI system is trying to complete the classic sentence: “The cat sat on the…“. Under the hood, there is a neural network architecture consisting of many layers and computation-intensive components. In the case of LLMs, the arduous work done by each of these elements culminates in calculating thousands of raw scores, or logits: one for every potential next word to generate. For simplicity, imagine this scenario with just three possible words or outputs to choose from to complete the previous example sentence:

• • “mat”: 4.2
  • “floor”: 1.5
  • “cloud”: -2.0

You may agree that, from a statistical viewpoint, these numbers look so raw and untidy that they might be a bit of a headache. First, their sum is not 1 or 100%. Second, they can be negative, as you can observe. Third, there is no clear sense of how much more likely “mat” is compared to “floor”, for instance, because the scale on which scores move is not bounded.

How can we convert these logits into a standard probability distribution where values lie between 0 and 1 and together add up to 1? With the softmax function, of course!

How Softmax Works: A Two-Step Transformation

Softmax brilliantly accomplishes the required transformation from messy logits to interpretable probabilities using two mathematical steps, neither of which is particularly complex. Indeed, you can easily implement them in a spreadsheet, for instance.

In the first step, the aim is to turn all scores positive. To do so, we use Euler’s number e, which is approximately 2.718, and raise it to the power of each of our scores x. Mathematically, raising e to any real-valued exponent, including a negative one, results in a positive number. Furthermore, doing this helps stretch the scale, especially when there are large differences between any two or more logits. Here is how to apply this to our three raw logits:

“mat” (4.2) $\rightarrow$ $e^{4.2} \approx 66.69$
“floor” (1.5) $\rightarrow$ $e^{1.5} \approx 4.48$
“cloud” (-2.0) $\rightarrow$ $e^{-2.0} \approx 0.14$

It may look like the differences between values have intensified, but wait until what comes next.

In the second step, called normalization, we turn these exponentials into slices of a pie or proportions, so that they all sum to 1 or 100%. The way to do this is to add up all the exponential values and use this sum as a denominator. Thus, for each exponential, we divide it by this sum, and voilà: we now have values that can be perfectly interpreted as probabilities!

In our cat example, the total sum of exponentials is: 66.69 + 4.48 + 0.14 = 71.31

Now, we divide each exponential by this total:

“mat”: 66.69 / 71.31 = 0.935 (93.5%)
“floor”: 4.48 / 71.31 = 0.063 (6.3%)
“cloud”: 0.14 / 71.31 = 0.002 (0.2%)

Why “Softmax”?

The reason softmax has this name is related to how it behaves. It normally assigns the vast majority of the probability to the highest original score, but keeps giving some chance to the runner-up options that follow it. This softness — in contrast with simply picking the option with the highest score and assigning it full probability — is essential for creativity in AI systems like LLMs. These models apply sampling strategies once these probabilities have been calculated with softmax, so that instead of always picking the highest-probability word, they use probability-based sampling. As a result, even though “mat” has the highest chance of being selected, “floor” and even “cloud” could also be chosen, although this is less probable!

LLMs use interesting parameters to control the behavior of this sampling process, such as temperature and top-k, among others. You can read more about them here.