A Quick Derivation of Entropy-Invariant Softmax

熵不变性Softmax的一个快速推导 苏剑林 · 2022-04-11 · Mathematics approximation entropy attention

In the post Viewing Attention's Scale Operation Through Entropy Invariance, we derived a version of the attention mechanism with the property of entropy invariance:

\begin{equation}Attention(Q,K,V) = softmax\left(\frac{\kappa \log n}{d}QK^{\top}\right)V\label{eq:a}\end{equation}

We can observe that this is achieved mainly by introducing a length-dependent scaling factor $\log n$ into the Softmax. The original derivation was rather involved and relied on quite a few assumptions, which is not ideal for building intuition, so here we supply a relatively simple and quick derivation.

Derivation

We can set aside the attention mechanism background entirely and just suppose we have $s_1,s_2,\cdots,s_n\in\mathbb{R}$, and define

$$p_i = \frac{e^{\lambda s_i}}{\sum\limits_{i=1}^n e^{\lambda s_i}}$$

Clearly this is just the result of Softmax applied to $s_1,s_2,\cdots,s_n$ after multiplying it by the scaling factor $\lambda$. Now let's compute its entropy:

\begin{equation}\begin{aligned}H =&\, -\sum_{i=1}^n p_i \log p_i = \log\sum_{i=1}^n e^{\lambda s_i} - \lambda\sum_{i=1}^n p_i s_i \\ =&\, \log n + \log\frac{1}{n}\sum_{i=1}^n e^{\lambda s_i} - \lambda\sum_{i=1}^n p_i s_i \end{aligned}\end{equation}

In the first term, $\log$ involves "exponentiate first, then average". We approximate this with "average first, then exponentiate" (a mean-field approximation):

\begin{equation} \log\frac{1}{n}\sum_{i=1}^n e^{\lambda s_i}\approx \log\exp\left(\frac{1}{n}\sum_{i=1}^n \lambda s_i\right) = \lambda \bar{s} \end{equation}

Then we know that Softmax tends to put most of its weight on the largest $\max$ (see Random Musings on Function Smoothing: Differentiable Approximations of Non-differentiable Functions), so we have the approximation

\begin{equation}\lambda\sum_{i=1}^n p_i s_i \approx \lambda s_{\max}\end{equation}

Hence

\begin{equation}H\approx \log n - \lambda(s_{\max} - \bar{s})\end{equation}

The idea of entropy invariance is to eliminate, as much as possible, the influence of the length $n$. So according to the equation above we need $\lambda\propto \log n$. If we put this back into the attention mechanism, then $s$ takes the form $\langle \boldsymbol{q}, \boldsymbol{k}\rangle\propto d$ (where $d$ is the vector dimension), so we need $\lambda\propto \frac{1}{d}$, and putting these together gives

\begin{equation}\lambda\propto \frac{\log n}{d}\end{equation}

which is exactly the result $\eqref{eq:a}$ quoted at the beginning of this post.

Summary

We've worked out a simple and clear derivation for the previously proposed "entropy-invariant Softmax".