Selective State-Space Adaptation and Retrieval for Language Model Reasoning

2026/07/21 by Atahan Dokme, Larry Heck
#cs.CL

paper · pdf

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

Citations

• LLMs Get Lost In Multi-Turn Conversation
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Related

• RE-Adapt: Reverse Engineered Adaptation of Large Language Models
• SBoRA: Low-Rank Adaptation with Regional Weight Updates
• Data Efficient Adaptation in Large Language Models via Continuous Low-Rank Fine-Tuning
• HEFT: A Coarse-to-Fine Hierarchy for Enhancing the Efficiency and Accuracy of Language Model Reasoning
• A Comparative Analysis of LLM Adaptation: SFT, LoRA, and ICL in Data-Scarce Scenarios