Loki: Low-Rank Keys for Efficient Sparse Attention

要約

大規模な言語モデルでの推論は、特に長いシーケンス長が使用される場合、関連するコンピューティング コストとメモリ コストの点で高価になる可能性があります。
特に、そのようなモデルで使用される自己注意メカニズムはこれらのコストに大きく寄与しており、その結果、推論のための疎な注意近似を提案するいくつかの最近の研究が行われています。
この研究では、アテンション ブロックで計算されるキー ベクトルの次元に焦点を当てることによってセルフ アテンションの計算を近似することを提案します。
私たちの分析により、重要なベクトルは、いくつかのデータセットとモデルにわたって一貫して、大幅に低次元の空間に存在することが明らかになりました。
この観察を利用して、低次元空間で計算された注意スコアに基づいて KV キャッシュ内のトークンをランク付けして選択する新しいスパース注意方法である Loki を提案します。
私たちの評価では、Loki が他の一般的な近似手法よりもモデルの有効性を維持できる一方で、データの移動 (ロード/ストア) と計算コストの削減によりアテンションの計算を高速化できることが示されています。

要約(オリジナル)

Inference on large language models can be expensive in terms of the compute and memory costs involved, especially when long sequence lengths are used. In particular, the self-attention mechanism used in such models contributes significantly to these costs, which has resulted in several recent works that propose sparse attention approximations for inference. In this work, we propose to approximate the self-attention computation by focusing on the dimensionality of key vectors computed in the attention block. Our analysis reveals that the key vectors lie in a significantly lower-dimensional space, consistently across several datasets and models. Exploiting this observation, we propose Loki, a novel sparse attention method that ranks and selects tokens in the KV-cache based on attention scores computed in low-dimensional space. Our evaluations show that Loki is able to maintain the efficacy of the models better than other popular approximation methods, while speeding up the attention computation due to reduced data movement (load/store) and compute costs.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google