Softmax vs the Dirichlet distribution

Question

As far as I understand one can in principle model the distribution over a set of $k$ categories using e.g.:

• the Dirichlet distribution
• A softmax model.

As far as I can tell, both use $k$ parameters to model the distribution over $k$ categories.

My questions:

• Does the "softmax distribution" have a name?
• Why isn't the softmax discussed more frequently in textbooks?
• If the Dirichlet distribution is a common distribution of choice for categorical outcomes, why do NNs use softmax layers instead of Dirichlet ones?

---

Context

Almost very textbook that I have encountered and that introduces common distributions (Binomial, Beta, Gaussian, Gamma, Poisson, Dirichlet) never mentions it (e.g. see this list of common distributions in Wikipedia).

Why isn't it presented more commonly in introductory books as type of common distribution? For almost any common distribution, Wikipedia (for example) covers their moments, marginal distributions, entropy, relationship to other known distributions, etc. Why not the same with the softmax distribution?

This is quite perplexing for me given the fact that softmax layers are actually extremely common in neural networks, and the converse is equally true, I have never encountered neural networks using "Dirichlet layers" for modeling categorical outcomes.

What explains this discrepancy? Is it a matter of tradition? Or is there a deeper connection and equivalency between them, and they just happen to go by different names?

Answer 1

The softmax function is a function. For a fixed input, you get a fixed output, just as for any other function. Someone might have invented a softmax distribution in some context, but that's not the meaning of "softmax function," as described in your post, or commonly used in neural networks.

The Dirichlet distribution is a distribution. Drawing from a Dirichlet distribution with some parameters will almost surely give a different result each time. The Dirichlet distribution is not a distribution over categories, it's a distribution over probability vectors.

The commonality that you've found between the softmax function and the Dirichlet distribution is that the results (function outputs, random draws) are always probability vecotrs: the vectors are non-negative, and sum to 1.

The normal distribution is a probability distribution which assigns positive probability to any real number. A linear function yields real number for any input. But you wouldn't conflate the two, because they're clearly different things in important ways. Likewise, the softmax function and Dirichlet distribution are distinct.

Answer 2

In fact, the distribution that is equivalent to the softmax function is the Boltzman distribution, which is in the list of common distributions you linked to. (The Wikipedia entry focusses on its original use in statistical mechanics and you would obviously ignore the $kT$ denominator for machine learning uses, though it's worth noting that softmax is not invariant under scale).

I can't answer why the Dirichlet PDF (as a function) is not as popular in NNs, though I suspect that this answer comparing softmax to normalisation is the reason. It suggests that softmax has the nice property of effectively cancelling out the log in cross entropy loss, leading to faster learning when the loss is larger.

Answer 3

A softmax function serves multiple purposes. The output of a softmax function is typically interpreted as the parameters of a Categorical distribution (a.k.a. Multinomial distribution with event count equals 1). However, it may be alternatively known as the mean of a Dirichlet distribution (Nandy et al., 2020). In addition, if a softmax is applied to a Gaussian random variable, the output of the softmax is a so-called Logistic Normal (Aitchison et al., 1980; Blei et al., 2007) random variable.

Literatures, indeed, admit the name "softmax distribution", e.g. (Dieng et al., 2020). You might need to infer its meaning from the context.