Two Interesting Observations about Attention and Softmax: Robustness and Information Content
Over the past few weeks I've been thinking about various properties of the attention mechanism, and in the process I've gained a deeper understanding of both attention and softmax. In this post I'll briefly share two of these insights:
1. Softmax attention is naturally resistant to a certain amount of noise perturbation;
2. The initialization problem can also be given an intuitive understanding from the perspective of information entropy.
Robustness
The attention mechanism, based on softmax normalization, can be written as
\begin{equation}o = \frac{\sum\limits_{i=1}^n e^{s_i} v_i}{\sum\limits_{i=1}^n e^{s_i}}\end{equation}
One day it occurred to me: what would happen if we added i.i.d. noise to $s_i$? To explore this, let's consider
\begin{equation}\tilde{o} = \frac{\sum\limits_{i=1}^n e^{s_i+\varepsilon_i} v_i}{\sum\limits_{i=1}^n e^{s_i+\varepsilon_i}}\end{equation}
where $\varepsilon_i$ is i.i.d. noise. However, after a bit of analysis I found that the conclusion is "not much happens" — the attention mechanism is naturally resistant to this kind of noise, i.e. $\tilde{o}\approx o$.
To see why, we just need to realize that:
\begin{equation}\tilde{o} = \frac{\frac{1}{n}\sum\limits_{i=1}^n e^{s_i+\varepsilon_i} v_i}{\frac{1}{n}\sum\limits_{i=1}^n e^{s_i+\varepsilon_i}}=\frac{\mathbb{E}_i[e^{s_i+\varepsilon_i} v_i]}{\mathbb{E}_i[e^{s_i+\varepsilon_i}]}\approx \frac{\mathbb{E}_i[e^{s_i}v_i]\mathbb{E}[e^{\varepsilon}]}{\mathbb{E}_i[e^{s_i}]\mathbb{E}[e^{\varepsilon}]}=\frac{\mathbb{E}_i[e^{s_i}v_i]}{\mathbb{E}_i[e^{s_i}]}=o\end{equation}
The approximate equality uses the fact that $\varepsilon_i$ and $s_i,v_i$ are mutually independent, so the expectation of the product equals the product of the expectations.
Information Content
If we let $p_i = e^{s_i}\left/\sum\limits_{i=1}^n e^{s_i}\right.$, then $p_i$ describes a discrete probability distribution, and we can compute its information entropy
\begin{equation}H = -\sum_{i=1}^n p_i\log p_i\quad\in[0,\log n]\end{equation}
In "Making Sense of Entropy: From Entropy and the Maximum Entropy Principle to the Maximum Entropy Model (Part 1)"], we discussed how entropy is a measure of uncertainty, and also a measure of information content. How should we understand the connection between the two? Entropy is essentially a measure of uniformity — the more uniform something is, the less certain it is, so entropy measures uncertainty. Since entropy has a lower bound of 0, this uncertainty also represents the maximum amount of information we can gain in going from "uncertain" to "completely certain."
We know that if we initialize $s_i$ to be very large, then $p_i$ approaches a one-hot distribution, at which point training becomes impossible due to vanishing gradients (see "A Brief Discussion on Transformer Initialization, Parameterization, and Normalization"]). I found that this can also be understood quite intuitively from the perspective of information content: training a model is essentially a process of going from uncertainty (a random model) to certainty (a trained model), and the optimizer is responsible for "extracting" information from the random model. A one-hot distribution has zero information content, so there's nothing left for the optimizer to "extract" — if anything, it might even have to "pay in" — and naturally optimization cannot proceed well. This is why we should initialize models to be as uniform as possible, to ensure that the amount of information available for extraction is as large as possible.
Of course, besides ensuring the upper bound on information content is large enough, we also need to ensure the lower bound is small enough, so that the amount of information actually extractable is maximized. Earlier, when discussing contrastive learning, some readers didn't quite understand the significance of the temperature parameter — this too can be understood in terms of information content. Let
\begin{equation}p_i = \frac{e^{(\cos\theta_i) / \tau}}{\sum\limits_{i=1}^n e^{(\cos\theta_i)/\tau}}\end{equation}
If $\tau=1$, then the upper bound on the information entropy is $\log n$, but the lower bound is approximately $\log n - 0.4745$ (see the comments section]), meaning the amount of information obtainable is too small. This is why we shrink $\tau$, so that the lower bound on entropy approaches 0, thereby increasing the amount of information that can be obtained.
In Short
Just a quick, casual blog post. As we can see, the conclusion still comes down to — "Turns Out Attention and Softmax Really Are Meant for Each Other"].
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.