Обнаружение драгоценностей в ранних слоях: ускорение LLM с длинным контекстом с сокращением входных токенов в 1000 раз

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. Важно отметить, что он достигает увеличения скорости в 2,4 раза и снижения использования памяти GPU на 30\% по сравнению с методами SOTA. Оценка на задаче "Иголка в стоге сена" показывает, что 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.