语言模型推理的选择性状态空间适配与检索

# 语言模型推理的选择性状态空间适配与检索
来源：https://arxiv.org/html/2607.19326
Atahan Dokme Larry Heck AI虚拟助手(AVA)实验室 佐治亚理工学院 \{adokme3,larryheck\}@gatech.edu

###### 摘要

低秩适配引入了一个静态学习更新，该更新对所有输入完全相同。该更新提供任务级别的适配，但没有显式表示分词级别或实例级别的状态变化。本文提出了一系列适配器，在两个互补的粒度上引入选择性状态空间循环。在分词级别，MaLoRA（Mamba 调制的低秩适配）使适配器的缩放因子成为跨分词具有循环状态的动态输入相关函数，与先前工作的无状态调制器形成对比。在上下文级别，MaRA（Mamba 检索适配器）跟踪跨片段状态，并在调制语言模型生成其答案之前，选择与查询最相关的片段。在三个冻结主干网络（Qwen-2.5-7B、Llama-3.1-8B、Gemma-2-9B）和两个推理基准（MuSiQue、2WikiMultihopQA）上，该系列在 3×2 网格的每个单元格上都提高了推理准确性，平均提高 +6.8 F1（相对 +10.5%），在最难的单元格上高达 +9.3 F1（相对 +18.2%），优于 LoRA 基线。分词级别的增益在长度压力下延续到 RULER QA-2。

# 语言模型推理的选择性状态空间适配与检索

Atahan Dokme Larry Heck\_AI 虚拟助手 (AVA) 实验室
佐治亚理工学院
\{adokme3,larryheck\}@gatech.edu

## 1 引言

