2025-11-15T12:52:11.146335

ReMamba: Equip Mamba with Effective Long-Sequence Modeling

Yuan, Liu, Li et al.
While the Mamba architecture demonstrates superior inference efficiency and competitive performance on short-context natural language processing (NLP) tasks, empirical evidence suggests its capacity to comprehend long contexts is limited compared to transformer-based models. In this study, we investigate the long-context efficiency issues of the Mamba models and propose ReMamba, which enhances Mamba's ability to comprehend long contexts. ReMamba incorporates selective compression and adaptation techniques within a two-stage re-forward process, incurring minimal additional inference costs overhead. Experimental results on the LongBench and L-Eval benchmarks demonstrate ReMamba's efficacy, improving over the baselines by 3.2 and 1.6 points, respectively, and attaining performance almost on par with same-size transformer models.
academic

ReMamba: Equip Mamba with Effective Long-Sequence Modeling

基本信息

  • 论文ID: 2408.15496
  • 标题: ReMamba: Equip Mamba with Effective Long-Sequence Modeling
  • 作者: Danlong Yuan, Jiahao Liu, Bei Li, Huishuai Zhang, Jingang Wang, Xunliang Cai, Dongyan Zhao
  • 分类: cs.CL (Computation and Language)
  • 发表时间: 2024年8月 (arXiv预印本)
  • 论文链接: https://arxiv.org/abs/2408.15496
  • 代码链接: https://github.com/lblankl/ReMamba

摘要

本文针对Mamba架构在长上下文理解任务中的性能不足问题，提出了ReMamba方法。虽然Mamba在短上下文NLP任务中表现出色且推理效率高，但在处理长上下文时其性能明显低于Transformer模型。ReMamba通过两阶段重前向过程中的选择性压缩和适应技术来增强Mamba的长上下文理解能力，仅引入最小的额外推理开销。在LongBench和L-Eval基准测试中，ReMamba分别比基线模型提升了3.2和1.6个点，性能接近同等规模的Transformer模型。

研究背景与动机

问题定义

  1. 核心问题: Mamba模型在处理长上下文（2k tokens以上）时性能显著下降，无法有效保持远距离信息
  2. 重要性: 长上下文理解是大语言模型发展的关键能力，对于文档理解、对话系统等应用至关重要
  3. 现有方法局限性:
    • Transformer面临二次计算复杂度和线性内存消耗问题
    • 混合架构虽然缓解了问题但降低了计算效率
    • 现有的Mamba改进方法（如LongMamba、DeciMamba）效果有限

研究动机

作者通过实验发现，Mamba在短上下文任务上超越了同等规模的Transformer，但在长上下文任务上存在显著性能差距。这种RNN-like架构的固定状态空间限制了其保存远距离信息的能力，导致信息遗忘问题严重。

核心贡献

  1. 识别问题根源: 通过初步研究发现Mamba的信息丢失问题严重，甚至随机压缩都能获得相似性能
  2. 提出ReMamba方法: 设计了两阶段选择性压缩和适应机制，有效缓解长上下文信息丢失
  3. 实现显著性能提升: 在LongBench和L-Eval上分别提升3.2和1.6个点，接近Transformer性能
  4. 保持效率优势: 仅增加一次前向传播的开销，维持恒定内存消耗和高推理速度
  5. 方法通用性: 成功扩展到Mamba2架构，证明方法的普适性

方法详解

任务定义

输入: 长上下文序列 {ti}^L_，其中L为序列长度 输出: 基于长上下文的自然语言生成结果 目标: 在保持Mamba推理效率的同时，提升其长上下文理解能力

模型架构

ReMamba采用两阶段架构设计：

Stage 1: 选择性压缩 (Selective Compression)

压缩范围定义:

  • 相对压缩范围: range := (s, e)，其中 e = s + p
  • 绝对索引集合: R := S, E，其中 S = L·s+1, E = L·(s+p)
  • 压缩比率: ρ，最终保留 K := |R|·ρ 个隐藏表示

重要性评分机制:

q = Query(hL)
{ki}^E_{i=S} = Key({hi}^E_{i=S})
cosi = (ki · q) / max(||ki||2 · ||q||2, ε)

Top-K选择:

G = argmax_{A⊂{S,S+1,...,E},|A|=K} Σ_{i∈A} cosi

压缩表示生成:

{vi}^K_{i=1} = Value({hj}, j ∈ G)
Tnew = Cat({ti}^{S-1}_{i=1}, {vi}^K_{i=1}, {ti}^L_{i=E+1})

Stage 2: 选择性适应 (Selective Adaptation)

对于选中的隐藏状态，修改Mamba的选择性机制：

