H1B-KV: Hybrid One-Bit Caches for Memory-Efficient Large Language Model Inference
作者: Harshil Vejendla
分类: cs.CL, cs.LG
发布日期: 2025-10-07
备注: MIT URTC 2025 Technical Paper (Oral), 5 pages, 1 figure
💡 一句话要点
H1B-KV:提出混合一位缓存,解决LLM长文本推理中的内存瓶颈问题。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 KV缓存 内存压缩 量化 二进制草图
📋 核心要点
- 长文本LLM推理面临内存瓶颈,现有方法如量化、token剔除或仅键草图压缩不彻底,效果受限。
- H1B-KV采用混合压缩策略,键使用1位二进制草图,值使用4位量化,大幅降低内存占用且保留上下文信息。
- 实验表明,H1B-KV在多种任务上经过微调后,性能与全精度模型相当,且优于现有压缩方法。
📝 摘要(中文)
大型语言模型(LLM)的自回归解码需要缓存不断增长的过去键值(KV)对,这使得长上下文推理成为一个内存受限的问题。虽然最近的方法探索了量化缓存、逐出token或使用二进制草图表示键(例如,Loki),但这些方法通常提供不完整的解决方案,因为它们将一个组件(如值)保持未压缩状态或丢弃上下文信息。本文介绍了混合一位KV缓存(H1B-KV),这是一种全面的压缩方案,可以大幅减少内存使用,而不会牺牲上下文。H1B-KV使用1位二进制草图表示每个键向量,从而实现硬件友好的按位注意力,并使用4位量化进一步压缩值向量。这种整体的混合方法使一个70亿参数的LLM能够处理8k token的上下文,而缓存内存低于60 MB,减少了70倍。我们证明,经过轻量级微调后,H1B-KV不仅在困惑度基准测试中,而且在复杂的下游任务(如数学推理(GSM8K)、多任务理解(MMLU)和代码生成(HumanEval))中都与全精度性能相匹配。我们的结果表明,H1B-KV在每字节质量方面显著优于领先的量化(KIVI)、token逐出(SparseLLM)和仅键草图(Loki)方法,使其成为在内存受限环境中部署LLM的强大解决方案。
🔬 方法详解
问题定义:大型语言模型在长文本推理时,需要缓存大量的Key-Value对,导致内存占用过高,成为性能瓶颈。现有的压缩方法,例如量化、Token Eviction和Key-only Sketching,要么压缩不彻底,要么会损失上下文信息,无法在保证性能的同时显著降低内存占用。
核心思路:H1B-KV的核心思路是采用一种混合压缩方案,对Key和Value分别使用不同的压缩方法。Key使用1-bit二进制草图进行压缩,Value使用4-bit量化进行压缩。这种混合方法可以在大幅降低内存占用的同时,尽可能地保留上下文信息,从而保证模型的性能。
技术框架:H1B-KV的技术框架主要包含两个部分:Key的1-bit二进制草图压缩和Value的4-bit量化压缩。在推理过程中,首先使用1-bit二进制草图计算Attention权重,然后使用4-bit量化的Value计算最终的输出。整个过程可以无缝集成到现有的LLM架构中。
关键创新:H1B-KV的关键创新在于提出了一种混合压缩方案,将1-bit二进制草图和4-bit量化结合起来,实现了在内存占用和性能之间的平衡。与现有方法相比,H1B-KV能够更有效地降低内存占用,同时保持模型的性能。
关键设计:Key的1-bit二进制草图压缩使用随机投影的方法,将高维的Key向量投影到1-bit空间。Value的4-bit量化使用线性量化的方法,将Value向量量化到4-bit空间。此外,论文还提出了一种轻量级的微调方法,用于恢复压缩带来的性能损失。
📊 实验亮点
H1B-KV在70亿参数的LLM上实现了70倍的内存压缩,使其能够在低于60MB的缓存中处理8k token的上下文。经过轻量级微调后,H1B-KV在GSM8K、MMLU和HumanEval等下游任务上与全精度模型性能相当,且显著优于KIVI、SparseLLM和Loki等现有方法。
🎯 应用场景
H1B-KV适用于内存受限的设备上部署大型语言模型,例如移动设备、嵌入式系统等。该技术可以降低LLM的部署成本,并使其能够在更广泛的场景中使用,例如智能助手、自动翻译、文本生成等。未来,H1B-KV可以进一步扩展到其他类型的模型和任务中。
📄 摘要(原文)
Autoregressive decoding in large language models (LLMs) requires caching a growing list of past key-value (KV) pairs, making long-context inference a memory-bound problem. While recent methods have explored quantizing the cache, evicting tokens, or using binary sketches for keys (e.g., Loki), these approaches often provide an incomplete solution by leaving one component (like values) uncompressed or by discarding context information. This paper introduces the Hybrid One-Bit KV Cache (H1B-KV), a comprehensive compression scheme that radically reduces memory usage without sacrificing context. H1B-KV represents each key vector using a 1-bit binary sketch, enabling hardware-friendly bitwise attention, and further compresses value vectors using 4-bit quantization. This holistic, hybrid approach allows a 7-billion parameter LLM to handle an 8k-token context with under 60 MB of cache memory - a 70x reduction. We demonstrate that after a lightweight finetuning, H1B-KV matches full-precision performance not only on perplexity benchmarks but also on complex downstream tasks like mathematical reasoning (GSM8K), multi-task understanding (MMLU), and code generation (HumanEval). Our results show H1B-KV significantly outperforms leading quantization (KIVI), token eviction (SparseLLM), and key-only sketching (Loki) methods in quality-per-byte, establishing it as a robust solution for deploying LLMs in memory-constrained environments.