Selective State-Space Adaptation and Retrieval for Language Model Reasoning
2026-07-21 • Computation and Language
Computation and Language
AI summaryⓘ
The authors propose a new method to improve how language models adapt to different tasks by making the adaptation depend on each word (token) and the overall context. Their approach, called MaLoRA, changes the way the model processes each token by adding memory across tokens, unlike previous methods that treated tokens independently. They also introduce MaRA, which helps the model pick relevant parts of the input before answering. Tests on several large language models and reasoning tasks show that their methods improve accuracy compared to traditional static adaptation techniques.
low-rank adaptationlanguage modelstoken-level adaptationrecurrent statescontext-level retrievalMaLoRAMaRAtask adaptationreasoning benchmarksfrozen backbone models
Authors
Atahan Dokme, Larry Heck
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.