α = ReLU(cos'_{t-1})
Δ^l_{t-1}' = Proj1(h^{l-1}_{t-1})
δ = Δ^l_{t-1}' · α + Θ^l
Δ^l_{t-1} = Softplus(δ)

其中Θ^l是可训练的层级偏置参数，控制重要性分数对状态更新的影响强度。

技术创新点

  1. 双阶段设计: 第一阶段压缩信息，第二阶段整合，避免了直接修改SSM扫描算法的复杂性
  2. 选择性机制融合: 巧妙利用Mamba原有的选择性机制整合重要性分数
  3. 可微分近似: 通过修改Δ值而非直接相乘，保证了训练的可微分性
  4. 梯度缩放策略: 根据重要性分数成比例缩放梯度，强调关键信息的学习

实验设置

数据集

  • 训练数据: LongOrca数据集（约50万样本）
    • OpenOrca数据集的长指令调优实例
    • LongAlpaca-12k长上下文对齐数据
    • 最大长度截断至6000 tokens
  • 评测数据:
    • LongBench-E（英文分支）：13个长上下文理解任务
    • L-Eval：6个闭式长上下文任务

评价指标

  • LongBench: 任务特定准确率（如ROUGE、EM、F1等）
  • L-Eval: 闭式任务准确率
  • 推理速度: tokens/second
  • 内存消耗: GPU内存使用量

对比方法

  • 基线模型: Mamba 2.8B（预训练和微调版本）
  • 对比方法:
    • DeciMamba 2.8B
    • Llama-3B（使用线性位置插值扩展上下文）
  • 消融实验: 随机选择、固定选择、乘性选择等变体

实现细节

  • 超参数: s=0, p=0.18, ρ=0.009（LongBench最优配置）
  • 训练策略: LoRA微调，rank=32
  • 优化器: AdamW，学习率2e-5
  • 硬件: 8×A100-80GB GPU，DeepSpeed Zero Stage 3

实验结果

主要结果

LongBench性能对比:

模型Average Score
Mamba (SFT)24.63
ReMamba (SFT)27.86
Llama-3B (SFT)28.99

L-Eval性能对比:

模型Average Score
Mamba (SFT)22.19
ReMamba (SFT)23.83
Llama-3B (SFT)22.69

消融实验

选择策略对比:

  • 随机选择: 与基线性能相近，证实信息丢失假设
  • 固定选择: 略优于随机选择
  • 乘性选择: 有一定改善
  • ReMamba完整方法: 显著优于所有变体

长度泛化性能:

  • ReMamba在2k-9k所有长度上均优于基线
  • 最优性能长度从4k扩展到6k
  • 性能差距随上下文长度增加而扩大

效率分析

内存消耗:

  • ReMamba相比Mamba仅增加少量恒定内存开销
  • 远低于Transformer的二次增长内存需求

推理速度:

  • 与原始Mamba速度相当
  • 显著快于Transformer（约2-3倍）

Mamba2扩展实验

在Mamba2上应用ReMamba方法，LongBench平均分数提升1.6个点，证明方法的通用性。

相关工作

长上下文建模

  1. Transformer扩展: 位置插值、RoPE等技术
  2. Mamba改进: LongMamba通过长上下文微调，DeciMamba通过无训练方法
  3. 混合架构: Jamba等结合attention和SSM的方法

上下文压缩

  1. KV缓存压缩: 针对Transformer的内存优化
  2. 提示压缩: 软提示和检索增强生成方法
  3. 选择性注意: 动态分配计算资源的方法

结论与讨论

主要结论

  1. 问题诊断准确: 成功识别Mamba长上下文性能不足的根本原因
  2. 方法有效性: ReMamba显著提升长上下文性能，接近Transformer水平
  3. 效率保持: 在提升性能的同时保持了Mamba的推理效率优势
  4. 方法通用: 成功扩展到Mamba2，显示良好的普适性

局限性

  1. 理论上限: 由于固定状态空间限制，Mamba难以在超长上下文上超越Transformer
  2. 方法局限: 主要通过压缩缓解信息丢失，未从根本上改变状态更新机制
  3. 超参数敏感: 需要针对不同任务调整压缩参数
  4. 评估范围: 主要在英文数据集上评估，多语言泛化性待验证

未来方向

  1. 状态机制改进: 直接修改状态空间更新机制
  2. 自适应压缩: 根据内容动态调整压缩策略
  3. 多模态扩展: 将方法扩展到视觉-语言任务
  4. 理论分析: 深入分析方法的理论基础和性能边界

深度评价

优点

  1. 问题洞察深刻: 通过随机压缩实验巧妙证明了Mamba的信息丢失问题
  2. 方法设计巧妙: 两阶段设计既保持了可微分性又有效利用了原有机制
  3. 实验全面充分: 包含多个基准、消融实验、效率分析等
  4. 工程实现优秀: 开源代码，便于复现和应用
  5. 写作清晰: 逻辑清晰，技术细节描述准确

不足

  1. 理论分析不足: 缺乏对为什么该方法有效的深层理论解释
  2. 评估局限: 主要在QA类任务上评估，其他类型长上下文任务覆盖不足
  3. 超参数复杂: 需要调整多个超参数，实际应用中可能需要大量调优
  4. 基线对比: DeciMamba的效果不佳可能与超参数设置有关

影响力

  1. 学术价值: 为Mamba长上下文建模提供了新思路和有效解决方案
  2. 实用价值: 方法简单有效，易于在实际系统中部署
  3. 可复现性: 提供完整代码和详细实验设置
  4. 启发意义: 为其他序列建模架构的改进提供了参考

适用场景

  1. 文档理解: 长文档问答、摘要生成等任务
  2. 对话系统: 需要维护长对话历史的场景
  3. 代码理解: 长代码文件的分析和生成
  4. 资源受限环境: 需要高效推理的边缘计算场景

参考文献

核心相关工作:

  1. Gu, A. and Dao, T. (2024). Mamba: Linear-time sequence modeling with selective state spaces.
  2. Dao, T. and Gu, A. (2024). Transformers are ssms: Generalized models and efficient algorithms through structured state space duality.
  3. Bai, Y. et al. (2024). Longbench: A bilingual, multitask benchmark for long context understanding.
  4. Chen, Y. et al. (2024). Longlora: Efficient fine-tuning of long-context large language models.

总体评价: 这是一篇高质量的研究论文，针对Mamba架构的长上下文理解问题提出了创新且有效的解决方案。方法设计巧妙，实验充分，具有良好的理论价值和实用价值。虽然存在一些局限性，但为相关领域的发展做出了重要贡献。