Hierarchical Decomposition of Position Encodings, Enabling BERT to Handle Ultra-Long Text

层次分解位置编码，让BERT可以处理超长文本 苏剑林 · 2020-12-04 · Machine Learning model optimization attention

As we all know, mainstream BERT models today can process at most 512 tokens of text. The root cause of this bottleneck is that BERT uses absolute position encodings trained from random initialization, with the maximum position typically set to 512, so it can handle at most 512 tokens — anything beyond that simply has no position encoding available. Of course, there's also another important reason: the $\mathcal{O}(n^2)$ complexity of attention, which causes GPU memory usage to increase dramatically for long sequences, making it unfinetunable on ordinary GPUs.

This post focuses on the former reason. Assuming there is enough GPU memory available, how can we make a simple modification to a BERT model whose current maximum length is 512, so that it can directly process longer text? The main idea is to hierarchically decompose the already-trained absolute position encodings so that they can be extended to longer positions. more

Position Encoding

BERT uses trained absolute position encodings. This approach is simple and direct, and works quite well, but since each position vector is trained by the model itself, we have no way to infer the encoding vectors for other positions, hence the length limit.

One mainstream approach to solving this problem is to switch to relative position encodings — this is a viable option. Huawei's NEZHA] model, for instance, is a BERT variant that has switched to relative position encodings. Relative position encodings typically truncate the positional difference, keeping the relative positions to be handled within a bounded range, and thus they aren't constrained by sequence length. But relative position encodings aren't a perfect solution either: for one thing, relative position encodings like NEZHA's increase computational cost (though this isn't the case for something like T5); for another, linear attention] can't use relative position encodings at all, meaning the approach isn't universal enough.

Readers might recall that Attention is All You Need] proposed a Sinusoidal absolute position encoding expressed using $\sin,\cos$ — couldn't we just use that directly to remove the length limit? In theory, yes, but the problem is that there's currently no open-source model using Sinusoidal position encodings — are we supposed to train one from scratch ourselves? That's clearly not very practical.

Hierarchical Decomposition

So, given limited resources, the most ideal approach is still to find a way to extend the position encodings of an already-trained BERT model, without retraining the model. Below I present a hierarchical decomposition scheme I came up with.

Diagram of hierarchical decomposition of position encodings

Specifically, suppose the already-trained absolute position encoding vectors are $\boldsymbol{p}_1,\boldsymbol{p}_2,\cdots,\boldsymbol{p}_n$. We want to construct, on this basis, a new set of encoding vectors $\boldsymbol{q}_1,\boldsymbol{q}_2,\cdots,\boldsymbol{q}_m$, where $m > n$. To do this, we set

\begin{equation}\boldsymbol{q}_{(i-1)\times n + j} = \alpha \boldsymbol{u}_i + (1 - \alpha) \boldsymbol{u}_j\label{eq:fenjie}\end{equation}

where $\alpha\in (0, 1)$, and $\alpha\neq 0.5$ is a hyperparameter, and $\boldsymbol{u}_1,\boldsymbol{u}_2,\cdots,\boldsymbol{u}_n$ is the "basis" of this set of position encodings. This representation has a very clear meaning: it hierarchically expresses the position $(i - 1)\times n + j$ as $(i, j)$, and then the position encodings corresponding to $i, j$ are respectively $\alpha \boldsymbol{u}_i$ and $(1 - \alpha) \boldsymbol{u}_j$, while the final encoding vector for $(i - 1)\times n + j$ is the sum of the two. The requirement that $\alpha\neq 0.5$ is to distinguish between the two different cases $(i, j)$ and $(j, i)$.

We want the position vectors to remain the same as the original ones when not exceeding $n$, so as to be compatible with the already-trained model. In other words, we want $\boldsymbol{q}_1=\boldsymbol{p}_1,\boldsymbol{q}_2=\boldsymbol{p}_2,\cdots,\boldsymbol{q}_n=\boldsymbol{p}_n$, which lets us back out each $\boldsymbol{u}_i$:

\begin{equation}\boldsymbol{u}_i = \frac{\boldsymbol{p}_i - \alpha\boldsymbol{p}_1}{1 - \alpha},\quad i = 1,2,\cdots,n\end{equation}

In this way, our parameters remain $\boldsymbol{p}_1,\boldsymbol{p}_2,\cdots,\boldsymbol{p}_n$, but we can now represent encodings for $n^2$ positions, and the first $n$ position encodings are compatible with the original model.

Self-Critique

Indeed, once you understand it, you might feel that this decomposition is really nothing technically sophisticated — just a purely ad-hoc result. And honestly, that's exactly what it is.

