arXiv:2602.02499cs.CL2026-02

通过检索历史相关片段,让模型高效处理长文本。

ROSA-Tuning: Enhancing Long-Context Modeling via Suffix Matching

  • 用可训练的检索模块找长文本中相关历史位置
  • 在1.7B模型上接近全局注意力效果,计算开销低
  • 适合需要长文本理解又怕显存不足的研究者

长上下文建模与计算效率是当前大语言模型的核心挑战。现有高效注意力方法虽降低复杂度,但通常覆盖历史信息有限。本文提出ROSA-Tuning,一种基于检索与召回机制的长上下文增强方法。在标准注意力之外,该方法并行引入基于CPU的ROSA(RWKV在线后缀自动机)检索模块,高效定位长上下文中与当前查询相关的历史位置,并以可训练方式将信息注入模型状态,后续由范围受限注意力进行加权融合。为支持端到端训练,采用二值化策略与反事实梯度算法,并通过异步CPU-GPU流水线优化整体效率。在Qwen3-Base-1.7B上的系统评估表明,ROSA-Tuning显著恢复了窗口注意力模型的长上下文建模能力,在LongBench等基准上性能接近甚至匹配全局注意力,同时保持与窗口注意力相当的计算开销和显存占用,为高效长上下文处理提供了新路径。示例代码见https://github.com/zyaaa-ux/ROSA-Tuning。

原文摘要 · Abstract (English)

Long-context capability and computational efficiency are among the central challenges facing today's large language models. Existing efficient attention methods reduce computational complexity, but they typically suffer from a limited coverage of the model state. This paper proposes ROSA-Tuning, a retrieval-and-recall mechanism for enhancing the long-context modeling ability of pretrained models. Beyond the standard attention mechanism, ROSA-Tuning leverages in parallel a CPU-based ROSA (RWKV Online Suffix Automaton) retrieval module, which efficiently locates historical positions in long contexts that are relevant to the current query, and injects the retrieved information into the model state in a trainable manner; subsequent weighted fusion can then be handled by range-restricted attention. To enable end-to-end training, we employ the binary discretization strategy and the counterfactual gradient algorithm, and further optimize overall execution efficiency via an asynchronous CPU-GPU pipeline. Systematic evaluations on Qwen3-Base-1.7B show that ROSA-Tuning substantially restores the long-context modeling ability of windowed-attention models, achieving performance close to and in some cases matching global attention on benchmarks such as LongBench, while maintaining computational efficiency and GPU memory usage that are nearly comparable to windowed-attention methods, offering a new technical path for efficient long-context processing. The example code can be found at https://github.com/zyaaa-ux/ROSA-Tuning.

长文本高效推理注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。