RealFormer: Moving the Residual Connection onto the Attention Matrix

RealFormer：把残差转移到Attention矩阵上面去 苏剑林 · 2020-12-24 · Machine Learning gradient attention

Everyone knows that Layer Normalization is one of the crucial components of the Transformer model, and it is used in two ways: PostLN and PreLN. The paper On Layer Normalization in the Transformer Architecture gives a fairly detailed comparison of the two. In short, PreLN is more friendly to gradient descent, converges faster, and is more robust to training hyperparameters such as the learning rate — everything about it is good, except for one glaring flaw: PreLN's performance always seems to be slightly worse than PostLN's. A recent paper from Google, RealFormer: Transformer Likes Residual Attention, proposes the RealFormer design, which successfully closes this gap: the model enjoys the same optimization-friendliness as PreLN, while performing even better than PostLN — truly having your cake and eating it too.

Formulation

RealFormer stands for "Residual Attention Layer Transformer," i.e., "a Transformer model with residual attention layers." As the name suggests, the idea is to move the residual connection into the attention mechanism.

There's a small anecdote about this name. When this blog post was first published, RealFormer was actually called Informer, standing for "Residual Attention Transformer," and the original paper was titled Informer: Transformer Likes Informed Attention. Clearly, it's hard to guess the full name from "Informer" alone, and the author of this blog even poked fun at Google's clumsy and arbitrary naming conventions at the time. A day later, it turned out the paper had been renamed to RealFormer, so this post was updated accordingly. Whether this was because the authors saw the criticism, or because the name "Informer" clashed with another paper published just a few days earlier, Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting, who knows — ha!

Diagram of PostLN, PreLN, and RealFormer structures

Back to the model. As shown in the figure above, RealFormer essentially moves the residual connection onto the attention matrix, while keeping the overall structure as PostLN. This preserves the performance of PostLN while also incorporating the optimization-friendliness of residual connections. Specifically, where the attention at layer $n$ used to be

\begin{equation}Attention(\boldsymbol{Q}_n,\boldsymbol{K}_n,\boldsymbol{V}_n) = softmax\left(\boldsymbol{A}_n\right)\boldsymbol{V}_n,\quad \boldsymbol{A}_n=\frac{\boldsymbol{Q}_n\boldsymbol{K}_n^{\top}}{\sqrt{d_k}}\end{equation}

it now becomes

\begin{equation}Attention(\boldsymbol{Q}_n,\boldsymbol{K}_n,\boldsymbol{V}_n) = softmax\left(\boldsymbol{A}_n\right)\boldsymbol{V}_n,\quad \boldsymbol{A}_n=\frac{\boldsymbol{Q}_n\boldsymbol{K}_n^{\top}}{\sqrt{d_k}} + \boldsymbol{A}_{n-1}\end{equation}

That's it. The end.

Experiments

Of course, it can't really be "the end" that quickly — we still need to run some experiments to see how well it works. But as far as the modification itself goes, that really is all there is to it — it's that simple. The original paper runs a great many experiments, and essentially all the results point to improved performance:

$$\text{RealFormer}\geq \text{PostLN} \geq \text{PreLN}$$

It seems that this time, PostLN might really be on its way out for good. Some of the experimental results are shown below:

MLM accuracy comparisonSQuAD evaluation comparisonGLUE evaluation comparisonComparison across different training steps

The first and fourth figures deserve special mention. From the first figure, we can see that for the RealFormer architecture, scaling up the model (from large to xlarge) brings a clear performance improvement, whereas the ALBERT paper previously noted that scaling up BERT's model size does not bring noticeable benefits. Taken together, this suggests that the issue may lie with PostLN rather than being an inherent flaw of BERT itself, and switching to RealFormer can fix it. From the fourth figure, we can see that RealFormer trained for 500K steps achieves performance comparable to PostLN trained for 1M steps, which indicates that RealFormer has very high training efficiency.

Besides the experiments above, the paper also compares the effects of different learning rates and different dropout ratios, showing that RealFormer is indeed fairly robust to these hyperparameters. The original paper also analyzes the distribution of RealFormer's attention values, showing that RealFormer produces more reasonable attention patterns.

Analysis

In this section, let's do a bit of simple thinking and analysis about RealFormer.

It's not hard to understand why RealFormer is friendlier to gradient descent: the design of $\boldsymbol{A}_n=\frac{\boldsymbol{Q}_n\boldsymbol{K}_n^{\top}}{\sqrt{d_k}} + \boldsymbol{A}_{n-1}$ genuinely provides a direct path, so that the attention from the first layer can flow straight through to the last layer, and there is naturally no risk of vanishing gradients. In contrast, PostLN has the structure $\text{LayerNorm}(x + f(x))$: it looks like $x+f(x)$ prevents vanishing gradients, but the step $\text{LayerNorm}$ reintroduces the risk of vanishing gradients. The consequence is that, at the start of training, the gradients of the earlier layers are very small while those of the later layers are very large. If a large learning rate is used, the later layers tend to blow up; if a small learning rate is used, the earlier layers fail to learn properly. This is why PostLN is harder to train and requires a small learning rate together with warmup for gradual training.

