Entropy Normalization for Probability Distributions

概率分布的熵归一化（Entropy Normalization） 苏剑林 · 2021-12-24 · Mathematics probability entropy

In the previous article, Attention Scale from the Perspective of Entropy Invariance, we derived a new Attention Scale from the standpoint of entropy invariance, and experiments showed that this new Scale, which possesses entropy invariance, indeed improves the extrapolation performance of Attention. This left me with a natural question:

Is there some operation, analogous to L2 Normalization, that can be applied directly to a probability distribution so as to transform it into one with a specified entropy value, while preserving the main characteristics of the original distribution?

Curious, I did some searching but couldn't find any similar work, so I decided to try deriving one myself. I ended up with a result I'm reasonably satisfied with, which I'll tentatively call "Entropy Normalization." I'm writing it up here for readers who might find it useful.

Power Transform

First, suppose we have an $n$-element distribution $(p_1,p_2,\cdots,p_n)$, whose entropy is defined as

\begin{equation}\mathcal{H} = -\sum_i p_i \log p_i = \mathbb{E}[-\log p_i]\end{equation}more

Since $p_i \in [0,1]$, we have $-p_i \log p_i \geq 0$, and hence $\mathcal{H} \geq 0$. The minimum value 0 is attained when one particular $p_i$ equals 1 and all the other $p_i$ equal 0 (i.e., a one-hot distribution); it can also be shown that $p_i$ attains its maximum value $\log n$ when all the $1/n$ are equal, so that $\mathcal{H}$'s range is $\mathcal{H}$, i.e. $[0,\log n]$.

So the first thing we need is a transform of the distribution that preserves its essential information while being able to move the entropy anywhere from $0$ to $\log n$. The transform we choose here is the power transform:

\begin{equation}p_i\quad\to\quad \tilde{p}_i = \frac{p_i^{\gamma}}{\sum\limits_i p_i^{\gamma}}\end{equation}

One reason for choosing the power transform is that it preserves the monotonicity of the distribution — that is, if $p_i > p_j$, then also $\tilde{p}_i > \tilde{p}_j$, which I consider one of the important properties a distribution transform should preserve. Moreover, when all the $p_i$ are nonzero and pairwise distinct, the power transform is indeed capable of moving the entropy across $0\sim \log n$. Without loss of generality, suppose $1 > p_1 > p_2 > \cdots > p_n > 0$; clearly when $\gamma = 0$, we have $\tilde{p_i}=1/n$, at which point the entropy attains its maximum value $\log n$, and when $\gamma \to\infty$, we have

\begin{equation}\tilde{p}_1 = \lim_{\gamma\to\infty}\frac{p_1^{\gamma}}{\sum\limits_i p_i^{\gamma}} = \lim_{\gamma\to\infty}\frac{1}{1 + \sum\limits_{i > 1} (p_i/p_1)^{\gamma}}=1\end{equation}

That is to say, at this point the distribution becomes the one-hot distribution $(1,0,\cdots,0)$, whose corresponding entropy is the minimum value 0. In fact, by taking a derivative one can further show that the entropy is monotonically decreasing with respect to $\gamma$, so that as $\gamma$ increases from $0$ to $\infty$, the entropy decreases monotonically from $\log n$ to $0$.

Iterative Solution

Having confirmed that the power transform is indeed a usable transformation, we now need to work out how to solve for it — that is, for any given $\mathcal{H}^*\in(0,\log n)$, we need to find the correct $\gamma$ such that the resulting entropy equals a specified target value $\mathcal{H}^*$.

First, let's write out

\begin{equation}\mathcal{H}_{\gamma} = -\sum_i\frac{p_i^{\gamma}}{\sum\limits_i p_i^{\gamma}}\log \frac{p_i^{\gamma}}{\sum\limits_i p_i^{\gamma}}=\log\sum_i p_i^{\gamma} - \frac{\gamma\sum\limits_i p_i^{\gamma}\log p_i}{\sum\limits_i p_i^{\gamma}}\end{equation}

