A Theoretical Flaw in Relative Position Encoding Transformers, and Some Countermeasures

相对位置编码Transformer的一个理论缺陷与对策 苏剑林 · 2022-06-07 · Machine Learning language model attention positional encoding

Positional encoding is a crucial component of the Transformer. In Positional Encodings That Have Racked Researchers' Brains] we already summarized a number of common positional encoding designs. Broadly speaking, we can split Transformer positional encodings into two categories: "absolute positional encoding" and "relative positional encoding," with the latter generally performing somewhat better across a range of NLP/CV experiments.

However, it turns out that virtually all current relative positional encoding schemes operate on the Attention matrix before the Softmax, and this way of injecting position information actually carries a theoretical flaw that prevents the Transformer from being a "universal approximator." This post analyzes this issue and explores some possible fixes.

A Simple Probe

As the name suggests, positional encoding is meant to supply the model with positional information. So how do we judge whether a model has sufficient ability to recognize position? I once devised a simple probe experiment for this purpose:

A model with the ability to recognize position should be able to accurately implement the following mapping:
\begin{equation}\begin{array}{lc} > \text{input:} & [0, 0, \cdots, 0, 0] \\ > & \downarrow\\ > \text{output:} & [1, 2, \cdots, n-1, n] > \end{array}\end{equation}

That is, given $n$ zeros as input, the model should be able to output the position indices $1\sim n$ in order. The idea behind this probe is simple: if a model can do this, then recognizing position is an intrinsic capability of the model itself, independent of the external input — which is exactly what we want. It's easy to see that absolute positional encoding, since it's applied directly to the input, can trivially pass this probe test.

Falling Short

However, when I applied this simple probe to Transformer models equipped with relative positional encoding, I found that almost none of them could accomplish the above task.

Specifically, aside from the design proposed in Self-Attention with Relative Position Representations], every other relative positional encoding scheme (including RoPE], which I proposed myself) only modifies the Attention matrix before the Softmax. As a result, the Attention matrix — despite now carrying relative position information — remains a probability matrix (i.e., each row still sums to 1).

On the other hand, for a Transformer model, the only source of interaction between tokens is the $\boldsymbol{A}\boldsymbol{V}$ step of Self Attention, which can be written as $\boldsymbol{o}_i = \sum\limits_j a_{i,j}\boldsymbol{v}_j$. Identical inputs mean every $\boldsymbol{v}_j$ is identical, so

\begin{equation}\boldsymbol{o}_i = \sum_j a_{i,j}\boldsymbol{v}_j = \sum_j a_{i,j}\boldsymbol{v} = \left(\sum_j a_{i,j}\right)\boldsymbol{v} = \boldsymbol{v}\end{equation}

which means every $\boldsymbol{o}_i$ is also identical. In other words, each position in the model outputs exactly the same result from start to finish, so there is no way for the model to output distinct values of $[1, 2, \cdots, n-1, n]$.

A similar observation appears in the recent paper Your Transformer May Not be as Powerful as You Expect], where the authors construct a slightly different example to demonstrate this same fitting-capacity deficiency in relative-position Transformers — arriving independently at essentially the same conclusion. Also, at the start of this post I mentioned "universal approximation" — does fixing this counterexample actually restore universal approximation ability? That paper also provides a corresponding theoretical analysis confirming this, which I won't repeat here.

An Initial Fix

With a bit of thought, it becomes clear that the root of the problem is that each row of the Attention matrix sums to 1. To fix this, we just need to find a way to break this constraint. To this end, Your Transformer May Not be as Powerful as You Expect], building on this finding, further proposes the following design:

\begin{equation}\boldsymbol{O} = (\boldsymbol{A}\odot \boldsymbol{C})\boldsymbol{V}\quad \text{or equivalently}\quad\boldsymbol{o}_i = \sum_j a_{i,j}c_{i,j}\boldsymbol{v}_j\end{equation}

where $\boldsymbol{C}$ is a trainable parameter matrix and $\odot$ denotes elementwise multiplication (the Hadamard product]). To keep the whole model relying purely on relative position information (since this post specifically concerns the deficiency of relative position encoding Transformers), we constrain $\boldsymbol{C}$ to be a Toeplitz matrix], i.e., $c_{i,j}=g(i-j)$.

With $\boldsymbol{C}$ introduced, $\boldsymbol{A}\odot \boldsymbol{C}$ as a whole clearly no longer necessarily sums to 1 along each row, which breaks the constraint and thus resolves the issue (see the original paper for more experimental results). But this comes at a cost: not only does it introduce a new parameter matrix, but since $\boldsymbol{C}$ itself has a fixed finite size, it doesn't naturally support variable-length inputs (or the matrix $\boldsymbol{C}$ needs to be truncated accordingly, i.e., in the form $c_{i,j}=g(\text{clip}(i-j, p_{\min}, p_{\max}))$). All in all, this solution feels a bit inelegant.

Dropping the Denominator

Let's return once more to the root cause: every row of the Attention matrix sums to 1. What operation causes this? Clearly, it's the Softmax:

