Processing Scores in Chunks with Online Softmax

In the previous article, we explored why a subset of scores cannot determine the final softmax probabilities. The denominator must also account for the remaining scores in the same row. This time, we will explore online softmax, a way to accumulate that information while reading scores a few at a time.

First, we will see why subtracting the maximum from every score makes computation stable even when the scores are large. Next, we will examine the roles of the maximum and the sum of exponentials. Finally, we will follow how each new chunk of scores either leaves the maximum unchanged or updates it, and how we rescale and add to the exponential sum using that reference. The key is to update these two values for the scores processed so far without revisiting all the earlier scores.

Subtracting the Maximum for Stable Computation

Let us briefly revisit softmax. We first apply the exponential function exp to each score to obtain positive values. We then add them all and divide each value by that sum to obtain probabilities that sum to one. For example, if the scores are [0, 0], both exponentials are 1 and their sum is 2, so each probability is 50%.

As scores grow, however, their exponentials grow very quickly. Applying exp directly to the scores [1000, 1001, 1002] in Figure 1 exceeds the representable range of FP32 or FP64. FP32 and FP64 are floating-point formats that represent real numbers using 32 and 64 bits, respectively. Even when the probabilities are mathematically defined, a computer cannot obtain valid results if it cannot represent intermediate values.

To avoid this, we subtract the same maximum from every score before applying exp. In Figure 1, subtracting the maximum of 1002 gives [-2, -1, 0]. Their exponentials are approximately [0.135, 0.368, 1], with a sum of about 1.503. Dividing each value by this sum gives approximately [9.0%, 24.5%, 66.5%]. Decimals in this article and its figures are rounded after calculation at full precision.

Why do the probabilities stay the same when the scores change? If a score is x and the maximum subtracted from every score is m, then exp(x − m) = exp(x) × exp(−m). Every exponential receives the same positive scale factor, so that factor applies both to each individual exponential in the numerator and to the full sum in the denominator. It cancels in the division, preserving the probabilities. This is a mathematical relationship; in the actual computation, we subtract m from the score and then apply exp, without first producing the original large exponentials.

After subtracting the maximum, all scores are zero or negative, and a score equal to the maximum becomes zero. The largest exponential is therefore exp(0) = 1, avoiding the overflow that can occur when applying exp to a large positive number. Here we consider a row of finite scores. This does not eliminate underflow, where very small exponentials round to zero, or all floating-point rounding errors. The numerically stable softmax formula

The Maximum and Exponential Sum for Normalization

Stable softmax computation needs two values: the maximum m, which sets the reference for computing exponentials, and the exponential sum ℓ, obtained by summing the exponentials of scores after subtracting that m. The symbol ℓ is a distinct form of the lowercase letter l.

In the previous article, we used the sum of the original scores’ exponentials as the denominator. From here on, the exponential sum ℓ means the sum of exponentials after subtracting the maximum. Changing the maximum changes the exponential values and their sum even for the same scores, so we must always think of m and ℓ as a pair.

The scores in Figure 2 are [1, 2]. The maximum m is 2, and the exponentials are exp(1 − 2) ≈ 0.368 and exp(2 − 2) = 1. The exponential sum ℓ is therefore about 1.368. The first probability is approximately 0.368 divided by that sum, or 26.9%; the second is 1 divided by that sum, or 73.1%.

In this example, the two scores make up the entire row, so these are the final probabilities. But if [1, 2] is only the first part of a longer row, the current m and ℓ describe only that part. Later scores may increase the maximum and contribute more values to the exponential sum. To obtain probabilities for the full row, we must incorporate the remaining scores and complete both values.

Updating Both Values with Each New Chunk

Now let us read the row [1, 2, 0, 1, 3, 2] two scores at a time. After processing the first chunk, [1, 2], we have m = 2 and ℓ ≈ 1.368, as in Figure 2. We keep these two values and move to the next chunk.

When we read a new chunk, we first choose the larger of the previous maximum and the new chunk’s maximum as our new reference. We then adjust the existing exponential sum to that reference and add the new scores’ exponentials computed with the same reference. Figure 3 shows a case where the maximum stays the same, followed by one where it increases.

When the Maximum Stays the Same

The maximum in the second chunk, [0, 1], is 1. Because it is smaller than the previous maximum of 2, m stays at 2. The existing exponential sum already uses this reference, so we can keep it unchanged.

We subtract the same m of 2 from the new scores before computing their exponentials. Adding exp(0 − 2) + exp(1 − 2) ≈ 0.503 to the previous sum of about 1.368 gives a new exponential sum of about 1.871. Now m = 2 and ℓ ≈ 1.871 account for all four scores read so far: [1, 2, 0, 1].

The important point is that the new chunk’s exponentials also use the maximum so far, 2. If we subtracted only the new chunk’s own maximum of 1 and added that exponential sum directly to the previous sum, the two sums would use different references. We are accumulating exponentials computed with the same reference, rather than forming probabilities separately for each chunk and concatenating them.

When the Maximum Increases

In the third chunk, [3, 2], we encounter a larger score, 3. We must update m from 2 to 3. But the previous exponential sum of about 1.871 was computed by subtracting the old reference of 2. Before adding the new scores’ sum, we must convert the previous sum to the reference obtained by subtracting the maximum of 3.

For an earlier score x, the two references are related as follows.

$$\begin{array}{l}\exp (x-3)\\ =\exp (x-2)\times \exp (2-3)\end{array}$$

Because we subtract a maximum that is one larger, we can multiply each earlier exponential by exp(−1) ≈ 0.368. This scale factor is the same for all earlier scores. Therefore, we can multiply the single accumulated exponential sum by that factor, without rereading and recomputing the earlier scores one by one. Multiplying the previous sum by exp(2 − 3) gives about 0.688.

We then process the new scores [3, 2] using the new maximum of 3. Their contribution is exp(3 − 3) + exp(2 − 3) ≈ 1.368. Adding this to the rescaled previous sum of about 0.688 gives a final exponential sum of about 2.056. The resulting m = 3 and ℓ ≈ 2.056 correspond to the maximum and exponential sum obtained by processing all six scores together from the start.

The case where the maximum stays unchanged follows the same rule. When the previous and new maxima are equal, the scale factor is exp(0) = 1, so the previous sum remains unchanged. A single rule handles both cases: choose the new maximum, rescale the previous sum, then add the new scores’ exponential sum. The online normalization update

Computing Final Probabilities from the Accumulated Values

After reading all chunks, we know the maximum and exponential sum for the entire row. For example, the final probability for the first score, 1, is exp(1 − 3) ÷ 2.056 ≈ 6.6%. Unlike the 26.9% obtained from [1, 2] alone, this probability accounts for the other four scores as well.

Completing m and ℓ does not, by itself, output the probabilities for every position. Each probability still requires applying the final m and ℓ to its corresponding score. If we accumulated only the statistics without retaining the scores, we must be able to reread or recompute the scores when producing the probabilities.

Online softmax updates this normalization information for each chunk of scores. Because it can rescale the previous exponential sum to a new reference when the maximum increases, it can obtain a shared reference and denominator for an entire row without keeping all scores in memory at once.