用前缀树共享缓存，让大模型生成更快更省内存。
Efficient Beam Search for Large Language Models Using Trie-Based Decoding
- 构建前缀树共享多个分支的键值缓存，减少重复存储。
- 内存降低4到8倍，推理速度最快提升2.4倍。
- 适合资源有限或大规模部署的生成任务。
本文提出一种基于前缀树（trie）的并行解码方法，解决批量束搜索中的内存效率问题。通过在具有相同前缀的多个束之间共享单一键值缓存，该方法显著降低内存占用并实现高效解码。我们在三种注意力架构上进行了评估：多头注意力（Phi-3.5-mini-instruct）、分组查询注意力（Llama-3.1-8B-Instruct）和滑动窗口注意力（Mistral-Small-24B-Instruct-2501），分别使用CNN/DailyMail数据集进行摘要生成，以及HumanEval数据集进行代码生成。实验表明，该方法在不损失生成质量的前提下，内存消耗减少4至8倍，解码速度最高提升2.4倍，证明其在内存受限环境及大规模部署中的适用性。
原文摘要 · Abstract (English)
This work presents a novel trie (prefix-tree)-based parallel decoding method that addresses the memory inefficiency of batch-based beam search. By sharing a single KV cache across beams with common prefixes, our approach dramatically reduces memory usage and enables efficient decoding. We evaluated our method across three attention architectures, Multi-Head Attention (Phi-3.5-mini-instruct), Grouped Query Attention (Llama-3.1-8B-Instruct), and Sliding Window Attention (Mistral-Small-24B-Instruct-2501), using CNN/DailyMail for abstractive summarization and HumanEval for code generation. Our experiments demonstrate substantial memory savings (4--8$\times$) and up to 2.4$\times$ faster decoding, without compromising generation quality. These results highlight our method's suitability for memory-constrained environments and large-scale deployments.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。