Transformer Upgrade Path: 14. When HWFA Meets ReRoPE

Transformer升级之路：14、当HWFA遇见ReRoPE 苏剑林 · 2023-08-24 · Machine Learning attention positional encoding extrapolation RoPE

In the previous article, Transformer Upgrade Path: 13. Inverting Leaky ReRoPE, I tried applying the idea of Leaky ReRoPE in reverse during training, so that the positional encoding at inference time reduces to ordinary RoPE — thereby achieving length extrapolation while avoiding the slower inference speed of ReRoPE. Unfortunately, the experimental results showed that "Leaky ReRoPE → RoPE" did not perform as well as "RoPE → ReRoPE/Leaky ReRoPE," so this problem has not been fully solved yet.

At this point, I recalled that HWFA, which I proposed earlier in Transformer Upgrade Path: 9. A New Idea for Global Length Extrapolation, already has some inherent length-extrapolation ability. Would combining it with ReRoPE — joining forces, so to speak — produce even better results? More importantly, adding HWFA could substantially reduce inference cost, thereby compensating for ReRoPE's main weakness!

A Quick Recap

First, let's go through the customary recap of HWFA. HWFA (Hybrid Window-Full Attention) is not a specific model per se, but rather a way of combining attention mechanisms that boosts a model's length-extrapolation ability while largely preserving performance, and also reduces both training and inference cost.

Specifically, HWFA consists of "$L-1$ layers of Window RoPE Attention + $1$ layers of Full NoPE Attention," meaning the first $L-1$ attention layers all use RoPE and restrict the receptive field via a window, which makes inference cost constant and, combined with block-parallel optimization, can also speed up training. As for the last attention layer, it retains the global form but drops the positional encoding (NoPE), while adding a $\log n$ scaling factor. With these modifications, and by choosing the window size appropriately, the model's training performance only degrades slightly, while it exhibits excellent length-extrapolation ability.

Coincidentally, Google later proposed FOT (Focused Transformer), which has a lot in common with HWFA: it likewise uses $L-1$ layers of Local Attention plus $1$ layers of Full Attention, with the Full Attention also being NoPE. The difference is that FOT places the Full Attention layer in the middle, and its Local Attention doesn't strictly restrict the receptive field, so it can't directly extrapolate to longer lengths — hence they proposed crossbatch training to extend the model's effective length. Afterward, I experimented with applying crossbatch training to HWFA as well, and it also produced good results.

New Insight

Back to the main topic of this article: how can HWFA "join forces" with ReRoPE? Recall that ReRoPE is applied to Full RoPE Attention — it's a post-hoc truncation of the relative position matrix at inference time:

$$\begin{pmatrix}0 & \\ 1 & 0 & \\ 2 & 1 & 0 &\\ \ddots & 2 & 1 & 0 & \\ \ddots & \ddots & 2 & 1 & 0 & \\ \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\ \small{L - 2} & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\ \small{L - 1} & \small{L - 2} & \ddots & \ddots & \ddots & 2 & 1 & 0 & \\ \end{pmatrix} \,\to\, \begin{pmatrix} \color{red}{0} & \\ \color{red}{1} & \color{red}{0} & \\ \color{red}{\ddots} & \color{red}{1} & \color{red}{0} & \\ \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{1} & \color{red}{0} & \\ \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\ddots} & \color{green}{w} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \\ \color{green}{w} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{1} & \color{red}{0} & \\ \end{pmatrix}$$

Surprisingly, this kind of post-processing exhibits excellent length-extrapolation ability. However, due to the special structure of RoPE, the original ReRoPE implementation requires computing the attention matrix twice, and it's incompatible with mainstream acceleration techniques like Flash Attention. Overall, this leads to a noticeable increase in inference cost.

However, adding HWFA can greatly alleviate this problem! To summarize: ReRoPE is only applied to Full RoPE Attention, while most of HWFA consists of Window RoPE Attention — so the "HWFA + ReRoPE" scheme naturally suggests itself: during training, replace HWFA's original Full NoPE Attention with Full RoPE Attention, and then during inference, switch it to Full ReRoPE Attention. This way, the extra cost of switching to ReRoPE at inference time becomes very small, while the benefit of converting the other layers to Window Attention becomes even more pronounced.

Beyond that, "HWFA + ReRoPE" can also make up for the performance loss of the original HWFA. Previously, in order to guarantee length-extrapolation ability, HWFA's Full Attention had to drop positional encoding entirely (i.e., NoPE), and the receptive field $\tilde{w}$ of the Window Attention had to satisfy $(\tilde{w}-1)(L-1)+1 = \alpha N$ (where $L$ is the number of layers, $N$ is the training length, and $0 < \alpha \leq 1$). These constraints limited the model's expressive power and degraded training performance. But once ReRoPE is introduced, the receptive field of Window Attention can be made somewhat larger, the Full Attention can use RoPE, and it can be placed in a middle layer rather than only at the end — it can even span more than $1$ layers of Full Attention. All of these changes can compensate for the performance loss, and thanks to ReRoPE, the length-extrapolation ability doesn't degrade either.

