A Bit of Pretraining Goes a Long Way: Transformer's Long-Sequence Scores Can Rise a Lot!

预训练一下，Transformer的长序列成绩还能涨不少！ 苏剑林 · 2023-10-08 · Machine Learning language model attention

As the mainstream model architecture for LLMs, Transformer performs excellently across the board on most tasks. In most cases, the main complaint about Transformer is its quadratic complexity, not its effectiveness—except on one particular Benchmark called Long Range Arena (LRA). LRA has long been the "home turf" of linear RNN-style models, on which Transformer shows a clear gap by comparison, to the point that people started to wonder whether this was some inherent flaw of Transformer.

Recently, however, the paper Never Train from Scratch: Fair Comparison of Long-Sequence Models Requires Data-Driven Priors has filled in this "missing piece." The paper points out that the lack of pretraining is the main reason Transformer performs worse on LRA, and that all architectures can gain improvements from pretraining—with Transformer benefiting even more noticeably.

The Old Background

Long Range Arena (LRA) is a benchmark for long-sequence modeling, proposed in the paper Long Range Arena: A Benchmark for Efficient Transformers. As the title suggests, LRA was built to test various Efficient variants of Transformer. It contains multiple types of data, with sequence lengths ranging from 1k to 16k, and many prior works on Efficient Transformers have been evaluated on LRA. Although there's some controversy about how representative it is, LRA remains a classic benchmark for testing the long-sequence capabilities of Efficient Transformers.more

LRA results from the MEGA paper

It may surprise some readers that the standard Transformer (XFM) doesn't perform particularly well on this Benchmark, lagging noticeably behind a series of linear RNN-style models, such as the classic SSMs (S4, S4D, S5) or the LRU we covered previously—even the previous SOTA model MEGA had to equip a linear RNN module (called EMA in the paper) on top of GAU. In short, the previous leaderboard on LRA has been sending a strong signal that "Attention is optional, but RNN is essential."

(Note: The full LRA leaderboard can be found at https://paperswithcode.com/sota/long-range-modeling-on-lra.)

The New Conclusion

Clearly, the appearance of Never Train from Scratch: Fair Comparison of Long-Sequence Models Requires Data-Driven Priors has broken this impression. It shows that pretraining on the training set can substantially narrow the gap between the two, further advancing the view that "without pretraining, comparisons aren't fair."

The improvement of "Transformer + pretraining" over Transformer and various Effective variants

The pretraining approach is quite simple: the task can be either MLM or GPT-style, and the dataset is still the original training set. This means that, aside from the additional compute cost, no extra source of knowledge is introduced, so the comparison remains fair. Indeed, whether it's Transformer or RNN, both gain noticeable improvements after pretraining—it's just that Transformer's improvement is more pronounced:

"Transformer + pretraining" vs. "S4 + pretraining"Comparison with SOTA models

In hindsight, the paper's conclusion isn't really surprising—it even feels "obviously true"—yet apparently no one had thought along these lines before (or perhaps they had, but didn't think it was the key factor?). So credit is very much due to the authors for being the first to recognize and demonstrate the importance of pretraining on LRA.

The importance of pretraining actually reveals the importance of inductive bias on LRA. Because LRA needs sequences to be sufficiently long, its tokens are very fine-grained—for instance, the text task tokenizes at the character level, and the image task tokenizes at the pixel level, flattening 2D images directly into 1D sequences. Clearly, these tasks require both long-range dependencies and pronounced locality, which fits the characteristics of linear RNN very well. Transformer, by comparison, doesn't have such a pronounced inductive bias—it needs additional positional encoding just to have positional information, and even with that, it lacks significant locality. Hence it needs pretraining even more to adapt to the characteristics of the data, or in other words, pretraining is used to supplement its inductive bias.

The End

This post quickly shared with everyone a relatively new experimental finding: pretraining can effectively improve various models' scores on LRA, especially for Transformer—after pretraining, its performance can basically approach the SOTA tier. This overturns the author's long-held impression that LRA absolutely requires adding a linear RNN.