提出一种无需额外参数的稀疏注意力机制,提升长文本推理时的压缩稳定性。
Lag-Relative Sparse Attention In Long Context Training
- 按窗口选取最相关键值对,实现高效稀疏注意力
- 在压缩上下文下显著减少性能下降,问答任务表现更优
- 适合需长文本推理且资源受限的模型微调场景
大语言模型在自然语言处理与生成方面取得显著进展,但其处理长上下文的能力仍受注意力计算的二次复杂度和键值记忆占用线性增长的限制。为降低计算成本与内存开销,通常在推理时采用键值缓存压缩技术,但常导致严重性能下降,因模型未针对压缩上下文进行训练。尽管存在更复杂的压缩方法,但多数不适用于后训练阶段,或因不兼容梯度优化、计算开销过高而受限。为此,我们提出基于LagKV压缩方法的滞后相对稀疏注意力(Lag-Relative Sparse Attention, LRSA),用于长上下文后训练。该方法通过分块预填充,从固定大小的滞后窗口中选择最相关的前K个键值对,使模型聚焦于关键历史信息,同时保持高效率。实验表明,该方法显著提升了模型在键值压缩下的鲁棒性,并在问答微调任务中取得了更优结果。
原文摘要 · Abstract (English)
Large Language Models (LLMs) have made significant strides in natural language processing and generation, yet their ability to handle long-context input remains constrained by the quadratic complexity of attention computation and linear-increasing key-value memory footprint. To reduce computational costs and memory, key-value cache compression techniques are commonly applied at inference time, but this often leads to severe performance degradation, as models are not trained to handle compressed context. Although there are more sophisticated compression methods, they are typically unsuitable for post-training because of their incompatibility with gradient-based optimization or high computation overhead. To fill this gap with no additional parameter and little computation overhead, we propose Lag-Relative Sparse Attention(LRSA) anchored by the LagKV compression method for long context post-training. Our method performs chunk-by-chunk prefilling, which selects the top K most relevant key-value pairs in a fixed-size lagging window, allowing the model to focus on salient historical context while maintaining efficiency. Experimental results show that our approach significantly enhances the robustness of the LLM with key-value compression and achieves better fine-tuned results in the question-answer tuning task.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。