One Pass Streaming Algorithm for Super Long Token Attention Approximation in Sublinear Space

要約

長いコンテキストを含むストリーミング アプリケーション、特に拡張ダイアログやテキスト分析に大規模言語モデル (LLM) を導入することは非常に重要ですが、2 つの大きな課題があります。
まず、以前のトークンのキーと値の状態 (KV) がキャッシュされるため、デコード段階では大量のメモリが消費されます。
第 2 に、アテンションの計算には時間がかかり、各トークンの生成の時間計算量は $O(n^2)$ になります。
最近の OpenAI DevDay (2023 年 11 月 6 日) で、OpenAI は 128K の長さのドキュメントをサポートできる新しいモデルをリリースしました。私たちの論文では、コンテキストの長さ $n$ が 128K よりはるかに大きい場合のメモリ効率の問題に焦点を当てています。
($n \gg 2^d$)。
クエリ、キー、および値の行列 $Q、K、V \in \mathbb{R}^{n \times d}$ を使用した単層セルフアテンションを考慮すると、多項式法はアテンション出力 $T \in \ を近似します。
mathbb{R}^{n \times d}$。
これは、$n^{1+o 内でのアテンション ${\sf Attn}(Q, K, V)$ の計算を促進するために $U_1, U_2 \in \mathbb{R}^{n \times t}$ を構築することでこれを実現します。
(1)}$ 回の実行。
それにもかかわらず、キー行列 $K と値行列 $K, V \in \mathbb{R}^{n \times d}$ を保存するには、依然として $O( n d)$ スペースが必要であり、大量のメモリ使用量につながります。
これらの課題に対応して、ストリーミング形式でデータの 1 パスのみを読み取る新しいアルゴリズムを導入しました。
この方法では、サブリニア空間 $o(n)$ を使用して 3 つのスケッチ行列を保存し、正確な $K、V$ ストレージの必要性を軽減します。
特に、私たちのアルゴリズムは、超長いトークンで優れたメモリ効率の高いパフォーマンスを示します。
トークンの長さ $n$ が増加すると、メモリ使用量はほぼ一定のままですが、エラー保証は減少します。
このユニークな属性は、ストリーミング アプリケーションで LLM を効率的に処理する際の私たちの技術の可能性を強調しています。

要約(オリジナル)

Deploying Large Language Models (LLMs) in streaming applications that involve long contexts, particularly for extended dialogues and text analysis, is of paramount importance but presents two significant challenges. Firstly, the memory consumption is substantial during the decoding phase due to the caching of Key and Value states (KV) of previous tokens. Secondly, attention computation is time-consuming with a time complexity of $O(n^2)$ for the generation of each token. In recent OpenAI DevDay (Nov 6, 2023), OpenAI released a new model that is able to support a 128K-long document, in our paper, we focus on the memory-efficient issue when context length $n$ is much greater than 128K ($n \gg 2^d$). Considering a single-layer self-attention with Query, Key, and Value matrices $Q, K, V \in \mathbb{R}^{n \times d}$, the polynomial method approximates the attention output $T \in \mathbb{R}^{n \times d}$. It accomplishes this by constructing $U_1, U_2 \in \mathbb{R}^{n \times t}$ to expedite attention ${\sf Attn}(Q, K, V)$ computation within $n^{1+o(1)}$ time executions. Despite this, storing the Key and Value matrices $K, V \in \mathbb{R}^{n \times d}$ still necessitates $O( n d)$ space, leading to significant memory usage. In response to these challenges, we introduce a new algorithm that only reads one pass of the data in streaming fashion. This method employs sublinear space $o(n)$ to store three sketch matrices, alleviating the need for exact $K, V$ storage. Notably, our algorithm exhibits exceptional memory-efficient performance with super-long tokens. As the token length $n$ increases, our error guarantee diminishes while the memory usage remains nearly constant. This unique attribute underscores the potential of our technique in efficiently handling LLMs in streaming applications.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google