GAU-α: A First Taste of the Next-Generation Attention That's Faster, Better, and Cheaper

GAU-α：尝鲜体验快好省的下一代Attention 苏剑林 · 2022-04-22 · Machine Learning language model attention pretraining

In FLASH: Perhaps the Most Interesting Efficient Transformer Design in Recent Times, we introduced GAU (Gated Attention Unit). I'm happy to call it "the most promising design for the next generation of attention", because it genuinely achieves the trifecta of being faster (speed), better (performance), and cheaper (GPU memory).

However, some readers got the opposite results in their own tests — slower convergence, worse performance, and so on — which is quite different from what I found in my own experiments. In this post I'll share my own training experience, and release an early-taste version, "GAU-α", for everyone to try out.

Open source repository: https://github.com/ZhuiyiTechnology/GAU-alpha

GAU-α

Let's start with the scorecard for the open-sourced "GAU-α" on CLUE tasks:

$$\small{\begin{array}{c|ccccccccccc} \hline & \text{iflytek} & \text{tnews} & \text{afqmc} & \text{cmnli} & \text{ocnli} & \text{wsc} & \text{csl} & \text{cmrc2018} & \text{c3} & \text{chid} & \text{cluener}\\ \hline \text{BERT} & 60.06 & 56.80 & 72.41 & 79.56 & 73.93 & 78.62 & 83.93 & 56.17 & 60.54 & 85.69 & 79.45 \\ \text{RoBERTa} & 60.64 & \textbf{58.06} & 74.05 & 81.24 & 76.00 & \textbf{87.50} & 84.50 & 56.54 & 67.66 & 86.71 & 79.47\\ \text{RoFormer} & 60.91 & 57.54 & 73.52 & 80.92 & \textbf{76.07} & 86.84 & 84.63 & 56.26 & 67.24 & 86.57 & 79.72\\ \text{RoFormerV2}^* & 60.87 & 56.54 & 72.75 & 80.34 & 75.36 & 80.92 & 84.67 & 57.91 & 64.62 & 85.09 & \textbf{81.08}\\ \hline \text{GAU-}\alpha & \textbf{61.41} & 57.76 & \textbf{74.17} & \textbf{81.82} & 75.86 & 79.93 & \textbf{85.67} & \textbf{58.09} & \textbf{68.24} & \textbf{87.91} & 80.01\\ \hline \end{array}}$$more

All models here are Base-sized. The table shows results on the validation sets of CLUE tasks; everyone's runs and comparisons were done under the same fair conditions, so this is a reasonable relative comparison. Note that RoFormerV2 here is not the multi-task version from RoFormerV2: Pushing the Limits of Natural Language Understanding*, but rather a version that has only gone through MLM pretraining (this version was not open-sourced). The reason for comparing against this version specifically is that GAU-α has likewise only undergone MLM pretraining.

As the table shows, aside from WSC — an outlier with an extremely small amount of data — GAU-α has an edge on most tasks, and its average score (excluding WSC) is the best of the lot. The comparison between RoFormerV2 and GAU-α is the fairest one here, since they share the same training script, training data, and overall architecture; the only difference is that GAU-α replaces the Attention+FFN combination in RoFormerV2 with two layers of GAU. This comparison fully demonstrates the "better" aspect of the GAU design.

Furthermore, as discussed in RoFormerV2: Pushing the Limits of Natural Language Understanding, RoFormerV2 simplified the architecture to achieve faster speed. GAU-α, sharing the same overall architecture, benefits from this too — so GAU-α is faster than BERT, RoBERTa, and RoFormer as listed in the table, while still achieving a better average performance. Further testing shows that once the sequence length exceeds 512, GAU-α's speed starts to overtake the similarly streamlined RoFormerV2, with lower GPU memory usage as well — and the longer the sequence, the greater GAU-α's advantage.

Training

Now let's go over the training details of the model. The full code has already been open-sourced on GitHub, so if anything is unclear, you can check it against the code.

Model architecture: GAU-α simply replaces the Attention+FFN block in RoFormerV2 with two layers of GAU. In a previous post we compared the computation and parameter counts of two GAU layers against an Attention+FFN combination and found them roughly equivalent, so this substitution is reasonable. RoFormerV2's distinguishing features are that it retains the Post-Norm structure, removes all bias terms, and replaces Layer Norm with the simplest variant of RMS Norm — and GAU-α follows the same recipe.

Normalization: In Turns out Attention Pairs Better with Softmax we discussed the normalization issue in attention. For GAU-α's attention normalization, we chose the entropy-invariant softmax that I proposed myself, which has good length-extrapolation properties (referred to as softmax_plus in bert4keras for now).

Training setup: For initialization, I made adjustments following What Exactly Is So Hard About Training a 1000-Layer Transformer?, which means training can proceed directly without any warmup. The optimizer used is LAMB, with a piecewise-linear learning rate decay schedule. The pretraining task is whole-word MLM, using Baidu's LAC for tokenization — all of this is aligned with RoFormerV2.

That about covers everything worth mentioning; there really weren't that many changes made. Apart from spending a bit of time testing different normalization schemes, I didn't spend much extra time on anything else — training directly still produced solid results.

Summary

GAU is, in my view, "the most promising design for the next generation of attention." This post shared some of my training experience with GAU, and released an early-taste version, "GAU-α", as open source.