To distinguish it from the original version of HWFA, we can refer to the "HWFA + ReRoPE" combination as "HWFA2."

Experiments

Below I share some experimental results for "HWFA + ReRoPE (HWFA2)." Since introducing ReRoPE gives HWFA a lot more degrees of freedom, the experiments below only cover combinations that I felt were intuitively reasonable, and cannot fully verify all possible permutations.

The experimental model is the same as before, for both HWFA and ReRoPE: a 100M-parameter GAU model with a training length of 512. Note that there are two window-related parameters here: one is ReRoPE's own $w$ parameter, which earlier ReRoPE experiments showed has little effect, so it's fixed at 256 below; the other is the receptive field of HWFA's Window Attention, denoted above as $\tilde{w}$, which is adjustable. So the main parameters of "HWFA + ReRoPE" are the Window Attention's $\tilde{w}$, along with the number and placement of Full Attention layers. Earlier comparative experiments showed that, in terms of training performance, placing Full Attention in the middle works better than placing it at the end. So if there is 1 layer of Full Attention, its default placement is at the layer with index = num_layers / 2; if there are 2 layers of Full Attention, the default placement is at the layers with index = num_layers / 3 and index = 2 * num_layers / 3, and so on.

Some of the experimental results are as follows:

$$\begin{array}{c|cc} \hline \text{test length} & 512(\text{training}) & 4096(\text{repeat}) & 4096(\text{no repeat})\\ \hline \text{Baseline} & 49.41\% & 24.17\% & 23.16\% \\ \text{Baseline-}\log n & 49.40\% & 24.60\% & 24.02\% \\ \hline \text{ReRoPE-w256} & 49.41\% & 77.90\% & 48.48\% \\ \text{ReRoPE-w256-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.40\% & 48.85\% \\ \text{ReRoPE-w256-}\log n & 49.40\% & \boldsymbol{85.12\%} & 49.07\% \\ \hline \text{InvLeaky ReRoPE-w128-}\log n & 49.38\% & 82.25\% & 48.32\% \\ \text{InvLeaky ReRoPE-w128-b8-}\log n & 49.62\% & 81.15\% & 48.85\% \\ \hline \text{HFWA} & 48.70\% & 80.84\% & 48.15\% \\ \hline \text{HFWA-ReRoPE-w32-f1} & 49.29\% & 83.13\% & 49.34\% \\ \text{HFWA-ReRoPE-w64-f1} & 49.32\% & 82.41\% & \boldsymbol{49.37\%} \\ \text{HFWA-ReRoPE-w128-f1} & 49.21\% & 80.18\% & 48.99\% \\ \text{HFWA-ReRoPE-w256-f1} & 49.00\% & 54.94\% & 47.64\% \\ \text{HFWA-ReRoPE-w32-f2} & \boldsymbol{49.50}\% & 84.09\% & 49.35\% \\ \text{HFWA-ReRoPE-w64-f2} & 49.46\% & 84.43\% & 49.36\% \\ \text{HFWA-ReRoPE-w128-f2} & 49.35\% & 83.09\% & 48.97\% \\ \text{HFWA-ReRoPE-w256-f2} & 49.37\% & 75.24\% & 48.42\% \\ \hline \end{array}$$

In the table above, the number following $\text{w}$ is the size of the Window Attention's receptive field $\tilde{w}$, and the number following $\text{f}$ is the number of Full Attention layers. Due to various constraints, the original HWFA could only take $\tilde{w}$ up to 16; any larger and the length-extrapolation ability would drop noticeably. As the table shows, once $\tilde{w}$ is increased, training performance quickly catches up to the baseline, and further increasing the number of Full Attention layers even surpasses the baseline. As for extrapolation performance, both of the $\text{w32},\text{w64}$ cases are quite good, clearly exceeding HWFA. Overall, the best combination for HWFA-ReRoPE is $\text{w64-f2}$, whose training performance and non-repetitive extrapolation performance both exceed those of the original ReRoPE. Combined with the fact that the training length $N$ is 512 and the number of layers $L$ is 24, we might guess that the optimal value of $\tilde{w}$ should be roughly $2\sim 4$ times $N/L$.

Summary

This article proposed a way of combining HWFA with ReRoPE. Small-scale experimental results show that this combination can achieve near-optimal length-extrapolation performance without sacrificing training performance, and thanks to HWFA's design, it can also noticeably reduce inference cost — effectively mitigating the original drawback of ReRoPE's increased inference cost.