低秩适配 (Hu et al., 2022 (https://arxiv.org/html/2607.19326#bib.bib1)) 已成为调整冻结语言模型以适应下游任务的标准方法。适配器只训练一次，并统一应用于每个输入。对于具有多个潜在状态的任务，例如多跳推理，这个静态更新在所有状态之间共享，鼓励依赖训练集模式而不是针对每个输入的调整。本研究的目标是使适配更具输入适应性，并更好地与查询的状态结构对齐，同时保持主干网络冻结。

适配可以在三个级别上运行：在*任务*级别，LoRA 已经在此运行；在*分词*级别，适配的强度可以成为当前分词的函数；在*上下文*级别，模型可以关注与查询最相关的输入片段。

本工作的第一个贡献是 MaLoRA（Mamba 调制的低秩适配），一种用于低秩适配的循环分词级调制器。与以前使用无状态每分词缩放因子的分词调制器不同，MaLoRA 使用基于 Mamba 的选择性状态空间模块 (Gu and Dao, 2024 (https://arxiv.org/html/2607.19326#bib.bib2)) 在分词之间传播状态。第二个贡献是 MaRA（Mamba 检索适配器），一种*上下文级别*适配器，它通过一个在片段嵌入上操作的状态空间模块对每个输入片段进行评分，并为适配模型选择最相关的片段。

选择性状态空间循环在两个层次的适配中都有贡献。在分词级别，MaLoRA 的循环状态在主干网络上比无状态调制产生一致的增益 (表 1 (https://arxiv.org/html/2607.19326#S3.T1))。在上下文级别，状态空间模块执行证据选择，这是分词调制器无法做到的，因为它在段落级别对相关性不敏感 (第 3.4 节 (https://arxiv.org/html/2607.19326#S3.SS4))。因此，相同的 Mamba 块在同一个系统中扮演两个不同的角色。

该系列在 MuSiQue (Trivedi et al., 2022 (https://arxiv.org/html/2607.19326#bib.bib9)) 和 2WikiMultihopQA (Ho et al., 2020 (https://arxiv.org/html/2607.19326#bib.bib30)) 上，跨越三个冻结主干网络（Qwen-2.5-7B、Llama-3.1-8B、Gemma-2-9B）进行评估。RULER (Hsieh et al., 2024 (https://arxiv.org/html/2607.19326#bib.bib33)) 衡量在长度压力下保留基础模型能力的情况。分词级调制器和 MaRA 可以组合，组合后的系统在 3×2 网格的每个单元上都比 LoRA 有所改进，两种机制通过互补的渠道做出贡献：调制根据分词角色重塑低秩更新，MaRA 在生成之前选择与查询相关的片段，这些增益累积而非重叠。

## 2 相关工作

#### 参数高效适配。

低秩适配 (Hu et al., 2022 (https://arxiv.org/html/2607.19326#bib.bib1)) 通过学习到的静态更新来适配冻结的基础模型。扩展包括权重分解 (Liu et al., 2024 (https://arxiv.org/html/2607.19326#bib.bib5))、自适应秩分配 (Zhang et al., 2023 (https://arxiv.org/html/2607.19326#bib.bib4))、循环适配器 (Nguyen et al., 2024 (https://arxiv.org/html/2607.19326#bib.bib7))、非对称共享 AA 多头架构 (Tian et al., 2024 (https://arxiv.org/html/2607.19326#bib.bib26)) 以及混合 LoRA 路由 (Dou et al., 2024 (https://arxiv.org/html/2607.19326#bib.bib28); Gao et al., 2025 (https://arxiv.org/html/2607.19326#bib.bib27))。所有这些都共享一个假设，即更新是基础权重的静态重新参数化，而下面介绍的输入相关变体则放松了这一假设。

#### 分词感知与门控适配。

最近探索了在分词级别使低秩更新与输入相关。TopLoRA (Li et al., 2025 (https://arxiv.org/html/2607.19326#bib.bib48)) 在 LoRA 因子之间学习一个依赖于分词的*对角*矩阵，并作为本工作中的无状态先前技术基线。在此基线之上，本工作将*循环*分词调制器 (MaLoRA) 在分词之间传播状态，并将状态空间循环也置于*上下文片段*级别，在此处它执行模型附加的证据选择，将冻结生成器自身的隐藏状态作为输入，这是所有这些分词级别方法都没有涉及的设置。分词级别分析 (第 3.4 节 (https://arxiv.org/html/2607.19326#S3.SS4)) 表明，那里的调制是*对角色敏感但对相关性盲目*的，这促使将证据相关性决策移至片段级别。

#### 语言建模中的状态空间模型。

状态空间模型提供了次二次序列建模 (Gu et al., 2021 (https://arxiv.org/html/2607.19326#bib.bib3); Smith et al., 2022 (https://arxiv.org/html/2607.19326#bib.bib24); Gu and Dao, 2024 (https://arxiv.org/html/2607.19326#bib.bib2); Dao and Gu, 2024 (https://arxiv.org/html/2607.19326#bib.bib23))。先前有两个适配器应用了这种动态机制，并且都将循环置于适配器堆栈的*跨层*位置：Nguyen et al. (2024 (https://arxiv.org/html/2607.19326#bib.bib7)) (READ) 在适配器堆栈中跨层传播状态，Yu et al. (2025 (https://arxiv.org/html/2607.19326#bib.bib49)) (SSMLoRA) 使用 SSM 来互连插入到不同层相同模块类型中的 LoRA 矩阵。MaLoRA 在*轴*上有所不同：选择性状态空间模块是 LoRA 更新的*每分词门控*，循环是*在相同输入的不同分词之间*，因此每个分词的适配器强度都携带着到达它之前的轨迹记忆。

#### 推理与证据选择。

多跳 QA 基准 (Ho et al., 2020 (https://arxiv.org/html/2607.19326#bib.bib30); Trivedi et al., 2022 (https://arxiv.org/html/2607.19326#bib.bib9)) 暴露了语言模型拥有答案和使用支持证据之间的差距；长上下文评估揭示了相关的转变 (Hsieh et al., 2024 (https://arxiv.org/html/2607.19326#bib.bib33); Bai et al., 2024 (https://arxiv.org/html/2607.19326#bib.bib29); Laban et al., 2025 (https://arxiv.org/html/2607.19326#bib.bib36))。

#### 检索基线。

开放域和多跳 QA 中使用的大多数检索基线是外部模块。BM25 (Robertson and Zaragoza, 2009 (https://arxiv.org/html/2607.19326#bib.bib44)) 是一个稀疏检索器，密集或后期交互系统如 DPR (Karpukhin et al., 2020 (https://arxiv.org/html/2607.19326#bib.bib45))、Contriever (Izacard et al., 2021 (https://arxiv.org/html/2607.19326#bib.bib42))、E5 (Wang et al., 2022 (https://arxiv.org/html/2607.19326#bib.bib43))、ColBERT (Khattab and Zaharia, 2020 (https://arxiv.org/html/2607.19326#bib.bib41)) 以及最近的嵌入编码器 (Zhang et al., 2025 (https://arxiv.org/html/2607.19326#bib.bib40)) 都训练一个单独的检索模型，而不是重复使用下游语言模型的隐藏状态。本工作引入的检索适配器 MaRA 则是*模型附加*的：它从用于生成的同一个冻结主干网络的中间隐藏状态中对片段进行评分，然后在其上应用一个小型片段级状态空间模块。因此，它更接近于一个上下文级别的适配器，而不是传统的外包检索器。

## 3 MaLoRA：有状态的分词调制

参考说明图 1：从 LoRA 到 MaLoRA。(a) LoRA：静态低秩更新 BA x_t 加到冻结输出 W x_t 上。(b) TopLoRA (Li et al., 2025 (https://arxiv.org/html/2607.19326#bib.bib48))：在 LoRA 因子之间插入一个每分词对角调制器 Λ(x_t) = diag(exp(RMSNorm(P x_t)))，得到 h_t = (W + B Λ(x_t) A) x_t。调制器是无状态的：Λ(x_t) 仅依赖于当前分词。(c) MaLoRA：将每分词头部替换为一个在投影序列上的选择性状态空间模块（Mamba），输出是一个标量 λ(x_t) = softplus(φ(Mamba(P x_t)_t))，用于缩放 r 维中间变量。

### 3.1 预备知识

#### 低秩适配。

低秩适配 (Hu et al., 2022 (https://arxiv.org/html/2607.19326#bib.bib1)) 将对冻结权重 W0 ∈ R^{d_out × d_in} 的更新重新参数化为一个低秩乘积。对于输入 x ∈ R^{d_in}，输出变为

h = (W0 + B A) x,   A ∈ R^{r × d_in}, B ∈ R^{d_out × r},   (1)

其中 r ≪ min(d_in, d_out) 是秩。基础权重 W0 保持冻结；只训练 A 和 B。相同更新应用于每个输入的每个分词，因此适配的幅度独立于模型对当前分词的表示。

#### TopLoRA：无状态的每分词调制。

TopLoRA (Li et al., 2025 (https://arxiv.org/html/2607.19326#bib.bib48)) 通过在 LoRA 因子之间插入一个学习到的每分词对角调制器 Λ(x_t) ∈ R^{r × r}_{>0} 来放松这个静态假设：

h_t = (W0 + B Λ(x_t) A) x_t,   (2)

其中 Λ(x_t) = diag(λ(x_t))，且 λ(x_t) = exp(RMSNorm(P x_t)) ∈ R^{r}_{>0}，P ∈ R^{r × d_in} 是一个学习到的低秩投影（图 1 (https://arxiv.org/html/2607.19326#S3.F1)b）。调制器是*无状态的*：Λ(x_t) 仅依赖于当前分词，没有跨位置的传播。

### 3.2 MaLoRA

MaLoRA 用有状态的标量门控替换了 TopLoRA 的无状态每分词对角门。公式 (1) (https://arxiv.org/html/2607.19326#S3.E1) 中的静态低秩更新乘以一个由选择性状态空间模块 (Gu and Dao, 2024 (https://arxiv.org/html/2607.19326#bib.bib2)) 在投影序列上产生的每分词标量 λ(x_t) ∈ R_{>0}：

h_t = (W0 + B λ(x_t) A) x_t,   (3)
λ(x_t) = softplus(φ(o_t)),

其中 P ∈ R^{r × d_in} 是一个学习到的低秩投影，与 A 共享维度，φ 是一个小型学习头部，o_t 是 Mamba 在位置 t 的输出。Mamba 模块维护一个具有输入相关动态的隐藏状态 s_t ∈ R^{d_s}：

s_t = \bar{A}_t s_{t-1} + \bar{B}_t P x_t,   o_t = \bar{C}_t s_t,   (4)

其中 \bar{A}_t, \bar{B}_t, \bar{C}_t 按照标准 Mamba 参数化从 P x_t 计算得到（图 1 (https://arxiv.org/html/2607.19326#S3.F1)c）。当 λ(x_t) = 1 时，公式 (3) (https://arxiv.org/html/2607.19326#S3.E3) 简化为公式 (1) (https://arxiv.org/html/2607.19326#S3.E1)。

Mamba 块是核心贡献：它将选择性状态空间循环带入参数高效适配器内部，使每分词门控拥有轨迹 x_1, ..., x_t 的记忆，而不仅仅对 x_t 做出反应。其余设计选择如下。(i) 输出是单个标量 λ(x_t)，而不是 TopLoRA 的对角矩阵 Λ(x_t)，这样可以暴露每个分词的一个可解释值，用于第 3.4 节 (https://arxiv.org/html/2607.19326#S3.SS4) 的行为分析。(ii) 激活函数是 softplus，因此调制器既可以抑制 (λ < 1) 也可以放大 (λ > 1) LoRA 更新；激活函数的消融在附录 D (https://arxiv.org/html/2607.19326#A4) 中。

MaLoRA 与 LoRA 矩阵联合训练，附加到 {q_proj, k_proj, v_proj, up_proj, down_proj}，具有投影和层特定的参数。Mamba 模块将 LoRA 参数计数增加了约 50%；TopLoRA 基线增加了约 54%（其对角头部比 MaLoRA 的标量输出更宽）。具体计数在附录 C (https://arxiv.org/html/2607.19326#A3) 中。

### 3.3 调制器结果

使用了两个多跳 QA 数据集：MuSiQue (Trivedi et al., 2022 (https://arxiv.org/html/2607.19326#bib.bib9))（在 20 个候选片段中 2-4 跳）和 2WikiMultihopQA (Ho et al., 2020 (https://arxiv.org/html/2607.19326#bib.bib30))（在 10 个候选片段中最多 4 跳）。所有三个冻结主干网络（Qwen-2.5-7B、Llama-3.1-8B、Gemma-2-9B）在秩 r=16 下使用 EM/F1、beam 4、三个种子进行评估；完整超参数在附录 B (https://arxiv.org/html/2607.19326#A2) 中。MaLoRA 在表 1 (https://arxiv.org/html/2607.19326#S3.T1) 的每个单元上都比 LoRA 和 TopLoRA 基线有所改进，在 MuSiQue 上比 LoRA 增益最大（Llama 上高达 +7.9 F1）。循环状态的贡献在第 3.5 节 (https://arxiv.org/html/2607.19326#S3.SS5) 中剖析，相同的模式延续到常识推理（附录 H (https://arxiv.org/html/2607.19326#A8)）。

表 1：三个主干网络上的分词调制器。MuSiQue 和 2WikiMultihopQA 上的 F1，以及使用 MuSiQue 训练检查点的 RULER QA-2 准确率（%）（4k 上下文）。三个种子的均值；每个种子的标准差在附录 D (https://arxiv.org/html/2607.19326#A4) 中（表 7 (https://arxiv.org/html/2607.19326#A4.T7)）。M