Selective State-Space Adaptation and Retrieval for Language Model Reasoning

2026-07-21

AI Summary

Existing low-rank adaptation methods struggle to capture dynamic state variations at both the token and instance levels, limiting the performance of frozen large language models on complex multi-hop reasoning tasks. This work proposes a novel adapter family that introduces a Mamba-based selective state space mechanism at the token level for dynamic recurrent state modulation (MaLoRA) and designs a cross-segment state tracking and retrieval module at the context level (MaRA). To our knowledge, this is the first approach to integrate selective state spaces into low-rank adaptation, enabling input-dependent dynamic state modeling across two granularities. Experiments show that when integrated into frozen models such as Qwen, Llama, and Gemma, the method achieves an average F1 improvement of 6.8 (+10.5%) on MuSiQue and 2WikiMultihopQA, with gains up to 9.3 F1 (+18.2%), while maintaining significant advantages on the RULER QA-2 long-context benchmark.

Abstract

Low-rank adaptation introduces a static learned update applied identically to every input. The update provides task-level adaptation but does not explicitly represent token-level or instance-level state variation. A family of adapters is proposed that introduces selective state-space recurrence at two complementary granularities. At the token level, \textbf{MaLoRA} (Mamba-modulated low-rank adaptation) makes the adapter's scaling factor a dynamic input-dependent function with recurrent state across tokens, in contrast to the stateless modulators of prior work. At the context level, \textbf{MaRA} (Mamba Retrieval Adapter) tracks cross-segment state and selects the segments most relevant to the query, before the modulated language model generates its answer. Across three frozen backbones (Qwen-2.5-7B, Llama-3.1-8B, Gemma-2-9B) and two reasoning benchmarks (MuSiQue, 2WikiMultihopQA), the family improves reasoning accuracy on every cell of the $3{\times}2$ grid, by $+6.8$ F1 ($+10.5\%$ relative) on average and up to $+9.3$ F1 ($+18.2\%$ relative) on the hardest cell over the LoRA baseline, and the token-level gains carry to RULER QA-2 under length stress.

Problem

• low-rank adaptation
• state-space recurrence
• token-level adaptation
• context-level state
• language model reasoning

Innovation

• Selective State-Space Adaptation
• MaLoRA
• MaRA
• Low-Rank Adaptation
• Reasoning Enhancement