So if PreLN improves the gradient situation, why does it still fall short of PostLN? My guess is this: in PreLN, each step takes the form $x+f(x)$, and by the last layer this has accumulated into the form $x + f_1(x) + f_2(x) + \cdots + f_n(x)$. Adding things up layer after layer like this may cause both the magnitude and the variance to grow very large, so in the end one is "forced" to add one more Layer Norm at the very end just to stabilize the output. In other words, although PreLN improves the gradient behavior, its design inherently carries some instability, and this may well be the reason its performance is slightly worse.

In fact, people noticed quite early on that this characteristic of residual connections can cause instability. When I was previously studying GANs, I noticed that the implementation in Which Training Methods for GANs do actually Converge? replaces $x + f(x)$ with $x + 0.1 f(x)$. Inspired by that implementation, I also tried replacing $x + f(x)$ with $x + \alpha f(x)$, where $\alpha$ is a trainable scalar parameter initialized to zero, and this achieved good results as well. A paper from earlier this year, ReZero is All You Need: Fast Convergence at Large Depth, formally proposed this method under the name ReZero, and the experiments there show that with ReZero one can even do away with Layer Norm entirely. Unfortunately, the ReZero paper does not run more extensive experiments on Transformers, and RealFormer likewise does not compare its performance against ReZero.

Readers might object: given that PreLN has this issue, doesn't RealFormer's $\boldsymbol{A}_n=\frac{\boldsymbol{Q}_n\boldsymbol{K}_n^{\top}}{\sqrt{d_k}} + \boldsymbol{A}_{n-1}$ suffer from the same kind of accumulation problem? If we only look at $\boldsymbol{A}$, then yes, this issue would indeed arise — but don't forget that $\boldsymbol{A}$ still goes through a softmax normalization before being used in the computation. That is to say, the model has built-in normalization with respect to the matrix $\boldsymbol{A}$, so there's no risk of numerical divergence. In fact, quite the opposite happens: as the number of layers increases, the accumulation of $\boldsymbol{A}$ may make the absolute values of the elements of $\boldsymbol{A}$ progressively larger, causing the attention distribution to gradually approach a one-hot form, which leads to vanishing gradients in the later layers. But recall that we just said that in PostLN, the earlier layers have small gradients while the later layers have large ones — and now this effect further shrinks the gradients of the later layers, which actually brings the two into better sync and makes optimization easier. On another note, the attention probabilities may tend to converge over layers, meaning the attention pattern may become increasingly stable, producing a regularization effect similar to ALBERT's parameter sharing, which could well be beneficial to the model's performance. Intuitively, it also seems that using the RealFormer structure for adaptive-depth improvements like FastBERT ought to work even better, since RealFormer's attention naturally tends to converge across layers, which aligns well with the motivation behind FastBERT's design.

Furthermore, we can also understand RealFormer as still using the conventional residual structure, except that the residual connection is applied only to $\boldsymbol{Q}, \boldsymbol{K}$ and not to $\boldsymbol{V}$:

\begin{equation}\begin{aligned} &Attention(\boldsymbol{Q}_n,\boldsymbol{K}_n,\boldsymbol{V}_n) = softmax\left(\boldsymbol{A}_n\right)\boldsymbol{V}_n\\ &\boldsymbol{A}_n=\frac{\tilde{\boldsymbol{Q}}_n\tilde{\boldsymbol{K}}_n^{\top}}{\sqrt{d_k}},\quad\tilde{\boldsymbol{Q}}_n = \boldsymbol{Q}_n + \tilde{\boldsymbol{Q}}_{n-1}, \quad\tilde{\boldsymbol{K}}_n = \boldsymbol{K}_n + \tilde{\boldsymbol{K}}_{n-1} \end{aligned}\end{equation}

This is, to some extent, equivalent to $\boldsymbol{A}_n=\frac{\boldsymbol{Q}_n\boldsymbol{K}_n^{\top}}{\sqrt{d_k}} + \boldsymbol{A}_{n-1}$, whereas PreLN is equivalent to adding a residual connection to both $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$. Why is it that $\boldsymbol{V}$ "doesn't deserve" a residual connection? Looking at some recent improvements to relative position encoding, I've noticed what seems to be a common trend: dropping the bias term on $\boldsymbol{V}$. For instance, NEZHA's relative position encoding applies it simultaneously to the attention matrix (i.e., $\boldsymbol{Q},\boldsymbol{K}$) and to $\boldsymbol{V}$, whereas the more recent relative position encodings of XLNet and T5 apply it only to the attention matrix. So it seems that dropping the unnecessary bias on $\boldsymbol{V}$ is a fairly good choice, and RealFormer once again reflects this same principle.

Summary

This post introduced Google's new Transformer design, RealFormer, along with my own thoughts and analysis. The experimental results show that RealFormer combines the strengths of both PostLN and PreLN, and even surpasses both — making it an improvement well worth adopting.