\begin{equation}a_{i,j} = \frac{e^{b_{i,j}}}{\sum\limits_j e^{b_{i,j}}}\end{equation}

Here $\boldsymbol{B}=(b_{i,j})$ is the matrix before the Softmax is applied. It's obvious that it is the "divide by $\sum\limits_j e^{b_{i,j}}$" step that leads to $\sum\limits_j a_{i,j}=1$, which suggests a very direct idea:

If I don't want $\sum\limits_j a_{i,j}=1$, why not just skip dividing by $\sum\limits_j e^{b_{i,j}}$ altogether?

And indeed, this works! Experiments show that a Transformer that skips this division does successfully pass the probe test described above. This makes me appreciate the "foresight" of GAU] all the more: the novel Attention it proposed applies a $\text{relu}^2$ activation and then simply normalizes by dividing by $n$, avoiding $\sum\limits_j a_{i,j}=1$ altogether — thereby increasing the model's theoretical capacity (though perhaps the authors weren't thinking that far ahead, and this is largely my own retrospective interpretation).

A New Normalization

However, we noted in It Turns Out Attention and Softmax Are a Perfect Match] that Attention designs like the one in GAU, which forgo probability normalization, may suffer from weaker length-generalization ability. In other words, applying probability normalization causes the theoretical flaw discussed above, while simply normalizing by dividing by $n$ may hurt length generalization. Is there a way to get the best of both worlds?

Let's think a bit more broadly here. From the perspective of norms, $\sum\limits_j e^{b_{i,j}}$ is actually the $l_1$-norm of the vector $e^{b_{i,:}}$, so Softmax is essentially a $l_1$-normalization operation on the vector's $e^{b_{i,:}}$. So, to avoid $\sum\limits_j a_{i,j}=1$ while still keeping some form of normalization, could we swap in a different kind of normalization? For instance, $l_2$ normalization:

\begin{equation}a_{i,j} = \frac{e^{b_{i,j}}}{\sqrt{\sum\limits_j e^{2b_{i,j}}}}\end{equation}

After testing this, I found that Attention using this $l_2$ normalization does indeed pass the probe experiment. So does this modification actually help with the NLP pretraining scenarios we care more about? I ran corresponding comparison experiments, and the results come in two parts:

1. For the standard Attention + FFN combination, applying $l_2$-normalized Attention requires first shrinking the initial variance of the Attention $\boldsymbol{W}_V,\boldsymbol{W}_O$; the resulting performance was slightly worse than ordinary $l_1$-normalized Attention.
2. For the full-GAU architecture, $l_2$-normalized Attention can be applied directly, with no need to change the initialization; the resulting performance was slightly better than ordinary $l_1$-normalized Attention.

The difference between the two probably stems from how differently they're initialized: in the standard Attention + FFN combination, the initial Attention matrix is close to a uniform matrix (all entries roughly equal), whereas in Does the Gated Attention Unit (GAU) Still Need Warmup?] we showed that GAU's initial Attention matrix is closer to (some multiple of) the identity matrix.

A Twist in the Plot

Looking back over everything above, we see that the reason "a model with $\sum\limits_j a_{i,j}=1$ can't pass the probe experiment" is that "every $\boldsymbol{v}_j$ is identical." But what if not every $\boldsymbol{v}_j$ were identical?

We know that, starting from BERT, mainstream Transformer models format their input like "[CLS] SENT [SEP]" — that is, they append certain marker tokens before and after the actual input. If we treat these marker tokens as part of the model rather than as part of the input (i.e., the input becomes "[CLS] 0 0 ⋯ 0 0 [SEP]" rather than all zeros), could the probe then be passed?

I ran this experiment too, and found that once such marker tokens are added to the input, the probe test can indeed be passed — with no need to modify any other part of the relative-position-encoding Transformer. This result is almost comical: it turns out the authors of BERT were quite prescient after all — the special tokens [CLS] and [SEP] they introduced also happen to serve as positional anchors. All this theoretical deficiency we've been agonizing over turns out to be resolved simply by two special tokens. This can't help but remind one of the conclusion from How Much Position Information Do Convolutional Neural Networks Encode?], which observed that "CNNs recognize absolute position through padding" — the two phenomena share a certain kinship.

Of course, this doesn't mean all our earlier analysis was pointless. For the GAU model, for instance, switching Attention to $l_2$ normalization does genuinely speed up convergence and yield a slight performance boost. Furthermore, since we can accept $l_2$ normalization, could $e^{b_{i,j}}$ also be replaced with a more general activation function (say, one that drops the non-negativity constraint)? I ran a brief experiment with "$\text{swish}(b_{i,j})$ + $l_2$ normalization" and found it to be somewhat feasible. From this angle, Attention under $l_2$ normalization actually opens up a broader space for exploration.

Curtain Call

This post has analyzed a hidden flaw in relative-position-encoding Transformers and explored several possible countermeasures, along the way raising broader questions about the non-negativity and normalization schemes used for the Attention matrix.