The complexity of the rightmost expression convinces us that there is unlikely to be a closed-form solution, so we have to resort to an iterative solving algorithm.

We expand it around $\gamma=1$ (making chief use of $p_i^{\gamma}\approx p_i + (\gamma-1)p_i\log p_i$):

\begin{equation}\begin{aligned} \mathcal{H}_{\gamma} \approx &\, -\sum_i p_i\log p_i + \left(\left(\sum_i p_i\log p_i\right)^2-\sum_i p_i\left(\log p_i\right)^2\right)(\gamma - 1)\\ =&\, \mathcal{H}_1 + \left(\mathcal{H}_1^2-\mathbb{E}[\left(\log p_i\right)^2]\right)(\gamma - 1) \end{aligned}\end{equation}

which gives

\begin{equation}\gamma \approx 1 + \frac{\mathcal{H}_{\gamma}-\mathcal{H}_1}{\mathcal{H}_1^2-\mathbb{E}[\left(\log p_i\right)^2]}\end{equation}

Based on this result, starting from $\gamma=1$ and repeatedly applying the above formula, we can iteratively solve for the final distribution:

\begin{equation} \mathcal{H}\leftarrow -\sum_i p_i \log p_i,\quad \gamma \leftarrow 1 + \frac{\mathcal{H}^*-\mathcal{H}}{\mathcal{H}^2-\mathbb{E}[\left(\log p_i\right)^2]},\quad p_i \leftarrow \frac{p_i^{\gamma}}{\sum\limits_i p_i^{\gamma}} \end{equation}

This is, in fact, just Newton's method for solving a nonlinear equation. In experiments, I found that 3–4 iterations already give quite good convergence; if in practice you only need to roughly control the entropy range, then 1–2 iterations will suffice.

Reference implementation in Numpy:

p = np.random.random(100)
p /= p.sum()  # 模拟分布
gamma = 1
H_f = np.log(30)  # 希望达到的熵

for i in range(10):
    H = -(p * np.log(p)).sum()
    gamma = 1 + (H_f - H) / (H**2 - (p * np.log(p)**2).sum())
    p = p**gamma
    p /= p.sum()

Possible Applications

The main reason I wrote this post is simply that I found the concept of "Entropy Normalization" interesting enough to work out the derivation. As for what good concrete applications it might have, I haven't fully thought that through yet.

The smaller the entropy, the more the probability mass is concentrated on a few positions — in other words, the probabilities at other positions get closer to zero. So in a sense, entropy is a measure of how sparse a probability distribution is. If we want a sparser prediction result, we can control this via entropy normalization. On the other hand, the sparser the distribution, the more likely the model is to suffer from vanishing gradients, so conversely, entropy normalization could also be used to keep the entropy from getting too small, thereby alleviating the vanishing gradient problem.

Speaking of sparsity, this naturally brings to mind Sparsemax and my own Sparse Softmax that I devised earlier. Sparsemax achieves sparsity by treating entropy as a penalty term, whereas Sparse Softmax introduces sparsity through direct truncation; both offer better interpretability or better performance in certain scenarios. So does sparsity induced directly through entropy normalization actually work well? This might be a question worth exploring.

Also, in random sampling for autoregressive models, we often use top-$k$ or top-$p$ truncation, which is essentially also a way of lowering the entropy of the distribution. Correspondingly, we could use entropy normalization to keep the sampling entropy consistent at every step, as a substitute for top-$k$ or top-$p$ sampling — this is another possible application.

The main issue with using entropy normalization is that there's no clear standard for "what value should we normalize the entropy to." I don't currently have a great idea for this either; for now, the best I can suggest is tuning it by observing existing experimental results, but that's far from an ideal answer.

Summary

This post introduced the concept of Entropy Normalization, a direct transform that allows the entropy of a distribution to be set to a specified value, and sketched out a few potential applications.