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

2026-01-14

AI Summary

Existing efficient attention mechanisms struggle to balance coverage and computational efficiency in long-context modeling. This work proposes ROSA-Tuning, the first approach to integrate a Run-Length Online Suffix Automaton (ROSA) with large language models. ROSA-Tuning constructs ROSA structures in parallel on the CPU to retrieve historically relevant token positions and injects this retrieved information into the model’s internal states in a trainable manner, fusing it with windowed attention through learned weighting. The method incorporates binary discretization, counterfactual gradients, and an asynchronous CPU-GPU pipeline to enable end-to-end training. Evaluated on Qwen3-Base-1.7B, ROSA-Tuning substantially recovers the long-range modeling capability of windowed attention, achieving performance on LongBench and other benchmarks that approaches or matches full global attention, while maintaining computational efficiency and GPU memory usage comparable to windowed attention.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• long-context modeling
• computational efficiency
• large language models
• attention mechanism
• windowed attention

Innovation

• ROSA-Tuning
• long-context modeling
• suffix automaton
• efficient attention
• asynchronous CPU-GPU pipeline