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

TL;DR

ROSA-Tuning addresses the tension between state coverage and computational efficiency in long-context processing of large language models by introducing a retrieval-and-recall pathway (ROSA) that runs on CPU in parallel with the attention block. It uses a suffix automaton-based retrieval to identify relevant historical suffixes and injects their values into the model state via a trainable injection term, while keeping the core attention to a windowed mechanism, yielding $O(TW)$ GPU complexity. The method includes a binary discretization scheme and a counterfactual gradient to enable end-to-end training, plus an asynchronous CPU–GPU pipeline to optimize throughput. Experiments on Qwen3-Base-1.7B show ROSA-Tuning substantially restores long-context modeling of windowed-attention models, approaching or matching global attention while preserving similar GPU memory and speed to windowed methods. The work provides a practical route to efficient long-sequence processing with strong recall capabilities.

Abstract

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 introduces 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 design a binary discretization strategy and a 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.

Figures (2)

• Figure 1: ROSA-Tuning architecture
• Figure 2: Runtime comparison.