As we know, Layer Normalization is one of the important components of the Transformer model. Its usage comes in two main variants: PostLN and PreLN. The paper "On Layer Normalization in the Transformer Architecture" provides a detailed analysis of both. Simply put, PreLN is more friendly to gradient descent, converges faster, and is more robust to hyperparameters like learning rate. In short, it seems better in every way except for one major drawback: the performance of PreLN always seems slightly worse than PostLN. Recently, a paper from Google titled "RealFormer: Transformer Likes Residual Attention" proposed the RealFormer design, which successfully bridges this gap. It allows the model to possess the optimization friendliness of PreLN while achieving better results than PostLN, truly "having one’s cake and eating it too."
Formulation
RealFormer stands for "Residual Attention Layer Transformer." As the name suggests, it places the residual connection inside the Attention mechanism.
There is a small anecdote regarding this name. When this blog post was first drafted, RealFormer was actually called Informer, standing for "Residual Attention Transformer," and the original paper was titled "Informer: Transformer Likes Informed Attention." Obviously, it was hard to imagine the full name from the acronym "Informer," and I even criticized Google’s awkward and arbitrary naming conventions. A day later, I discovered it had been renamed to RealFormer and updated the post accordingly. I’m not sure if the authors saw my critique or if it was because the name "Informer" conflicted with another paper released a few days earlier titled "Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting." Haha!
Returning to the model, as shown in the figure above, RealFormer primarily places the residual connection on the Attention matrix while maintaining the overall PostLN structure. Therefore, it preserves the performance of PostLN while integrating the training friendliness of residuals. Specifically, the Attention in the original n-th layer is: Attention(\boldsymbol{Q}_n,\boldsymbol{K}_n,\boldsymbol{V}_n) = \text{softmax}\left(\boldsymbol{A}_n\right)\boldsymbol{V}_n,\quad \boldsymbol{A}_n=\frac{\boldsymbol{Q}_n\boldsymbol{K}_n^{\top}}{\sqrt{d_k}} Now it is changed to: Attention(\boldsymbol{Q}_n,\boldsymbol{K}_n,\boldsymbol{V}_n) = \text{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} And that’s it. End of article.
Experiments
Of course, it’s impossible to "end the article" that quickly; we must at least look at the experimental results. However, if we only look at the architectural changes, it really is that simple. The original paper conducted many experiments, and basically all results show that in terms of performance: \text{RealFormer} \geq \text{PostLN} \geq \text{PreLN} It seems that this time, PostLN might finally be able to exit the stage. Some experimental results are as follows:
It is worth highlighting the first and fourth figures. From the first figure, we can see that for the RealFormer structure, increasing the model scale (from large to xlarge) brings a significant performance boost. In contrast, the ALBERT paper once mentioned that increasing the BERT model scale does not yield significant benefits. Combining these two points suggests that this might be a flaw of PostLN rather than an inherent issue with BERT, and switching to RealFormer can improve this. From the fourth figure, we can see that training the RealFormer structure for 500,000 steps yields results equivalent to training PostLN for 1,000,000 steps, indicating that RealFormer has very high training efficiency.
In addition to the above experiments, the paper also compared the effects of different learning rates and dropout ratios, showing that RealFormer is indeed more robust to these parameters. The original paper also analyzed the distribution of Attention values, showing that RealFormer’s Attention results are more reasonable.
Analysis
In this section, we provide a brief analysis and reflection on RealFormer.
It is not difficult to understand why RealFormer is more friendly to gradient descent. The design \boldsymbol{A}_n=\frac{\boldsymbol{Q}_n\boldsymbol{K}_n^{\top}}{\sqrt{d_k}} + \boldsymbol{A}_{n-1} indeed provides a direct path, allowing the Attention of the first layer to reach the last layer directly, which naturally eliminates the risk of vanishing gradients. In contrast, PostLN uses the structure \text{LayerNorm}(x + f(x)). While x+f(x) prevents gradient vanishing, the \text{LayerNorm} step re-introduces the risk. The consequence is that in the initial stages, the gradients of the earlier layers are very small while the gradients of the later layers are very large. If a large learning rate is used, the later layers are prone to collapse; if a small learning rate is used, the earlier layers do not learn well. Therefore, PostLN is harder to train and requires a small learning rate combined with a warmup period.
So, if PreLN improves the gradient situation, why is it still inferior to PostLN? According to my guess, every step in PreLN is in the form of x+f(x). By the last layer, it becomes x + f_1(x) + f_2(x) + \cdots + f_n(x). This layer-by-layer accumulation might lead to very large values and variances, so a final Layer Norm is "forced" at the end to stabilize the output. Thus, although PreLN improves the gradient situation, its design inherently contains some unstable factors, which might be the reason for its slightly worse performance.
In fact, it was noticed long ago that this characteristic of residuals could cause instability. When I was researching GANs, I found that the implementation in "Which Training Methods for GANs do actually Converge?" replaced x + f(x) with x + 0.1 f(x). Inspired by their implementation, I also tried replacing x + f(x) with x + \alpha f(x), where \alpha is a trainable scalar parameter initialized to 0, which also achieved good results. Earlier this year, the paper "ReZero is All You Need: Fast Convergence at Large Depth" formally proposed this method, naming it ReZero. Experiments therein showed that using ReZero allows for the complete removal of Layer Norm. Unfortunately, the ReZero paper did not conduct more experiments on Transformers, and RealFormer did not compare its performance with ReZero.
Readers might argue: if PreLN has issues, 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 accumulation problem? If we only look at \boldsymbol{A}, then yes, such a problem exists. But don’t forget that \boldsymbol{A} is followed by a softmax normalization before participating in the calculation. That is to say, the model has a built-in normalization function for the matrix \boldsymbol{A}, so it does not have the risk of numerical divergence. On the contrary, as the number of layers increases, the accumulation of \boldsymbol{A} might cause the absolute values of its elements to grow larger, making the Attention gradually tend towards a one-hot form. This might cause gradients in later layers to vanish, but remember, we just said that in PostLN, the earlier layers have small gradients and the later layers have large ones. Now, by further reducing the gradients of the later layers, the two become more synchronized and thus easier to optimize. On the other hand, the Attention probability values might have a tendency to converge, meaning the Attention patterns might become increasingly stable, bringing a regularization effect similar to ALBERT’s parameter sharing, which could be beneficial for model performance. At the same time, intuitively, using the RealFormer structure for improvements like FastBERT (adaptive number of layers) might yield better results because RealFormer’s Attention itself has a convergence trend, which aligns better with the design philosophy of FastBERT.
Furthermore, we can also understand RealFormer as still using a conventional residual structure, but the residual structure is only applied to \boldsymbol{Q} and \boldsymbol{K} and not to \boldsymbol{V}: \begin{aligned} &Attention(\boldsymbol{Q}_n,\boldsymbol{K}_n,\boldsymbol{V}_n) = \text{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} To some extent, this is 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 residuals to \boldsymbol{Q}, \boldsymbol{K}, \text{ and } \boldsymbol{V}. Why is \boldsymbol{V} "not worthy" of a residual? From some recent improvements in relative position encoding, I have noticed a common trend: removing the bias of \boldsymbol{V}. For example, the relative position encoding in NEZHA is applied to both the Attention matrix (i.e., \boldsymbol{Q}, \boldsymbol{K}) and \boldsymbol{V}, while the newer relative position encodings in XLNET and T5 are only applied to the Attention matrix. Therefore, it seems that removing unnecessary biases from \boldsymbol{V} is a better choice, and RealFormer reflects this once again.
Conclusion
This article introduced Google’s new design for the Transformer, RealFormer, and provided my own reflections and analysis. Experimental results show that RealFormer possesses the advantages of both PostLN and PreLN and even outperforms both, making it a worthwhile improvement to use.
Original address: https://kexue.fm/archives/8027
For more details on reposting, please refer to: "Scientific Space FAQ"