MILO: Efficient Many-shot In-Context Learning with Block-wise Low-rank Compression

📄 arXiv: 2609.29913v1 📥 PDF

作者: Youpeng Zhao, Tian Tan, Liqian Peng, Jun Wang, Alec Go

分类: cs.CL

发布日期: 2026-09-24

备注: Technical Report

💡 一句话要点

提出MILO以解决多示例上下文学习中的KV缓存效率问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 多示例学习 低秩压缩 键值缓存 推理效率 动态秩预算

📋 核心要点

  1. 现有的多示例上下文学习方法在推理过程中面临KV缓存内存的效率瓶颈，导致在线服务和设备部署的挑战。
  2. 本文提出的MILO框架通过块级低秩压缩策略，利用多示例上下文的低秩冗余，显著提高KV缓存的存储效率。
  3. 实验结果显示，MILO在Qwen2.5模型上实现了50%的KV缓存内存减少和1.8倍的吞吐量提升，且性能降幅极小。

📝 摘要（中文）

多示例上下文学习（ICL）使大型语言模型能够通过数千个示例适应复杂任务，但这种方法使得推理效率瓶颈转移至键值（KV）缓存内存。由于KV缓存的线性扩展特性，存储这些中间张量成为在线服务和设备部署的重大挑战。为了解决这一问题，本文提出了一种新颖的压缩框架MILO，利用多示例上下文中固有的低秩冗余。MILO采用块级低秩压缩策略，在块粒度上压缩KV缓存，并根据信息熵动态分配秩预算，以保持关键块的保真度，同时对冗余块进行强压缩。实验结果表明，MILO在Qwen2.5模型上实现了KV缓存内存减少50%和吞吐量提升1.8倍，分类和推理基准测试中性能降幅微乎其微，显著优于现有基线。

🔬 方法详解

问题定义：本文旨在解决多示例上下文学习中KV缓存内存效率低下的问题。现有方法在处理大量示例时，KV缓存的线性扩展特性导致内存占用过高，影响在线服务和设备部署的可行性。

核心思路：MILO框架的核心思路是利用多示例上下文中的低秩冗余，通过块级低秩压缩策略来减少KV缓存的内存占用，同时保持重要信息的完整性。

技术框架：MILO的整体架构包括数据预处理、块级压缩和动态秩预算分配三个主要模块。首先，对输入的多示例上下文进行分块处理，然后应用低秩压缩技术，最后根据信息熵动态调整每个块的秩预算。

关键创新：MILO的主要创新在于块级低秩压缩策略和动态秩预算分配机制。这种设计使得在压缩冗余块的同时，能够有效保留关键块的信息，显著提高了压缩效率。

关键设计：在MILO中，块的大小和秩预算的分配是关键设计参数。通过信息熵评估每个块的重要性，MILO能够在保证性能的前提下，灵活调整压缩策略。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MILO在Qwen2.5模型上实现了高达50%的KV缓存内存减少，同时吞吐量提升了1.8倍。与现有基线相比，MILO在分类和推理基准测试中几乎没有性能降幅，显示出其优越性。

🎯 应用场景

MILO框架在大型语言模型的推理过程中具有广泛的应用潜力，尤其适用于需要处理大量示例的复杂任务，如自然语言处理、机器翻译和对话系统。通过提高KV缓存的存储效率，MILO能够促进在线服务和移动设备上的AI应用，降低资源消耗，提升用户体验。

📄 摘要（原文）

Many-shot in-context learning (ICL) enables large language models (LLMs) to adapt to complex tasks by conditioning on thousands of demonstration examples, but this paradigm shifts the inference efficiency bottleneck to the key-value (KV) cache memory. Due to the linear scaling behavior of the KV cache, storing these intermediate tensors has become a paramount challenge for both online serving and on-device deployment. To address this issue, we propose a novel compression framework, termed MILO, that exploits the low-rank redundancy inherent in many-shot contexts. Specifically, MILO features a block-wise low-rank compression strategy that compresses the KV cache at the block granularity, where each block contains multiple many-shot examples. Furthermore, to handle the heterogeneous context density across different blocks, MILO dynamically allocates rank budgets based on the information entropy, preserving the fidelity of critical blocks while aggressively compressing redundant ones. Experimental results on Qwen2.5 models demonstrate that our method achieves up to 50% reduction in KV cache memory and 1.8x throughput improvement, with negligible performance degradation on classification and reasoning benchmarks, significantly outperforming prior baselines.