C: softmax and why the largest score is subtracted

Convert scores to probabilities and show that the naive formula overflows for large scores.

#include <math.h>
#include <stdio.h>

static void softmax_naive(const double *z, double *out, int n) {
  double total = 0.0;
  for (int i = 0; i < n; i++) {
    out[i] = exp(z[i]);
    total += out[i];
  }
  for (int i = 0; i < n; i++) {
    out[i] /= total;
  }
}

static void softmax_stable(const double *z, double *out, int n) {
  double largest = z[0];
  for (int i = 1; i < n; i++) {
    if (z[i] > largest) {
      largest = z[i];
    }
  }
  double total = 0.0;
  for (int i = 0; i < n; i++) {
    out[i] = exp(z[i] - largest);
    total += out[i];
  }
  for (int i = 0; i < n; i++) {
    out[i] /= total;
  }
}

int main(void) {
  double small[3] = {2.0, 1.0, 0.1};
  double out[3];
  softmax_stable(small, out, 3);
  printf("stable: %.3f %.3f %.3f\n", out[0], out[1], out[2]);
  softmax_naive(small, out, 3);
  printf("naive:  %.3f %.3f %.3f\n", out[0], out[1], out[2]);

  double big[2] = {1000.0, 1001.0};
  double big_out[2];
  softmax_naive(big, big_out, 2);
  printf("naive on large scores is NaN: %s\n", isnan(big_out[0]) ? "yes" : "no");
  softmax_stable(big, big_out, 2);
  printf("stable on large scores: %.3f %.3f\n", big_out[0], big_out[1]);
  return 0;
}

Output

stable: 0.659 0.242 0.099
naive:  0.659 0.242 0.099
naive on large scores is NaN: yes
stable on large scores: 0.269 0.731

Explanation

Softmax turns a vector of scores into probabilities: each output is exp(z_i) divided by the sum of exp(z_j). The outputs are positive and sum to 1. A network usually ends in softmax when it must choose one of several classes.

exp grows very fast. In double precision exp(1000) overflows to infinity, and infinity divided by infinity is NaN. Because softmax(z) equals softmax(z − c) for any constant c, subtracting the largest score first keeps every exponent at or below zero and gives the same probabilities. The naive and stable versions agree for small scores and only the stable one survives large ones.

Explore the dry run

Guided walkthroughReady

1#include <math.h>

2#include <stdio.h>

4static void softmax_naive(const double *z, double *out, int n) {

5 double total = 0.0;

6 for (int i = 0; i < n; i++) {

7 out[i] = exp(z[i]);

8 total += out[i];

9 }

10 for (int i = 0; i < n; i++) {

11 out[i] /= total;

12 }

13}

15static void softmax_stable(const double *z, double *out, int n) {

16 double largest = z[0];

17 for (int i = 1; i < n; i++) {

18 if (z[i] > largest) {

19 largest = z[i];

20 }

21 }

22 double total = 0.0;

23 for (int i = 0; i < n; i++) {

24 out[i] = exp(z[i] - largest);

25 total += out[i];

26 }

27 for (int i = 0; i < n; i++) {

28 out[i] /= total;

29 }

30}

32int main(void) {

33 double small[3] = {2.0, 1.0, 0.1};

34 double out[3];

35 softmax_stable(small, out, 3);

36 printf("stable: %.3f %.3f %.3f\n", out[0], out[1], out[2]);

37 softmax_naive(small, out, 3);

38 printf("naive: %.3f %.3f %.3f\n", out[0], out[1], out[2]);

40 double big[2] = {1000.0, 1001.0};

41 double big_out[2];

42 softmax_naive(big, big_out, 2);

43 printf("naive on large scores is NaN: %s\n", isnan(big_out[0]) ? "yes" : "no");

44 softmax_stable(big, big_out, 2);

45 printf("stable on large scores: %.3f %.3f\n", big_out[0], big_out[1]);

46 return 0;

47}

Follow the example

Start to inspect the source, values, and output. These prepared steps explain the published sample; grouped loop steps describe multiple iterations.

View the complete dry-run table

Prepared trace of the published program. Some rows group repeated iterations.
What happens	State after step	Output so far
Every exponent is at most 0, so exp never overflows.	`—`	—
The probabilities for scores 2.0, 1.0, 0.1.	`—`	stable: 0.659 0.242 0.099
For small scores the naive version gives the same answer.	`—`	stable: 0.659 0.242 0.099 naive: 0.659 0.242 0.099
exp(1000) is infinity, and infinity / infinity is NaN.	`—`	stable: 0.659 0.242 0.099 naive: 0.659 0.242 0.099 naive on large scores is NaN: yes
After subtracting 1001 the scores become −1 and 0, giving the same result as for {0, 1}.	`—`	stable: 0.659 0.242 0.099 naive: 0.659 0.242 0.099 naive on large scores is NaN: yes stable on large scores: 0.269 0.731

Cost and common mistakes

O(n) time for n scores.

Link with -lm on Linux, or the call to exp fails to link.
Printing a NaN gives platform-specific text such as nan or -nan. Test with isnan instead of printing it.
Training code usually combines softmax and the logarithm into one stable log-softmax.

Try it yourself

What does softmax give for equal scores?

Reveal answer and reasoning

Equal probabilities: with four equal scores each probability is 0.25.

C: softmax and why the largest score is subtracted

Explanation

Follow the example

Cost and common mistakes

Try it yourself

Further reading