Softmax Optimization

  • Last Updated 14 June, 2026
  • by David Spuler, Ph.D.
The Softmax function is a significant cost in Transformer inference because it is part of the attention mechanism, whereas it was less of a bottleneck in earlier neural network architectures. A vanilla Softmax implementation is very expensive because it involves computing the exponentials of all of the logits. Various attempts have been made to optimize and approximate Softmax calculations, including:
  • Softmax optimizations
  • Softmax approximations
  • Integer-only Softmax
  • Pruned Softmax (removal)
  • Fused Softmax (kernel fusion)
  • Softmax replacements (use different functions)

Note that there are several other areas of theory that are relevant to Softmax optimizations and approximation. The denominator of the Softmax formula is a "sum of exponentials", and this type of calculation also appears in Logarithmic Number System (LNS) addition. Also, the sum of exponentials calculation, appears in "log-sum-exp networks", which are somewhat related to "tropical algebra". The approximations of "max-plus networks" may also be relevant to Softmax approximations.

What is Softmax?

The purpose of Softmax is to take a set of values in a vector of calculated values, and normalize them into probabilities. After Softmax, the output vector contains a new normalized set of values which all add up to 1, and they are intended to represent probabilities of the likelihood of each token/word associated with each vector element.

So why do we need all the exponentials? The idea is that the input vector contains "logits", which are logarithms of probabilities, so we are exponentiating each one to bring it out of log-domain into real-domain. Then we are normalizing them all so that they are probabilities that total exactly 1.

Softmax and Temperature: At the end of each decoder sequence, the Softmax function is used to normalize the logits before processing by a decoding algorithm to choose an output token with the highest probability. As part of this method, the Softmax function is usually changed to a "scaled Softmax" that uses a parameter called the "temperature".

What is the temperature? The purpose of the temperature parameter is as a hyper-parameter that influences the level of randomness or unpredictability in the output. A higher setting for temperature means that the decoder is more likely to output the lower-probability tokens. If the temperature is very large, the decoder is mostly going to output the highest probability token, meaning it is much less random.

What is the value of the temperature? The temperature is a non-zero positive floating point number that can be between 0 and 1, or can also be greater than 1. A temperature of zero cannot be used as it would cause divide-by-zero errors. If the temperature equals 1.0, it doesn't change the Softmax function at all (i.e. continues harmlessly without scaling). Since the Softmax function is scaled by the reciprocal of the temperature, the effect is to make randomness higher with a larger temperature setting (so it runs "hotter" and gets more "bubbly"). If the temperature is below 1.0, making it a fraction, the effect is to spread out the logits more, which has the effect of reducing randomness of the output. If the temperature is greater than 1.0, this contracts the logits towards each other, making the decoder more likely to choose each of them (although still with some randomness), thereby increasing output randomness. Read more about decoding algorithms and temperature.

Softmax Optimizations: Book Excerpts and Blog Articles

Free online book excerpts with full text chapters online and free PDF downloads, and the Aussie AI blog, including related articles:

Softmax Optimization

Example: Basic Softmax in C++: The Softmax function is an inefficient computation by default. First, the whole vector is scanned to compute the sum of the exponential of each value. Then the whole vector is re-scanned to divide each vector element by this sum-of-exponentials.

Here is a naive implementation of Softmax in C++: 

    #include <math.h>  // Declare expf()
    ....
    float yapi_vector_sum_of_exponentials(float v[], int n)
    {
	float sum = 0.0;
	for (int i = 0; i < n; i++) {
		float e = expf(v[i]);
		yassert(e > 0.0);
		sum += e;
	}
	return sum;
    }

    void yapi_vector_softmax_basic(float v[], int n)
    {
	float denom = yapi_vector_sum_of_exponentials(v, n);  // Denominator...
	if (denom == 0.0) {
		yassert(denom != 0.0);
		return;  // fail (should not occur)
	}
	for (int i = 0; i < n; i++) {
		v[i] = expf(v[i]) / denom;
	}
    }

Not only is this computation very slow (all those calls to exp!), but it's also prone to overflow and underflow. The real computation of Softmax needs to be further optimized and scaled to avoid these problems. Further optimizations may include the use of API calls to hardware acceleration APIs, pre-computed tables as approximations for the "exp" function, and converting the loop to pointer arithmetic.

Softmax Optimization

Research papers on Softmax optimizations in general:

Softmax Approximation

Softmax can be approximated in various ways, and there are also integer-only Softmax methods. Research papers on approximating the Softmax algorithm include:

Integer-Only Softmax

Softmax can be optimized by conversion to integer arithmetic. See also the overview of integer arithmetic in neural networks.

Research papers with an integer-only Softmax function:

Softmax Pruning

In addition to approximating Softmax, consideration has been given to "pruning" the Softmax component, i.e. removing it completely.

Softmax Alternatives

There is also research on replacing Softmax with a different component:

More AI Research

Read more about: