初期レイヤーの中の宝石を発見する：1000倍の入力トークン削減による長文脈LLMの高速化

Discovering the Gems in Early Layers: Accelerating Long-Context LLMs with 1000x Input Token Reduction

September 25, 2024

著者: Zhenmei Shi, Yifei Ming, Xuan-Phi Nguyen, Yingyu Liang, Shafiq Joty

cs.AI

要旨

大規模言語モデル（LLM）は、長いコンテキスト入力を処理する際に顕著な能力を示していますが、これには計算リソースとレイテンシーの増加が伴います。当研究では、LLMの推論を加速し、GPUメモリ消費を削減するための長いコンテキストのボトルネックに対処する革新的なアプローチを紹介しています。当研究では、LLMがクエリに回答する前に初期層で関連するトークンを特定できることを実証しています。この洞察を活用して、LLMの初期層をフィルターとして使用し、入力トークンを選択して圧縮するアルゴリズムを提案しています。これにより、後続の処理のためのコンテキスト長が大幅に削減されます。当メソッドであるGemFilterは、標準の注意機構やSnapKV/H2Oなどの既存の手法と比較して、速度とメモリ効率の両方で著しい改善が示されます。特に、SOTA手法と比較して2.4倍の高速化とGPUメモリ使用量の30％削減が実現されます。針の中のハイスタックタスクでの評価によると、GemFilterは標準の注意機構やSnapKVを大幅に上回り、LongBenchチャレンジでも同等のパフォーマンスを示します。GemFilterはシンプルでトレーニング不要であり、異なるLLMに広く適用可能です。重要なのは、選択された入力シーケンスを人間が検査できるようにすることで、解釈可能性を提供します。これらの知見は、LLMの展開に実用的な利点をもたらすだけでなく、LLMの内部メカニズムの理解を向上させ、LLMの設計と推論のさらなる最適化の道を開くものです。当コードはhttps://github.com/SalesforceAIResearch/GemFilter で入手可能です。

English

Large Language Models (LLMs) have demonstrated remarkable capabilities in handling long context inputs, but this comes at the cost of increased computational resources and latency. Our research introduces a novel approach for the long context bottleneck to accelerate LLM inference and reduce GPU memory consumption. Our research demonstrates that LLMs can identify relevant tokens in the early layers before generating answers to a query. Leveraging this insight, we propose an algorithm that uses early layers of an LLM as filters to select and compress input tokens, significantly reducing the context length for subsequent processing. Our method, GemFilter, demonstrates substantial improvements in both speed and memory efficiency compared to existing techniques, such as standard attention and SnapKV/H2O. Notably, it achieves a 2.4times speedup and 30\% reduction in GPU memory usage compared to SOTA methods. Evaluation on the Needle in a Haystack task shows that GemFilter significantly outperforms standard attention, SnapKV and demonstrates comparable performance on the LongBench challenge. GemFilter is simple, training-free, and broadly applicable across different LLMs. Crucially, it provides interpretability by allowing humans to inspect the selected input sequence. These findings not only offer practical benefits for LLM deployment, but also enhance our understanding of LLM internal mechanisms, paving the way for further optimizations in LLM design and inference. Our code is available at https://github.com/SalesforceAIResearch/GemFilter.