So why would one expect this to work? First, because the hierarchical decomposition has strong interpretability, we can expect our result to have a certain degree of extrapolation ability — at the very least, it should be a decent initialization for positions greater than $n$. Second, the experiments in the next section validate this — after all, experiments are the only criterion for proving whether a trick actually works. Fundamentally, what we're doing here is quite simple: we construct an extension scheme for the position encodings that is compatible with the original first $n$ encodings, while also being able to extrapolate to more positions, leaving the rest for the model to adapt to. There are certainly countless ways to do this kind of thing — I've simply chosen the one I personally find most interpretable, offering one possibility. It's not necessarily the optimal scheme, nor is it guaranteed to work.

Additionally, let's discuss the choice of $\alpha$. My default choice is $\alpha=0.4$. In theory, both $\alpha\in (0, 1)$ and $\alpha\neq 0.5$ hold, but in practice, I'd recommend choosing a value of $0 < \alpha < 0.5$ that is on the larger side. This is because we rarely encounter sequences with tens of thousands of tokens; for personal GPUs, being able to handle up to 2048 is already quite generous. If $n=512$, that means $i = 1, 2, 3, 4$ while $j=1,2,\cdots,512$; if $\alpha > 0.5$, then according to the decomposition $\eqref{eq:fenjie}$, $\alpha \boldsymbol{u}_i$ would dominate, and thus the differences between position encodings would shrink (since $i$ only has 4 candidate values), making it harder for the model to distinguish between positions, which would slow down convergence. If instead $\alpha < 0.5$, then what dominates is $(1-\alpha) \boldsymbol{u}_j$, giving better distinguishability between position encodings ($j$ has 512 candidate values), so the model converges faster.

Experimental Verification

In summary, we can extend BERT's absolute position encodings at essentially zero cost, allowing its maximum length to reach $n^2=512^2=262144\approx 26\text{ten thousand}$! That should surely satisfy our needs, right? This modification is already built into bert4keras>=0.9.5] — users only need to pass the parameter hierarchical_position=True in build_transformer_model to enable it. True can also be set to a floating-point value between 0 and 1, representing the value of $\alpha$ mentioned above; when it's True, the default is $\alpha=0.4$.

As for the results, I first tested the MLM task, directly setting the max length to 1536, then loading pretrained RoBERTa weights, and found that the MLM accuracy was roughly 38% (compared to around 55% when truncated to 512). After finetuning, the accuracy could quickly recover (in about 3000 steps) to above 55%. This result shows that the position encodings extended this way are indeed effective for the MLM task. If you have spare compute, it's probably a good idea to continue pretraining on MLM for a while before moving on to other tasks. We also ran experiments with different values of $\alpha$, showing that $\alpha=0.4$ is indeed a good default value, as shown in the figure below.

MLM training accuracy under different alpha values

I then tested two long-text classification problems, setting the length to 512 and 1024 respectively, keeping all other parameters the same for finetuning (finetuning directly, without first continuing pretraining on MLM). For one of the datasets, there was no noticeable change; for the other, the 1024 setting achieved about 0.5% higher accuracy on the validation set than the 512 setting. This again shows that the hierarchically decomposed position encoding proposed here does work. So, if you have a GPU with enough memory, feel free to give it a try — especially for long-text sequence labeling tasks, where it seems like it should be quite well-suited. In bert4keras, it's really just one extra line of code — if it helps, great; if not, you haven't wasted much effort. Feel free to report your own test results.

Finally, here's a reference table of maximum length vs. maximum batch size during training (RoBERTa Base, on a 24GB TITAN RTX):

$$\begin{array}{c|c} \hline \text{sequence length} & \text{batch_size}\\ \hline 512 & 22\\ 1024 & 9\\ 1536 & 5\\ \hline \end{array}$$

From this table, we can see that when the sequence length doubles, the GPU memory usage roughly doubles as well (a bit more) — which seems to contradict the legendary $\mathcal{O}(n^2)$ complexity? In fact, $\mathcal{O}(n^2)$ complexity really shows up for sufficiently long sequences — and "sufficiently long" here means several thousand to tens of thousands of tokens. For sequences no longer than 2048, BERT's complexity is actually still nearly linear. So in this kind of scenario, directly using the "BERT + extended position encoding" approach is much more convenient than designs like "sentence splitting + BERT + LSTM."

Summary

This post has shared a scheme I came up with for extending position encodings based on hierarchical decomposition. Through this extension, BERT can theoretically handle text up to 260,000 tokens long — as long as you have enough GPU memory, there's no long text BERT can't handle.

So, is your GPU memory ready?