Scoprire i Gioielli nei Primi Strati: Accelerare LLM a Lungo Contesto con una Riduzione di 1000 Volte dei Token di Input

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

September 25, 2024

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

cs.AI

Abstract

I Large Language Models (LLM) hanno dimostrato capacità notevoli nel gestire input con contesti lunghi, ma ciò comporta un aumento delle risorse computazionali e della latenza. La nostra ricerca introduce un approccio innovativo per superare il collo di bottiglia del contesto lungo al fine di accelerare l'elaborazione inferenziale degli LLM e ridurre il consumo di memoria GPU. La nostra ricerca dimostra che gli LLM possono identificare i token rilevanti nei primi strati prima di generare risposte a una query. Sfruttando questa intuizione, proponiamo un algoritmo che utilizza i primi strati di un LLM come filtri per selezionare e comprimere i token di input, riducendo significativamente la lunghezza del contesto per l'elaborazione successiva. Il nostro metodo, GemFilter, mostra miglioramenti sostanziali sia in velocità che in efficienza della memoria rispetto alle tecniche esistenti, come l'attenzione standard e SnapKV/H2O. In particolare, ottiene un aumento di velocità del 2,4 volte e una riduzione del 30\% nell'uso della memoria GPU rispetto ai metodi SOTA. La valutazione sul compito Needle in a Haystack mostra che GemFilter supera significativamente l'attenzione standard, SnapKV e dimostra prestazioni comparabili nella sfida LongBench. GemFilter è semplice, non richiede addestramento ed è ampiamente applicabile su diversi LLM. In modo cruciale, fornisce interpretabilità consentendo agli esseri umani di ispezionare la sequenza di input selezionata. Questi risultati non solo offrono benefici pratici per l'implementazione degli LLM, ma migliorano anche la nostra comprensione dei meccanismi interni degli LLM, aprendo la strada a ulteriori ottimizzazioni nel design e nell'elaborazione inferenziale degli LLM. Il nostro codice è disponibile su 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.