My question comes from another question answered on Stackoverflow; the Keras implemantion of softmax activation function is customized to subtract the maximum value of the different logits:

def softmax(x, axis=-1):
    # when x is a 2 dimensional tensor
    e = K.exp(x - K.max(x, axis=axis, keepdims=True))
    s = K.sum(e, axis=axis, keepdims=True)
    return e / s

When using Tensorflow with the normal implementation of softmax, it means doing the following operation prior to the softmax call:

logits = logits - tf.expand_dims(tf.reduce_max(logits, axis=-1), -1)

Question: How is it that subtracting the max of the logits helps learning? I don't get it.

Subsidiary question: How come the Keras implement knew this hack? Is there any publication related to this discovery?


1 Answer 1


This is a simple trick to improve the numerical stability. As you probably know, exponential function grows very fast, and so does the magnitude of any numerical errors. This trick is based on the following equality:

$$\frac{e^{x+c}}{e^{x+c}+e^{y+c}} = \frac{e^x e^c}{e^x e^c+e^y e^c} = \frac{e^x e^c}{e^c (e^x+e^y)} = \frac{e^x}{e^x+e^y},$$

where $c$ is the maximum which you are subtracting. As you can see, you can subtract any value without changing the softmax output. Selecting the maximum is a convenient way to ensure numerical stability.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.