Vcc: 重要なトークンを優先することでTransformerを128Kトークン以上にスケーリング

Vcc: Scaling Transformers to 128K Tokens or More by Prioritizing Important Tokens

May 7, 2023

著者: Zhanpeng Zeng, Cole Hawkins, Mingyi Hong, Aston Zhang, Nikolaos Pappas, Vikas Singh, Shuai Zheng

cs.AI

要旨

Transformerモデルは自然言語処理（NLP）やコンピュータビジョンの基盤となっています。近年、シーケンス長nの関数としての二次コストを削減するための様々な研究が行われていますが、超長いシーケンス（例えば16Kトークンを超えるもの）を効率的に扱うことは依然として困難です。例えば、本全体に基づいて質問に答えることや科学記事を要約するといったアプリケーションは非効率的か、あるいは実現不可能です。本論文では、Transformerモデルの複雑さをnに依存しないサイズrの表現に圧縮することで、nへの依存を大幅に削減することを提案します。具体的には、多くのタスクにおいて最終的な予測に最も関連するのは特殊なトークンの小さなサブセット（我々はこれをVIPトークンと呼ぶ）であるという事実を利用し、これらのVIPトークンの表現を近似する際の影響に基づいて入力シーケンスを選択的に圧縮するVIPトークン中心圧縮（Vcc）スキームを提案します。競合するベースラインと比較して、提案アルゴリズムは効率的であるだけでなく（4Kおよび16Kの長さでベースラインよりも3倍以上の効率向上を達成）、多数のタスクで競争力のあるまたはそれ以上の性能を達成します。さらに、我々のアルゴリズムは128Kトークン（またはそれ以上）にスケールアップ可能であり、一貫して精度向上を提供することを示します。

English

Transformer models are foundational to natural language processing (NLP) and computer vision. Despite various recent works devoted to reducing the quadratic cost of such models (as a function of the sequence length n), dealing with ultra long sequences efficiently (e.g., with more than 16K tokens) remains challenging. Applications such as answering questions based on an entire book or summarizing a scientific article are inefficient or infeasible. In this paper, we propose to significantly reduce the dependency of a Transformer model's complexity on n, by compressing the input into a representation whose size r is independent of n at each layer. Specifically, by exploiting the fact that in many tasks, only a small subset of special tokens (we call VIP-tokens) are most relevant to the final prediction, we propose a VIP-token centric compression (Vcc) scheme which selectively compresses the input sequence based on their impact on approximating the representation of these VIP-tokens. Compared with competitive baselines, the proposed algorithm not only is efficient (achieving more than 3times efficiency improvement compared to baselines on 4K and 16K lengths), but also achieves competitive or better performance on a large number of tasks. Further, we show that our algorithm can be scaled to 128K tokens (or more) while consistently offering accuracy improvement.