A Unified View of Long-Sequence Models towards Modeling Million-Scale Dependencies

TL;DR

This paper reviews and unifies existing long-sequence modeling methods, demonstrates the importance of long context, and proposes a scalable multi-head attention system capable of handling million-scale dependencies efficiently.

Contribution

It provides a unified mathematical framework for long-sequence models, benchmarks their performance, and introduces a scalable attention algorithm for million-scale dependencies.

Findings

Long context length improves performance in many tasks.

Traditional Transformers struggle with long-range dependencies.

The proposed distributed multi-head attention scales efficiently on GPUs.

Abstract

Ever since their conception, Transformers have taken over traditional sequence models in many tasks, such as NLP, image classification, and video/audio processing, for their fast training and superior performance. Much of the merit is attributable to positional encoding and multi-head attention. However, Transformers fall short in learning long-range dependencies mainly due to the quadratic complexity scaled with context length, in terms of both time and space. Consequently, over the past five years, a myriad of methods has been proposed to make Transformers more efficient. In this work, we first take a step back, study and compare existing solutions to long-sequence modeling in terms of their pure mathematical formulation. Specifically, we summarize them using a unified template, given their shared nature of token mixing. Through benchmarks, we then demonstrate that long context length does yield better performance, albeit application-dependent, and traditional Transformer models fall short in taking advantage of long-range dependencies. Next, inspired by emerging sparse models of huge capacity, we propose a machine learning system for handling million-scale dependencies. As a proof of concept, we evaluate the performance of one essential component of this system, namely, the distributed multi-head attention. We show that our algorithm can scale up attention computation by almost $40\times$ using four GeForce RTX 4090 GPUs, compared to vanilla multi-head attention mechanism. We believe this study is an instrumental step towards modeling million-scale dependencies.