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

TL;DR

ROSA-Tuning introduces a retrieval-based mechanism that significantly improves long-context modeling in pretrained language models, achieving near-global attention performance with high efficiency.

Contribution

It proposes ROSA-Tuning, a novel retrieval-and-recall method that enhances long-context understanding in models without high computational costs.

Findings

Restores long-context modeling close to global attention.

Maintains computational efficiency comparable to windowed-attention.

Achieves strong performance on LongBench benchmarks.

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 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.